diff --git a/.github/workflows/test_python_cplusplus.yml b/.github/workflows/test_python_cplusplus.yml index 887bccf325..bb6b9d70ed 100644 --- a/.github/workflows/test_python_cplusplus.yml +++ b/.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_*/ diff --git a/.gitignore b/.gitignore index 96bf825460..75f7ba71a5 100644 --- a/.gitignore +++ b/.gitignore @@ -200,3 +200,4 @@ venv/* .coverage tests/sciml/models/* amici_models +doc/examples/example_petab_sciml/Dandekar_Patterns2020/ diff --git a/doc/examples/example_petab_sciml/ExamplePEtabSciML.ipynb b/doc/examples/example_petab_sciml/ExamplePEtabSciML.ipynb new file mode 100644 index 0000000000..6b9181411a --- /dev/null +++ b/doc/examples/example_petab_sciml/ExamplePEtabSciML.ipynb @@ -0,0 +1,238 @@ +{ + "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 +} diff --git a/doc/python_examples.rst b/doc/python_examples.rst index 14303da030..7dd65e3bfb 100644 --- a/doc/python_examples.rst +++ b/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 diff --git a/doc/rtd_requirements.txt b/doc/rtd_requirements.txt index 2c1e4784a6..972007ecc9 100644 --- a/doc/rtd_requirements.txt +++ b/doc/rtd_requirements.txt @@ -29,4 +29,5 @@ ipykernel -e python/sdist/[jax] antimony>=2.13 fiddy>=0.0.3 +petab_sciml[torch] @ git+https://github.com/petab-dev/petab_sciml.git@main -e git+https://github.com/petab-dev/libpetab-python.git@main#egg=petab