PEtab-SciML: Training a Universal Differential Equation

PEtab-SciML extends PEtab with neural networks, allowing hybrid mechanistic/machine-learning models (universal differential equations, UDEs) to be specified in a solver-independent format. This notebook loads a PEtab-SciML problem with AMICI and trains it with a basic JAX/optax training loop.

We use the Dandekar_Patterns2020 problem from the PEtab-SciML benchmark collection: an SIR model of the early COVID-19 outbreak in the UK, where a neural network replaces the quarantine control term (Dandekar et al., Patterns 1(9), 2020). The benchmark ships separate train/validation measurement tables, which we use as-is. For a general introduction to AMICI’s JAX/PEtab interface, see the JAX PEtab notebook.

import os
import urllib.request
from contextlib import contextmanager
from pathlib import Path

import equinox as eqx
import jax
import jax.random as jr
import matplotlib.pyplot as plt
import optax
from amici.importers.petab import PetabImporter
from amici.sim.jax import run_simulations

from petab import v2

jax.config.update("jax_enable_x64", True)

Download the benchmark problem

PEtab-SciML problems reference several files (SBML model, PEtab tables, neural network definition, …) via paths relative to the problem YAML, which AMICI currently expects to resolve locally rather than from a URL. So instead of loading the YAML directly by URL (as in the JAX PEtab notebook), we first download all files listed for the model.

MODEL = "Dandekar_Patterns2020"
BASE_URL = (
    "https://cdn.jsdelivr.net/gh/sebapersson/Benchmark-Models-PEtab-SciML"
    # branch matching the PEtab-SciML extension version AMICI supports
    f"@828d8942b177fe740b75bfe0620f00f7ef867687/Benchmark-Models/{MODEL}"
)
FILES = [
    f"model_{MODEL}.xml",
    f"parameters_{MODEL}.tsv",
    f"observables_{MODEL}.tsv",
    f"mapping_{MODEL}.tsv",
    f"hybridization_{MODEL}.tsv",
    f"net1_{MODEL}.yaml",
    f"net1_ps_{MODEL}.hdf5",
    f"measurements_train_{MODEL}.tsv",
    f"measurements_validation_{MODEL}.tsv",
    f"train_{MODEL}.yaml",
    f"validation_{MODEL}.yaml",
]

model_dir = Path(MODEL)
model_dir.mkdir(exist_ok=True)
for file in FILES:
    urllib.request.urlretrieve(f"{BASE_URL}/{file}", model_dir / file)


@contextmanager
def change_directory(path):
    cwd = os.getcwd()
    os.chdir(path)
    try:
        yield
    finally:
        os.chdir(cwd)

Import the train/validation problems

Both splits share the same model and neural network (net1), but provide different measurements. We import each split as its own JAXProblem.

from IPython.display import clear_output


def import_split(split):
    with change_directory(model_dir):
        petab_problem = v2.Problem.from_yaml(f"{split}_{MODEL}.yaml")
        pi = PetabImporter(
            petab_problem=petab_problem,
            module_name=f"{MODEL}_{split}",
            compile_=True,
            jax=True,
            validate=False,  # PEtab validation does not yet support array-valued parameters
        )
        return pi.create_simulator(force_import=True)


jax_train = import_split("train")
jax_val = import_split("validation")

llh_train, _ = run_simulations(jax_train)
llh_val, _ = run_simulations(jax_val)
print(f"nominal train llh: {llh_train:.2f}, validation llh: {llh_val:.2f}")
clear_output()

Training loop

The nominal parameters already correspond to a well-trained model, so to demonstrate training we first reinitialize the neural network to random weights, then re-fit it to the training data with Adam, tracking log-likelihood on the held-out validation split.

def sync_parameters(target, source):
    """Copy mechanistic and neural network parameters from `source` to `target`."""
    target = target.update_parameters(source.parameters)
    return eqx.tree_at(lambda p: p.model.nns, target, source.model.nns)


# randomly reinitialize the neural network
net_cls = type(jax_train.model.nns["net1"])
jax_train = eqx.tree_at(
    lambda p: p.model.nns["net1"], jax_train, net_cls(jr.PRNGKey(0))
)


def loss_fn(problem):
    llh, aux = run_simulations(problem)
    return -llh, aux


loss_and_grad = eqx.filter_value_and_grad(loss_fn, has_aux=True)
optimizer = optax.adam(3e-3)
opt_state = optimizer.init(eqx.filter(jax_train, eqx.is_inexact_array))


@eqx.filter_jit
def train_step(problem, opt_state):
    (nllh, _), grads = loss_and_grad(problem)
    updates, opt_state = optimizer.update(grads, opt_state)
    problem = eqx.apply_updates(problem, updates)
    return nllh, problem, opt_state

n_steps = 150
train_llh = []
val_llh = []

for step in range(n_steps):
    nllh, jax_train, opt_state = train_step(jax_train, opt_state)
    train_llh.append(-nllh.item())
    jax_val = sync_parameters(jax_val, jax_train)
    val_llh.append(run_simulations(jax_val)[0].item())

plt.plot(train_llh, label="train")
plt.plot(val_llh, label="validation")
plt.xlabel("training step")
plt.ylabel("log-likelihood")
plt.legend()
plt.show()