The “Ahhh” Shower

Rodrigo never intended to have an especially busy day on the 12th. Of course, he had to make a trip to the farmer’s market to get those specialty onions he loved (they were sweeter than the ones in the grocery store), but otherwise the Saturday was completely free. It was a much needed reprieve from the onslaught of work recently with another twelve hour shift on the 13th.

The day started out as planned: waking up with the sun’s rays gently urging his eye lids to make way, with no particular rush. The drive to the market was a smooth twenty minutes. Somehow, traffic was especially clear, and so was the line for the produce. Within an hour, Rodrigo was back home caramelizing the onions for soup and toasting bread. The rest of the day was a blur of mindless entertainment from the multitude of streaming sources while sipping on the French onion soup.

Trish’s Saturday was quite the opposite. Since winter still maintained its grasp on the mornings, she wanted to hike along the foothills in the weak heat of the afternoon. The morning therefore was filled with scrubbing and cleaning. She always loved the sense of accomplishment from finishing chores.

As the hike was finishing up, she got a text from Frank, that cute guy she’s been texting.

“hey, lost our fourth from trivia in 30 @ gecko’s, you in? sorry for short notice!”

The deliberation was short. Whilst her legs were jelly, she didn’t smell too bad since it was still chilly and Frank was, by all accounts, a catch.

“ofc, see you soon :)”

Trish and Frank’s team never came close to winning. No one on the team knew any Madonna songs nor read Julius Caesar. She got home more than ten hours after she left for the trailhead.

Though their days progressed thoroughly differently, Rodrigo and Trish both ended with a hot shower, one where “ahhhh”s were spoken, the true mark of a good day.

April 27, 2022December 9, 2024

New Mexico Rock Formations

I’ve been in New Mexico for awhile now, and while Albuquerque has not been exactly growing on me (more on this sooner or later), the surrounding nature has. There are two particular places in the last few weeks that were quite unique: City of Rocks (the NM version) and Valley of Dreams.

City of Rocks:

Valley of Dreams:

April 10, 2022December 9, 2024

Arrested Development cross Bob’s Burger

Louise Belcher is a never nude. Just with her bunny ears.

April 7, 2022December 9, 2024

Nonlinear Piola Transform

This is an extension for Linear Piola Transform, where we discuss the cases where the transformation from the reference element to the physical element is nonlinear such as a parametric element. The main tool here is actually from the exterior calculus framework, and will just be from the FEEC book by Arnold.

Let $\Omega, \hat \Omega$ by the physical element and reference element with an orientation-preserving, differentiable map $\phi: \hat \Omega \to \Omega$. Suppose $p: \Omega \to \mathbb{R}, v: \Omega \to \mathbb{R}^3$ are the functions on the physical element. The quantity $(p, \nabla \cdot v) = \int_\Omega p \nabla \cdot v \, dx$ is of importance.

The key observation here is that by letting $\nu = v^1 dx^2 \wedge dx^3 – v^2 dx^1 \wedge dx^3 + v^3 dx^1 \wedge dx^2$, we have $d\nu = v$ where $d$ is the exterior derivative. A key fact here is that with integral of differential forms is it preserves pullbacks (see second answer) hence
\begin{align*}
(p, \nabla \cdot v) = \int_\Omega p \nabla \cdot v \, dx = \int_\Omega p \wedge d\nu = \int_{\hat \Omega} \phi^*(p \wedge d\nu)
\end{align*}
where $\phi^*$ is the pullback of differential form $\phi$.

Now, it’s just a matter of algebra
\begin{align*}
\int_{\hat \Omega} \phi^*(p \wedge d\nu) = \int_{\hat \Omega} \phi^*p \wedge \phi^* d\nu = \int_{\hat \Omega} \phi^*p \wedge d\phi^*\nu.
\end{align*}
Now, we need to look up what the pullback does to a 0-form $p$ and on the 2-form $\nu$. Well, we can look this up or just algebraically do it (which I’ll do it sooner or later because I couldn’t find a good source online), we have $\phi^* p = p \cdot \phi$ and
\begin{align*}
\phi^* \nu = (\det \phi’ )(\phi’)^{-1} (v \cdot \phi): \hat \Omega \to \mathbb{R}^3
\end{align*}
which is exactly the Piola transform.

March 25, 2022December 9, 2024

Ne plus ultra

1. the highest point capable of being attained

2. the most profound degree of a quality or state

March 20, 2022December 9, 2024

Helplessness

The last three weeks has been terrible in every sense of the word with the outbreak of war in Europe. Even though I’m more than eight time zones way from the action, my anxiety level is through the roof, and has caused my heart rate to spike. Part of the reason is that I have absolutely no control over the situation. Nothing I do can effect the news.

I think one indication of maturity level is how one handles helplessness. As my world opened up while aging, there’s more and more things which I can’t change. The way I fundamentally look, the way mustard taste, the way someone else feels about you. And now, the way global geopolitics is shaping up for the foreseeable future.

My approach right now is best encapsulated by Vonnegut’s “so it goes.” It’s okay to be frustrated at what’s going on. In fact, one should absolutely be aghast at what’s happening to the climate crisis, Ukraine war, the pandemic… but alas, what can I do besides my little part?

It’s odd how I think about this right now. Perhaps it’s a strategy to deal with overwhelming anxiety, but I always believed that I have an internal locus of control. I firmly believe that my future is under my control, but having moved to ABQ, home of a nuclear lab, I fully realized that if nuclear war were to break out, I would be incinerated instantly. Furthermore, there’s also issues of health in my family that I can’t control. It just sucks that as we age, chaos increasingly dominates our life and helplessness goes up.

So it goes.

February 19, 2022December 9, 2024

Weather

Did it really snow if the snow is not there at the end of the day?

February 9, 2022December 9, 2024

Simple Neural ODE Code

This is a quick and dirty demonstration of training Neural ODEs using Jax. The paper is here. I will be following their notation and functional syntax for easy comparison.

In [1]:

import jax.numpy as jnp
import numpy as np
from jax import grad, jacfwd, jit
from matplotlib import pyplot as plt
from typing import Callable
from tqdm.notebook import tqdm


from IPython import display

Let’s first define some a simple ODE and solve/plot it using a forward Euler time step as proof of concept. Note that all the sophisticated explicit time steppers basically are combinations of Euler’s. At the same time, we will not be using any sophisticated Jax functions for now. Our ODE is of the form $z_t = \theta z$ where $\theta$ is a 2 by 2 matrix.

In [2]:

theta_true = jnp.array([[0.0, 1.0], [-1.0, 0.0]])

def f(u: jnp.ndarray, t: jnp.float32, theta: jnp.ndarray = theta_true) -> jnp.ndarray: 
    return theta @ u

def ODESolve(z_t0: jnp.ndarray, f: Callable, t0: jnp.float32, t1: jnp.float32, theta: jnp.ndarray) -> jnp.ndarray:
    delta_t = 0.001
    N = int((t1 - t0) / delta_t)
    z_holder = jnp.array(z_t0)
    for i in range(N): 
        z_holder += delta_t * f(z_holder, i * delta_t, theta)
    
    return z_holder

WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)

First, let’s make some benchmark data and fake data with only the initial condition off. Are we able to recover the true $z_{t_0}$?

In [3]:

z_t0_true = jnp.array([0, 1.0])
z_t1_true = ODESolve(z_t0_true, f, 0.0, 1.0, theta_true)
z_t0_noisy = jnp.array([0, 1.1])
z_t1_noisy = ODESolve(z_t0_noisy, f, 0.0, 1.0, theta_true)

Now, let’s make a loss function, and see if we can make a gradient decent algorithm such that we can find the true initial condition.

In [4]:

def L(z_t1: jnp.ndarray): 
    return jnp.linalg.norm(z_t1 - z_t1_true) ** 2 / 2

Before proceeding, let’s note that Jax allows one to find gradients and Jacobians very easily!

In [5]:

# partial L/ partial z_t1
plpzt1_noisy = grad(L)(z_t1_noisy)

# Jacobians with respect to z, and theta respectively
f_z = jacfwd(f, 0)
f_theta = jacfwd(f, 2)

Now, we need to find the gradients of L and f at specific points and places. This is simply algorithm 1 of the paper

In [6]:

def algorithm1(theta: jnp.ndarray, t0: jnp.float32, t1: jnp.float32, z_t1: jnp.ndarray, plpzt1: jnp.ndarray):
    # Define initial augmented state
    s0 = jnp.concatenate((z_t1, plpzt1, jnp.zeros_like(theta).flatten()))

    def aug_dynamics(u: jnp.ndarray, t: jnp.float32, theta: jnp.ndarray): 
        """
        u consists of z(t), a(t) and theta (4 by 4), also negative since I can't do backwards lol
        """
        return -jnp.concatenate((f(u[0:2], t, theta), 
                                - u[2:4].T @ f_z(u[0:2], t, theta) , 
                                (-u[2:4].T @ f_theta(u[0:2], t, theta)).flatten()))
    
    output = ODESolve(s0, aug_dynamics, 0.0, 1.0, theta_true)
    
    # Split the augmented data back to our regular inputs
    return output[0:2], output[2:4], output[4:].reshape((2,2))
    
print(algorithm1(theta_true, 0.0, 1.0, z_t1_noisy, plpzt1_noisy))

(DeviceArray([1.0890653e-06, 1.1010987e+00], dtype=float32), DeviceArray([1.3050158e-07, 1.0009895e-01], dtype=float32), DeviceArray([[0.03007106, 0.03902269],
             [0.03902263, 0.08009262]], dtype=float32))

Note that the derivative with respect to the initial condition $\partial L/ \partial z(t_0)$ is essentially $(0, .1)$ which indicates that we have succesfully obtained the correct gradient. Note that two points is not enough to properly define a ODE and we can certainly change the dynamics to go from z_t0_noisy to z_t1_true by simply changing theta.

Let’s encapsulate the above in a gradient descent algorithm… which is a bit slow…

In [7]:

my_theta = jnp.array(theta_true)
learning_rate = 0.1

def train(initial_z0, initial_theta): 
    # Make copies because I don't know how Jax works
    z_t0 = jnp.array(initial_z0)
    theta = jnp.array(initial_theta)
    
    for i in tqdm(range(5)): 
        # Given current initial conditions, let's 
        z_t1 = ODESolve(z_t0, f, 0.0, 1.0, theta)

        plpzt1 = grad(L)(z_t1)

        _, plpzt0, plptheta = algorithm1(theta, 0.0, 1.0, z_t1, plpzt1)

        z_t0 = z_t0 - learning_rate * plpzt0
        theta = theta - learning_rate * plptheta.reshape((2,2))
    
    return z_t0, theta
    
print(train(z_t0_noisy, theta_true))

(DeviceArray([8.1869512e-06, 1.0662646e+00], dtype=float32), DeviceArray([[-0.00988956,  0.9871685 ],
             [-1.0128397 , -0.02635211]], dtype=float32))

This seems to work! Success.

Note that the above code is rather slow, so let’s see how we can use jax things to speed it up! Using the XLA, and compiled code with jit, we see that using a 3rd order solver is even faster than the Euler time step by 10 folds!

In [8]:

from jax import lax

@jit
def f(u: jnp.ndarray, t: jnp.float32, theta: jnp.ndarray = theta_true) -> jnp.ndarray: 
    return theta @ u

def ODESolve_lax(z_t0: jnp.ndarray, f: Callable, t0: jnp.float32, t1: jnp.float32, theta: jnp.ndarray) -> jnp.ndarray:
    """
    We use the lax loop instead, which makes things faster so we use a third order RK3... which is still faster
    """
    delta_t = 0.001
    N = jnp.floor_divide(t1 - t0, delta_t).astype(jnp.int32)
    z_holder = jnp.array(z_t0)
    
    def rk3(i: jnp.int32, val: jnp.ndarray):
        k1 = f(val, i * delta_t, theta)
        k2 = f(val + delta_t * k1, i * delta_t + delta_t, theta)
        k3 = f(val + delta_t * (1 / 4 * k1 + 1 / 4 * k2), i * delta_t + 1 / 2 * delta_t, theta)
        return val + delta_t * (k1 / 6 + k2 / 6 + 2 * k3 / 3)
    
    return lax.fori_loop(0, N, rk3, z_t0)

In [9]:

647 ms ± 20.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
57.1 ms ± 1.44 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

In [10]:

print(ODESolve(z_t0_noisy, f, 0.0, 2 * jnp.pi, theta_true))
print(ODESolve_lax(z_t0_noisy, f, 0.0, 2 * jnp.pi, theta_true))

[-2.0616804e-04  1.1034613e+00]
[-2.0406539e-04  1.1000000e+00]

Now that ODESolve_lax is much faster (this is only on CPU… performance should be better if we export to GPU), we can try to apply jit to other functions to make it faster, and try to fit an actual ODE with mulitple points. First let’s define a loss function, and use jit to define the Jacobians:

In [11]:

@jit
def L(z_t1: jnp.ndarray, z_t1_true: jnp.ndarray): 
    """
    Simple function; easy to jit
    """
    return jnp.linalg.norm(z_t1 - z_t1_true) ** 2 / 2

In [12]:

f_z = jit(jacfwd(f, 0))
f_theta = jit(jacfwd(f, 2))

Now, rewriting algorithm1 with jit is fairly easy…

In [13]:

@jit
def algorithm1(theta: jnp.ndarray, t0: jnp.float32, t1: jnp.float32, z_t1: jnp.ndarray, plpzt1: jnp.ndarray):
    # Define initial augmented state
    s0 = jnp.concatenate((z_t1, plpzt1, jnp.zeros_like(theta).flatten()))
    
    # Just an index for the length of the problem 
    z_size = len(z_t1)
    
    @jit
    def aug_dynamics(u: jnp.ndarray, t: jnp.float32, theta: jnp.ndarray): 
        """
        u consists of z(t), a(t) and theta, also negative since I can't do backwards lol
        """
        z = u[0:z_size]
        a = u[z_size:2 * z_size]
        return -jnp.concatenate((f(z, t, theta), 
                                - a.T @ f_z(z, t, theta) , 
                                (-a.T @ f_theta(z, t, theta)).flatten()))
    

    output = ODESolve_lax(s0, aug_dynamics, t0, t1, theta_true)
    return output[0:z_size], output[z_size:2 * z_size], output[2 * z_size:].reshape((z_size, z_size))

print(algorithm1(theta_true, 0.0, 1.0, z_t1_noisy, plpzt1_noisy))

(DeviceArray([1.1002590e-03, 1.1005477e+00], dtype=float32), DeviceArray([0.00010013, 0.10004898], dtype=float32), DeviceArray([[0.03002402, 0.03898256],
             [0.03898243, 0.07997467]], dtype=float32))

As written the above algorithm doesn’t support batching, which is a huge pain in the butt. I couldn’t get vmap from jax to work out of the box but we won’t need it here.

Now the gradient formula is much faster. Let’s make a cooler true ODE solution, and plot it

In [14]:

theta_true = jnp.array([[-.1, .9], [-.9, -.2]])
init_true = jnp.array([0.0, 1.0])

total_time = 4 * jnp.pi
num_data = 20

@jit
def predict(initial_point: jnp.ndarray, theta: jnp.ndarray, total_time = 4 * jnp.pi, num_data = 20): 
    # jnp arrays are immutable
    prediction = jnp.zeros((num_data, 2))
    prediction = prediction.at[0, :].set(initial_point)
    for i in range(1, num_data): 
        prediction = prediction.at[i, :].set(ODESolve_lax(prediction[i - 1, :], f, total_time / num_data * i, 
                                        total_time / num_data * (i + 1), theta))

    return jnp.array(prediction)


all_data = predict(init_true, theta_true)

# TODO: we might want to add some sort of noise in the future
plt.plot(all_data[:, 0], all_data[:, 1], '-o')

Out[14]:

[<matplotlib.lines.Line2D at 0x7f30fc134c18>]

With the data done, we want to recover the both the dynamics and the initial condition. I think the ODE demo on the author github only recovers dynamics.

Since we care about the exact initial conditions for the above dynamics, it only makes sense that all our samples which alters the initial conditions should contain our dirty initial conditions. The following training loop does exactly that; it takes some sample of $(t_0, t_K)$ where $K$ is some random number (could be a batch) and then alters the parameters due to that.

The initial theta is very sensitive; this sort of model can’t really optimize over bifurcations very well it seems so if the model doesn’t converge, just try to run again.

In [15]:

# We start with some initial parameter and points
my_theta = jnp.array([[0.0, 1.0], [-1.0, 0.0]]) + np.random.normal(0, .4, size=(2,2))
initial_point = init_true + np.random.normal(0, .2, size=(2,))

# Custom gradient descent; no need for Adams in this case hopefully
learning_rate = 0.1

gradL = jit(grad(L))
batch_size = 3
epochs = 1000

for epoch in tqdm(range(epochs)): 
    if epoch 
        z = predict(initial_point, my_theta)
        plt.plot(z[:, 0], z[:, 1], alpha=(epoch / epochs) * .7 + 0.3 )
        plt.xlim([-.9, 0.9])
        plt.ylim([-1.0, 1.4])
        plt.title(f"Epoch {epoch}; Loss {L(z, all_data)}")
        display.clear_output(wait=True)
        display.display(plt.gcf())
    
    # Once we get close enough it's fine
    if L(z, all_data) < 1e-4: 
        break
    
    # Choose batch of time 
    times_indices = np.random.choice(np.arange(1, len(all_data)), size=batch_size, replace=False)
    
    for ind in times_indices:
        my_point = ODESolve_lax(initial_point, f, 0, ind * total_time / num_data, my_theta)
        
        # Get gradient # Need to batch this ultimately at some point 
        _, plpz0, plptheta = algorithm1(
            my_theta, 
            0.0, 
            ind * total_time / num_data, 
            my_point, 
            gradL(my_point, all_data[ind, :])
        )
        
        # Apply gradient 
        initial_point = initial_point - learning_rate * plpz0
        my_theta = my_theta - learning_rate * plptheta.reshape((2,2))
        
plt.plot(all_data[:, 0], all_data[:, 1], '-o', linestyle='dashed', alpha=.9)

Out[15]:

February 6, 2022December 9, 2024

Hitting the Slopes

I recently had an epitome the other day: having fun skiing entails skiing less.

The logic is simple. A tired skier is an unhappy skier. A tired skier is a an injury-prone skier. A tired skier is a miserable skier.

Another particular benefit of skiing less is I can save money, since many ski parks offer half-day passes for a discount. By only opting for the afternoon session, I also get to experience a warmer day with less crowds since après ski tend to set in around lunch time. There’s also the fact that I could sleep in more, and not contend with the early morning rush.

I finally understand less is more now.

January 23, 2022December 9, 2024

A Walk in the Woods

My third Bill Bryson book after a Short History of Nearly Everything and At Home. As usual, I loved his writing style and humor. I found myself laughing and snickering more while reading this book about struggling through the Appalachian Trail then the Anxious People novel.

Pros:

Lots of facts, and his usage of vocabulary has always been strong.
Dry, British humor

Cons:

No sources; there is a fact about the average American walking less than ____ miles a week which I couldn’t verify.

The novel was published back in 1998, I do wonder if anything has changed since. There were quite a few sections where Bryson crucified certain governmental actions (or lack thereof) which I suspect will have gained a lot of attention due to the book/movie. It’s just so depressing sometimes to read about the accelerating decline of nature.