Add PEtab SciML example notebook

.github/workflows/test_python_cplusplus.yml

@@ -239,6 +239,11 @@ jobs:
     - name: Install python package
       run: scripts/installAmiciSource.sh
 
+    - name: Install torch (for PEtab-SciML notebook)
+      run: |
+        source ./venv/bin/activate \
+          && python -m pip install "petab_sciml[torch] @ git+https://github.com/petab-dev/petab_sciml.git@main"
+
     - name: example notebooks
       run: scripts/runNotebook.sh doc/examples/example_*/
 

.gitignore

@@ -200,3 +200,4 @@ venv/*
 .coverage
 tests/sciml/models/*
 amici_models
+doc/examples/example_petab_sciml/Dandekar_Patterns2020/

doc/examples/example_petab_sciml/ExamplePEtabSciML.ipynb

@@ -0,0 +1,237 @@
+{
+ "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",
+    "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://raw.githubusercontent.com/sebapersson/Benchmark-Models-PEtab-SciML/\"\n",
+    "    # branch matching the PEtab-SciML extension version AMICI supports\n",
+    "    f\"bes/networks/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
+}

doc/python_examples.rst

@@ -19,4 +19,5 @@ Various example notebooks.
    examples/example_large_models/example_performance_optimization.ipynb
    examples/example_jax/ExampleJax.ipynb
    examples/example_jax_petab/ExampleJaxPEtab.ipynb
+   examples/example_petab_sciml/ExamplePEtabSciML.ipynb
    examples/example_splines/ExampleSplines.ipynb