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

import os, sys, subprocess, importlib, pathlib


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)


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)