A Coding Guide to Implement Advanced Differential Equation Solvers, Stochastic Simulations, and Neural Ordinary Differential Equations Using Diffrax and JAX

SENTINEL = “/tmp/diffrax_colab_ready_v3”

def _run(cmd):
subprocess.check_call(cmd)

def _need_install():
try:
import numpy
import jax
import diffrax
import equinox
import optax
import matplotlib
return False
except Exception:
return True

if not os.path.exists(SENTINEL) or _need_install():
_run([sys.executable, “-m”, “pip”, “uninstall”, “-y”, “numpy”, “jax”, “jaxlib”, “diffrax”, “equinox”, “optax”])
_run([sys.executable, “-m”, “pip”, “install”, “-q”, “–upgrade”, “pip”])
_run([
sys.executable, “-m”, “pip”, “install”, “-q”,
“numpy==1.26.4”,
“jax[cpu]==0.4.38”,
“jaxlib==0.4.38”,
“diffrax”,
“equinox”,
“optax”,
“matplotlib”
])
pathlib.Path(SENTINEL).write_text(“ready”)
print(“Packages installed cleanly. Runtime will restart now. After reconnect, run this same cell again.”)
os._exit(0)

- Advertisement -

import time
import math
import numpy as np
import jax
import jax.numpy as jnp
import jax.random as jr
import diffrax
import equinox as eqx
import optax
import matplotlib.pyplot as plt

print(“NumPy:”, np.__version__)
print(“JAX:”, jax.__version__)
print(“Backend:”, jax.default_backend())

def logistic(t, y, args):
r, k = args
return r * y * (1 – y / k)

t0, t1 = 0.0, 10.0
ts = jnp.linspace(t0, t1, 300)
y0 = jnp.array(0.4)
args = (2.0, 5.0)

sol_logistic = diffrax.diffeqsolve(
diffrax.ODETerm(logistic),
diffrax.Tsit5(),
t0=t0,
t1=t1,
dt0=0.05,
y0=y0,
args=args,
saveat=diffrax.SaveAt(ts=ts, dense=True),
stepsize_controller=diffrax.PIDController(rtol=1e-6, atol=1e-8),
max_steps=100000,
)

query_ts = jnp.array([0.7, 2.35, 4.8, 9.2])
query_ys = jax.vmap(sol_logistic.evaluate)(query_ts)

print(“n=== Example 1: Logistic growth ===”)
print(“Saved solution shape:”, sol_logistic.ys.shape)
print(“Interpolated values:”)
for t_, y_ in zip(query_ts, query_ys):
print(f”t={float(t_):.3f} -> y={float(y_):.6f}”)

def lotka_volterra(t, y, args):
alpha, beta, delta, gamma = args
prey, predator = y
dprey = alpha * prey – beta * prey * predator
dpred = delta * prey * predator – gamma * predator
return jnp.array([dprey, dpred])

lv_y0 = jnp.array([10.0, 2.0])
lv_args = (1.5, 1.0, 0.75, 1.0)
lv_ts = jnp.linspace(0.0, 15.0, 500)

sol_lv = diffrax.diffeqsolve(
diffrax.ODETerm(lotka_volterra),
diffrax.Dopri5(),
t0=0.0,
t1=15.0,
dt0=0.02,
y0=lv_y0,
args=lv_args,
saveat=diffrax.SaveAt(ts=lv_ts),
stepsize_controller=diffrax.PIDController(rtol=1e-6, atol=1e-8),
max_steps=100000,
)

print(“n=== Example 2: Lotka-Volterra ===”)
print(“Shape:”, sol_lv.ys.shape)