Take 4 at a stable blog
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:
Louise Belcher is a never nude. Just with her bunny ears.
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.
1. the highest point capable of being attained
2. the most profound degree of a quality or state
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.
Did it really snow if the snow is not there at the end of the day?
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.
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.
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}$?
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.
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!
# 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
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…
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!
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)
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)
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:
@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
f_z = jit(jacfwd(f, 0))
f_theta = jit(jacfwd(f, 2))
Now, rewriting algorithm1 with jit is fairly easy…
@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
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')
[<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.
# 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)
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.
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:
Cons:
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.
Nowadays, home is not only where the heart is, but also where my office is. I spend over twenty one hours per day on average inside the confines of the four walls. This is partly why I splurged for the two bedroom instead of the one, and why I’ve been sort of obsessing over air quality recently.
Ever since college, I’ve developed some sort of allergic reaction whenever I visit my parent’s house in Florida. I’m convinced it’s due to some sort of ragweed floating around, or the sudden changes in stress leads to a floundering immune system. When there’s nobody around, I just stuck tissues up my nose… luckily, my sinus passage has so far enjoyed New Mexico… for the most part.
I’ve noticed in the afternoons that my nose are constantly flared, and that repeated vacuuming of my carpets still leads to large amounts of dust piling in the bin. The rays of sunshine streaming into my apartments laid bare the amount of dust floating around. It makes sense; the dust comes from, well, the twenty one hours I spend at home and the fact that the dry, high desert environment readily kicks up the sand and particulates from the ground. Since no home is sealed, those little particulate matter diffuse into my castle.
Ever since I got my Airmega 200M air purifier, that nose flaring has gone away for the most part. A side effect is that I have to dust my home less with that HEPA filter running. It has also made me more aware of the effects on indoor air quality of cooking and vacuuming. The filter plus my humidifier really makes my indoor air quality much better. This translates to a happier Marshall.
Now, I’m concerned about my water quality…