Universal Differential Equations¶
A Universal Differential Equation (UDE) combines a mechanistic ODE with a neural network to learn unknown or partially-known dynamics from data. The ODE encodes prior biological knowledge; the neural network (called the NDE component) learns residual dynamics the ODE cannot explain on its own.
This notebook demonstrates UDE training on a Lotka-Volterra predator-prey system with an unknown external forcing term.
import diffrax
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import optax
import pandas as pd
from example_models import get_lotka_volterra
from mxlpy import plot
from mxlpy.jax.models import (
Anode,
AugUde,
FluxAnode,
FluxAugUde,
FluxOde,
HardLatentMapper,
Node,
Ode,
Ude,
)
from mxlpy.jax.train import IntegrationSettings, train
lv = Ode.from_mxlpy(get_lotka_volterra())
ts = jnp.linspace(0, 50, 100)
y0 = jnp.array([10.0, 10.0])
ys = lv.integrate(ts, y0, max_steps=8192)
fig, ax = plt.subplots()
ax.plot(ts, ys)
ax.legend(["prey", "pred"])
plt.show()
INFO:jax._src.xla_bridge:Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: libtpu.so: cannot open shared object file: No such file or directory
Baseline: mechanistic ODE¶
The Lotka-Volterra system describes predator-prey dynamics with four parameters (a, b, c, d). This serves as the mechanistic prior — the part of the system we already understand and can encode explicitly.
ys_forced = (ys.T + 0.14 * jnp.sin(ts)).T
fig, ax = plt.subplots()
ax.plot(ts, ys_forced)
ax.legend(["prey", "pred"])
plt.show()
Training the UDE¶
We construct a Ude by combining the mechanistic Ode with a small neural network (Node). The op="+" means the NDE output is added to the ODE right-hand side at each time step.
Training uses train, which freezes the ODE parameters and only optimises the neural network weights. A curriculum schedule gradually increases the proportion of the trajectory used for loss computation, which helps avoid poor local minima early in training.
ude, losses, grad_norms = train(
Ude(
ode=lv,
nn=Node(
variable_names=lv.variable_names,
width=8,
depth=3,
key=jax.random.PRNGKey(42),
out_scale=jnp.array([0.1]),
),
op="+",
),
ts=ts,
ys=ys_forced,
training_steps=[
(500, 0.2),
(500, 0.5),
(10_000, 1.0), # in practice longer, shortened for docs
],
avg_every=50,
integration_settings=IntegrationSettings(
max_steps=8192,
method=diffrax.Tsit5,
),
optim=optax.adabelief(learning_rate=2e-4),
)
ys_pred = ude.integrate(ts, y0, max_steps=8192)
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4), layout="constrained")
# ax1
ax1.plot(ts, ys_forced)
plot.reset_prop_cycle(ax=ax1)
ax1.plot(ts, ys_pred, ls="dashed")
ax1.legend(["prey", "pred"])
# ax2
pd.concat(losses, ignore_index=True)[12:].plot(ax=ax2)
plt.show()
simulate_time_course returns the same trajectory as a named result: arrays (ys, readouts) that also work inside jax.grad, and DataFrames for inspection.
sim = ude.simulate_time_course(ts, y0, max_steps=8192)
sim.get_combined().head()
| Prey | Predator | prey_growth | predation | predator_death | |
|---|---|---|---|---|---|
| 0.000000 | 10.000000 | 10.000000 | 1.000000 | 100.000000 | 4.000000 |
| 0.505050 | 9.550498 | 9.020997 | 0.955050 | 86.155014 | 3.608399 |
| 1.010101 | 9.210168 | 8.105610 | 0.921017 | 74.654030 | 3.242244 |
| 1.515152 | 8.962448 | 7.261772 | 0.896245 | 65.083252 | 2.904709 |
| 2.020202 | 8.793919 | 6.492410 | 0.879392 | 57.093723 | 2.596964 |
Separate the contribution of each model part¶
After training we can evaluate the ODE and NDE components independently along the predicted trajectory. jax.vmap maps each component over all time points in one call. This decomposition reveals what the neural network actually learned — ideally it should recover the unknown forcing signal.
rhs_ode = jax.vmap(ude.ode, in_axes=(0, 0, None))(ts, ys_pred, jnp.array([]))
rhs_nde = jax.vmap(ude.nn, in_axes=(0, 0, None))(ts, ys_pred, jnp.array([]))
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4), layout="constrained")
ax1.set(title="ODE")
ax2.set(title="NDE")
ax1.plot(ts, rhs_ode)
ax2.plot(ts, rhs_nde)
for ax in (ax1, ax2):
ax.legend(["dpreydt", "dpreddt"])
The ODE panel shows the mechanistic dynamics encoded in the prior model. The NDE panel shows the learned correction — the neural network's contribution that accounts for dynamics the ODE could not explain. A well-trained UDE produces an NDE signal that approximates the unknown forcing term.
Neural ODEs without a prior: Node¶
Node is the neural component on its own, without any mechanistic prior:
$$\frac{dy}{dt} = s \cdot \text{MLP}(y, \text{args})$$
where $s$ is out_scale. We train it on the same data, with the same budget and solver settings as the UDE above. To compare the models fairly, the remaining sections share one helper.
def fit(model):
trained, losses, _ = train(
model,
ts=ts,
ys=ys_forced,
training_steps=[
(500, 0.2),
(500, 0.5),
(10_000, 1.0), # in practice longer, shortened for docs
],
avg_every=50,
integration_settings=IntegrationSettings(
max_steps=8192,
method=diffrax.Tsit5,
),
optim=optax.adabelief(learning_rate=2e-4),
)
return trained, pd.concat(losses, ignore_index=True)
def plot_fit(model, losses):
ys_pred = model.integrate(ts, y0, max_steps=8192)
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4), layout="constrained")
ax1.plot(ts, ys_forced)
plot.reset_prop_cycle(ax=ax1)
ax1.plot(ts, ys_pred, ls="dashed")
ax1.legend(["prey", "pred"])
losses[12:].plot(ax=ax2)
plt.show()
node, node_losses = fit(
Node(
variable_names=lv.variable_names,
width=8,
depth=3,
key=jax.random.PRNGKey(42),
)
)
plot_fit(node, node_losses)
The pure neural ODE settles into a poor solution: it follows the first few time points, then drifts off monotonically instead of oscillating. Without a prior, it has to discover the oscillation from data alone, and here it doesn't within the training budget.
Part of the difficulty is structural. A Node is an autonomous ODE on the two observed variables, so exactly one trajectory passes through each point of the (prey, predator) plane, and trajectories can never cross. Dynamics close to that limit need a very contorted vector field, which is hard to learn, and dynamics beyond it, such as time-dependent forcing that makes the data revisit a state with a different velocity, can't be represented at all.
Augmented neural ODEs: Anode¶
Augmented neural ODEs (Dupont et al., 2019) lift this restriction by integrating in a larger space. The state $y$ (n_obs dimensions) is extended by n_hidden hidden dimensions $h$. The ODE runs on the latent state $z = (y, h)$, and trajectories that cross in observation space can be separated along $h$.
A latent mapper encodes the initial observation into $z_0$ and decodes the latent trajectory back to observations. integrate and simulate_* always take and return observations; the latent state stays internal.
| Mapper | Encode $y_0 \to z_0$ | Decode $z \to y$ |
|---|---|---|
SoftLatentMapper |
$(y_0, \text{linear}(y_0))$ | $\text{linear}(z)$ |
HardLatentMapper |
$(y_0, \text{linear}(y_0))$ | $z_{[:n_{obs}]}$ |
FullLatentMapper |
$\text{linear}(y_0)$ | $\text{linear}(z)$ |
MixedLatentMapper |
$(y_0, \text{linear}(y_0))$ | $(1-\alpha)\,z_{[:n_{obs}]} + \alpha\,\text{linear}(z)$ |
TrainableZ0LatentMapper |
$(y_0, z_0^{\text{learned}})$ | as MixedLatentMapper |
Anode defaults to SoftLatentMapper. The two annealed mappers need alpha, which you pass through integrate*(alpha=...) or IntegrationSettings(alpha=...). Lowering alpha over a training curriculum moves the decode from the learned map towards the plain slice.
anode, anode_losses = fit(
Anode(
variable_names=lv.variable_names,
n_hidden=2,
width=8,
depth=3,
key=jax.random.PRNGKey(42),
)
)
plot_fit(anode, anode_losses)
The network, width and training budget are the same as for the Node, plus two hidden dimensions, and the fit is now close.
Augmented UDEs: AugUde¶
AugUde combines a mechanistic Ode with an Anode, in the same way that Ude combines it with a Node. With $d = \text{Anode}(t, z, \text{args})$ on the latent state $z = (y, h)$:
$$ \begin{aligned} \frac{dy}{dt} &= \text{op}\big(\text{ODE}(t, y, \text{args}),\ d_{[:n_{obs}]}\big) \\ \frac{dh}{dt} &= d_{[n_{obs}:]} \end{aligned} $$
The mechanistic model only sees the observed slice $y = z_{[:n_{obs}]}$, so the hidden dimensions reach the observed dynamics only through the neural correction. This follows the hybrid-modelling literature. Augmented NODEs define the hidden dimensions for the neural model. Physics-plus-correction hybrids such as APHYNITY (Yin et al., 2021) keep the mechanistic term on the physical state. In latent hybridisation (Qian et al., 2021), the expert ODE runs on its own variables and drives the ML variables, not the other way round. If unobserved variables have known mechanisms, add them to the Ode itself rather than as hidden dimensions.
Because the mechanistic part reads $z_{[:n_{obs}]}$, the latent mapper has to keep the observation there:
HardLatentMapperis the interpretable choice: the decoded observations are exactly the state the mechanistic model sees.MixedLatentMapperandTrainableZ0LatentMapperanneal towards it.SoftLatentMapperis allowed, but its learned decode means the reported observations differ from what the ODE acts on.FullLatentMapperdoesn't keep the observation and is rejected.
aug_ude, aug_ude_losses = fit(
AugUde(
ode=lv,
nn=Anode(
variable_names=lv.variable_names,
n_hidden=2,
width=8,
depth=3,
key=jax.random.PRNGKey(42),
latent_mapper=HardLatentMapper,
),
op="+",
)
)
plot_fit(aug_ude, aug_ude_losses)
On this small example, AugUde reaches train's default target loss early and fits about as well as the plain Ude above. The hidden dimensions earn their keep when the correction depends on something the observations don't contain, such as an unobserved species or an internal clock.
Flux-based variants¶
Every model above has a flux-based counterpart for models built with reactions and stoichiometry. They combine the mechanistic and neural fluxes, then map them through the mechanistic stoichiometry nv:
| State-based | Flux-based | Neural part |
|---|---|---|
Ode |
FluxOde |
none |
Node |
FluxNode |
flux MLP |
Ude |
FluxUde |
FluxNode |
Anode |
FluxAnode |
flux + Markov MLP |
AugUde |
FluxAugUde |
FluxAnode |
In FluxAnode and FluxAugUde the hidden dimensions can't go through nv, which only maps onto the observed variables. Their dynamics come from a separate Markov MLP:
$$ \begin{aligned} v &= \text{op}\big(\text{fluxes}_{\text{ODE}}(t, y, \text{args}),\ \text{fluxes}_{\text{NN}}(t, z, \text{args})\big) \\ \frac{dy}{dt} &= N v \\ \frac{dh}{dt} &= \text{Markov}(z, \text{args}) \end{aligned} $$
FluxNode.from_flux_ode and FluxAnode.from_flux_ode build the neural part with the mechanistic model's variable names, flux names, stoichiometry and readouts. Simulations of FluxUde/FluxAugUde report the combined fluxes, and additionally the ode_fluxes and nn_fluxes that went into them. The model below is untrained; it only shows the construction.
flux_lv = FluxOde.from_mxlpy(get_lotka_volterra())
flux_aug_ude = FluxAugUde(
flux_ode=flux_lv,
flux_nn=FluxAnode.from_flux_ode(
flux_lv,
n_hidden=2,
key=jax.random.PRNGKey(42),
latent_mapper=HardLatentMapper,
),
op="+",
)
sim = flux_aug_ude.simulate_time_course(ts, y0, max_steps=8192)
sim.fluxes_df.head()
| prey_growth | predation | predator_death | |
|---|---|---|---|
| 0.000000 | 0.967600 | 99.998016 | 4.005514 |
| 0.505050 | 0.923526 | 85.999863 | 3.610024 |
| 1.010101 | 0.890310 | 74.379555 | 3.240556 |
| 1.515152 | 0.866265 | 64.715103 | 2.900183 |
| 2.020202 | 0.850023 | 56.651234 | 2.589992 |