{ "cells": [ { "cell_type": "markdown", "id": "fdf170e781de4e98ac98", "metadata": {}, "source": [ "# PEtab-SciML: Training a Universal Differential Equation\n", "\n", "[PEtab-SciML](https://github.com/petab-dev/petab_sciml) extends [PEtab](https://petab.readthedocs.io) 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](https://optax.readthedocs.io) training loop.\n", "\n", "We use the [Dandekar_Patterns2020](https://github.com/sebapersson/Benchmark-Models-PEtab-SciML/tree/main/Benchmark-Models/Dandekar_Patterns2020) problem from the [PEtab-SciML benchmark collection](https://github.com/sebapersson/Benchmark-Models-PEtab-SciML): 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](../example_jax_petab/ExampleJaxPEtab.ipynb)." ] }, { "cell_type": "code", "execution_count": null, "id": "47a582e4e962401295d8", "metadata": {}, "outputs": [], "source": [ "import os\n", "import urllib.request\n", "from contextlib import contextmanager\n", "from pathlib import Path\n", "\n", "import equinox as eqx\n", "import jax\n", "import jax.random as jr\n", "import matplotlib.pyplot as plt\n", "import optax\n", "from amici.importers.petab import PetabImporter\n", "from amici.sim.jax import run_simulations\n", "\n", "from petab import v2\n", "\n", "jax.config.update(\"jax_enable_x64\", True)" ] }, { "cell_type": "markdown", "id": "c41fa3e76c0140d6b678", "metadata": {}, "source": [ "## Download the benchmark problem\n", "\n", "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](../example_jax_petab/ExampleJaxPEtab.ipynb)), we first download all files listed for the model." ] }, { "cell_type": "code", "execution_count": null, "id": "132a165298a24249a26d", "metadata": {}, "outputs": [], "source": [ "MODEL = \"Dandekar_Patterns2020\"\n", "BASE_URL = (\n", " \"https://cdn.jsdelivr.net/gh/sebapersson/Benchmark-Models-PEtab-SciML\"\n", " # branch matching the PEtab-SciML extension version AMICI supports\n", " f\"@828d8942b177fe740b75bfe0620f00f7ef867687/Benchmark-Models/{MODEL}\"\n", ")\n", "FILES = [\n", " f\"model_{MODEL}.xml\",\n", " f\"parameters_{MODEL}.tsv\",\n", " f\"observables_{MODEL}.tsv\",\n", " f\"mapping_{MODEL}.tsv\",\n", " f\"hybridization_{MODEL}.tsv\",\n", " f\"net1_{MODEL}.yaml\",\n", " f\"net1_ps_{MODEL}.hdf5\",\n", " f\"measurements_train_{MODEL}.tsv\",\n", " f\"measurements_validation_{MODEL}.tsv\",\n", " f\"train_{MODEL}.yaml\",\n", " f\"validation_{MODEL}.yaml\",\n", "]\n", "\n", "model_dir = Path(MODEL)\n", "model_dir.mkdir(exist_ok=True)\n", "for file in FILES:\n", " urllib.request.urlretrieve(f\"{BASE_URL}/{file}\", model_dir / file)\n", "\n", "\n", "@contextmanager\n", "def change_directory(path):\n", " cwd = os.getcwd()\n", " os.chdir(path)\n", " try:\n", " yield\n", " finally:\n", " os.chdir(cwd)" ] }, { "cell_type": "markdown", "id": "c1d505df13ed41329e73", "metadata": {}, "source": [ "## Import the train/validation problems\n", "\n", "Both splits share the same model and neural network (`net1`), but provide different measurements. We import each split as its own [JAXProblem](https://amici.readthedocs.io/en/latest/generated/amici.sim.jax.html#amici.sim.jax.JAXProblem)." ] }, { "cell_type": "code", "execution_count": null, "id": "c859ba3f124b4bb986fc", "metadata": {}, "outputs": [], "source": [ "from IPython.display import clear_output\n", "\n", "\n", "def import_split(split):\n", " with change_directory(model_dir):\n", " petab_problem = v2.Problem.from_yaml(f\"{split}_{MODEL}.yaml\")\n", " pi = PetabImporter(\n", " petab_problem=petab_problem,\n", " module_name=f\"{MODEL}_{split}\",\n", " compile_=True,\n", " jax=True,\n", " validate=False, # PEtab validation does not yet support array-valued parameters\n", " )\n", " return pi.create_simulator(force_import=True)\n", "\n", "\n", "jax_train = import_split(\"train\")\n", "jax_val = import_split(\"validation\")\n", "\n", "llh_train, _ = run_simulations(jax_train)\n", "llh_val, _ = run_simulations(jax_val)\n", "print(f\"nominal train llh: {llh_train:.2f}, validation llh: {llh_val:.2f}\")\n", "clear_output()" ] }, { "cell_type": "markdown", "id": "f78ac1aa1c4e4e3d9bb9", "metadata": {}, "source": [ "## Training loop\n", "\n", "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." ] }, { "cell_type": "code", "execution_count": null, "id": "9cdde64d1dab411fa863", "metadata": {}, "outputs": [], "source": [ "def sync_parameters(target, source):\n", " \"\"\"Copy mechanistic and neural network parameters from `source` to `target`.\"\"\"\n", " target = target.update_parameters(source.parameters)\n", " return eqx.tree_at(lambda p: p.model.nns, target, source.model.nns)\n", "\n", "\n", "# randomly reinitialize the neural network\n", "net_cls = type(jax_train.model.nns[\"net1\"])\n", "jax_train = eqx.tree_at(\n", " lambda p: p.model.nns[\"net1\"], jax_train, net_cls(jr.PRNGKey(0))\n", ")\n", "\n", "\n", "def loss_fn(problem):\n", " llh, aux = run_simulations(problem)\n", " return -llh, aux\n", "\n", "\n", "loss_and_grad = eqx.filter_value_and_grad(loss_fn, has_aux=True)\n", "optimizer = optax.adam(3e-3)\n", "opt_state = optimizer.init(eqx.filter(jax_train, eqx.is_inexact_array))\n", "\n", "\n", "@eqx.filter_jit\n", "def train_step(problem, opt_state):\n", " (nllh, _), grads = loss_and_grad(problem)\n", " updates, opt_state = optimizer.update(grads, opt_state)\n", " problem = eqx.apply_updates(problem, updates)\n", " return nllh, problem, opt_state" ] }, { "cell_type": "code", "execution_count": null, "id": "52d7677e09274ce9a7cb", "metadata": {}, "outputs": [], "source": [ "n_steps = 150\n", "train_llh = []\n", "val_llh = []\n", "\n", "for step in range(n_steps):\n", " nllh, jax_train, opt_state = train_step(jax_train, opt_state)\n", " train_llh.append(-nllh.item())\n", " jax_val = sync_parameters(jax_val, jax_train)\n", " val_llh.append(run_simulations(jax_val)[0].item())" ] }, { "cell_type": "code", "execution_count": null, "id": "c6287e9353d14ca7a748", "metadata": {}, "outputs": [], "source": [ "plt.plot(train_llh, label=\"train\")\n", "plt.plot(val_llh, label=\"validation\")\n", "plt.xlabel(\"training step\")\n", "plt.ylabel(\"log-likelihood\")\n", "plt.legend()\n", "plt.show()" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.12.3" } }, "nbformat": 4, "nbformat_minor": 5 }