From 42d4767047c2875ab67b5ebee70f6eebacdcb997 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fabian=20Fr=C3=B6hlich?= Date: Wed, 4 Dec 2024 17:32:07 +0000 Subject: [PATCH] disentangle sim & preeq --- documentation/ExampleJaxPEtab.ipynb | 674 +++++++++++++++++++++++++++- 1 file changed, 673 insertions(+), 1 deletion(-) mode change 120000 => 100644 documentation/ExampleJaxPEtab.ipynb diff --git a/documentation/ExampleJaxPEtab.ipynb b/documentation/ExampleJaxPEtab.ipynb deleted file mode 120000 index 821b14f21f..0000000000 --- a/documentation/ExampleJaxPEtab.ipynb +++ /dev/null @@ -1 +0,0 @@ -../python/examples/example_jax_petab/ExampleJaxPEtab.ipynb \ No newline at end of file diff --git a/documentation/ExampleJaxPEtab.ipynb b/documentation/ExampleJaxPEtab.ipynb new file mode 100644 index 0000000000..855860e242 --- /dev/null +++ b/documentation/ExampleJaxPEtab.ipynb @@ -0,0 +1,673 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "d4d2bc5c", + "metadata": {}, + "source": [ + "# Simulating AMICI models using JAX\n", + "\n", + "## Overview\n", + "\n", + "This guide demonstrates how to use AMICI to export models in a format compatible with the [JAX](https://jax.readthedocs.io/en/latest/) ecosystem, enabling simulations with the [diffrax](https://docs.kidger.site/diffrax/) library. " + ] + }, + { + "cell_type": "markdown", + "id": "fb2fe897", + "metadata": {}, + "source": [ + "## Preparation\n", + "\n", + "To begin, we will import a model using [PEtab](https://petab.readthedocs.io). For this demonstration, we will utilize the [Benchmark Collection](https://github.com/Benchmarking-Initiative/Benchmark-Models-PEtab), which provides a diverse set of models. For more information on importing PEtab models, refer to the corresponding [PEtab notebook](https://amici.readthedocs.io/en/latest/petab.html).\n", + "\n", + "In this tutorial, we will import the Böhm model from the Benchmark Collection. Using [amici.petab_import](https://amici.readthedocs.io/en/latest/generated/amici.petab_import.html#amici.petab_import.import_petab_problem), we will load the PEtab problem. To create a [JAXModel](https://amici.readthedocs.io/en/latest/generated/amici.jax.html#amici.jax.JAXModel) instead of a standard AMICI model, we set the `jax` parameter to `True`.\n" + ] + }, + { + "metadata": {}, + "cell_type": "code", + "outputs": [], + "execution_count": null, + "source": [ + "from amici.petab.petab_import import import_petab_problem\n", + "import petab.v1 as petab\n", + "\n", + "# Define the model name and YAML file location\n", + "model_name = \"Boehm_JProteomeRes2014\"\n", + "yaml_url = (\n", + " f\"https://raw.githubusercontent.com/Benchmarking-Initiative/Benchmark-Models-PEtab/\"\n", + " f\"master/Benchmark-Models/{model_name}/{model_name}.yaml\"\n", + ")\n", + "\n", + "# Load the PEtab problem from the YAML file\n", + "petab_problem = petab.Problem.from_yaml(yaml_url)\n", + "\n", + "# Import the PEtab problem as a JAX-compatible AMICI model\n", + "jax_model = import_petab_problem(\n", + " petab_problem,\n", + " verbose=False, # no text output\n", + " jax=True, # return jax model\n", + ")" + ], + "id": "c71c96da0da3144a" + }, + { + "metadata": {}, + "cell_type": "markdown", + "source": [ + "## Simulation\n", + "\n", + "In principle, we can already use this model for simulation using the [simulate_condition](https://amici.readthedocs.io/en/latest/generated/amici.jax.html#amici.jax.JAXModel.simulate_condition) method. However, this approach can be cumbersome as timepoints, data etc. need to be specified manually. Instead, we process the PEtab problem into a [JAXProblem](https://amici.readthedocs.io/en/latest/generated/amici.jax.html#amici.jax.JAXProblem), which enables efficient simulation using [amici.jax.run_simulations]((https://amici.readthedocs.io/en/latest/generated/amici.jax.html#amici.jax.run_simulations)." + ], + "id": "7e0f1c27bd71ee1f" + }, + { + "metadata": {}, + "cell_type": "code", + "outputs": [], + "execution_count": null, + "source": [ + "from amici.jax import JAXProblem, run_simulations\n", + "\n", + "# Create a JAXProblem from the JAX model and PEtab problem\n", + "jax_problem = JAXProblem(jax_model, petab_problem)\n", + "\n", + "# Run simulations and compute the log-likelihood\n", + "llh, results = run_simulations(jax_problem)" + ], + "id": "ccecc9a29acc7b73" + }, + { + "metadata": {}, + "cell_type": "markdown", + "source": "This simulates the model for all conditions using the nominal parameter values. Simple, right? Now, let’s take a look at the simulation results.", + "id": "415962751301c64a" + }, + { + "metadata": {}, + "cell_type": "code", + "outputs": [], + "execution_count": null, + "source": [ + "# Define the simulation condition\n", + "simulation_condition = (\"model1_data1\",)\n", + "\n", + "# Access the results for the specified condition\n", + "results[simulation_condition]" + ], + "id": "596b86e45e18fe3d" + }, + { + "metadata": {}, + "cell_type": "markdown", + "source": [ + "Unfortunately, the simulation failed! As seen in the output, the simulation broke down after the initial timepoint, indicated by the `inf` values in the state variables `results[simulation_condition][1].x` and the `nan` likelihood value. A closer inspection of this variable provides additional clues about what might have gone wrong.\n", + "\n", + "The issue stems from using single precision, as indicated by the `float32` dtype of state variables. Single precision is generally a [bad idea](https://docs.kidger.site/diffrax/examples/stiff_ode/) for stiff systems like the Böhm model. Let’s retry the simulation with double precision." + ], + "id": "a1b173e013f9210a" + }, + { + "metadata": {}, + "cell_type": "code", + "outputs": [], + "execution_count": null, + "source": [ + "import jax\n", + "\n", + "# Enable double precision in JAX\n", + "jax.config.update(\"jax_enable_x64\", True)\n", + "\n", + "# Re-run simulations with double precision\n", + "llh, results = run_simulations(jax_problem)\n", + "\n", + "results" + ], + "id": "f4f5ff705a3f7402" + }, + { + "metadata": {}, + "cell_type": "markdown", + "source": "Success! The simulation completed successfully, and we can now plot the resulting state trajectories.", + "id": "fe4d3b40ee3efdf2" + }, + { + "metadata": {}, + "cell_type": "code", + "outputs": [], + "execution_count": null, + "source": [ + "%matplotlib inline\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "\n", + "\n", + "def plot_simulation(results):\n", + " \"\"\"\n", + " Plot the state trajectories from the simulation results.\n", + "\n", + " Parameters:\n", + " results (dict): Simulation results from run_simulations.\n", + " \"\"\"\n", + " # Extract the simulation results for the specific condition\n", + " sim_results = results[simulation_condition][1]\n", + "\n", + " # Create a new figure for the state trajectories\n", + " plt.figure(figsize=(8, 6))\n", + " for idx in range(sim_results[\"x\"].shape[1]):\n", + " time_points = np.array(sim_results[\"ts\"])\n", + " state_values = np.array(sim_results[\"x\"][:, idx])\n", + " plt.plot(time_points, state_values, label=jax_model.state_ids[idx])\n", + "\n", + " # Add labels, legend, and grid\n", + " plt.xlabel(\"Time\")\n", + " plt.ylabel(\"State Values\")\n", + " plt.title(simulation_condition)\n", + " plt.legend()\n", + " plt.grid(True)\n", + " plt.show()\n", + "\n", + "\n", + "# Plot the simulation results\n", + "plot_simulation(results)" + ], + "id": "72f1ed397105e14a" + }, + { + "metadata": {}, + "cell_type": "markdown", + "source": "`run_simulations` enables users to specify the simulation conditions to be executed. For more complex models, this allows for restricting simulations to a subset of conditions. Since the Böhm model includes only a single condition, we demonstrate this functionality by simulating no condition at all.", + "id": "4fa97c33719c2277" + }, + { + "metadata": {}, + "cell_type": "code", + "outputs": [], + "execution_count": null, + "source": [ + "llh, results = run_simulations(jax_problem, simulation_conditions=tuple())\n", + "results" + ], + "id": "7950774a3e989042" + }, + { + "metadata": {}, + "cell_type": "markdown", + "source": [ + "## Updating Parameters\n", + "\n", + "As next step, we will update the parameter values used for simulation. However, if we attempt to directly modify the values in `JAXModel.parameters`, we encounter a `FrozenInstanceError`." + ], + "id": "98b8516a75ce4d12" + }, + { + "metadata": {}, + "cell_type": "code", + "outputs": [], + "execution_count": null, + "source": [ + "from dataclasses import FrozenInstanceError\n", + "import jax\n", + "\n", + "# Generate random noise to update the parameters\n", + "noise = (\n", + " jax.random.normal(\n", + " key=jax.random.PRNGKey(0), shape=jax_problem.parameters.shape\n", + " )\n", + " / 10\n", + ")\n", + "\n", + "# Attempt to update the parameters\n", + "try:\n", + " jax_problem.parameters += noise\n", + "except FrozenInstanceError as e:\n", + " print(\"Error:\", e)" + ], + "id": "3d278a3d21e709d" + }, + { + "metadata": {}, + "cell_type": "markdown", + "source": [ + "The root cause of this error lies in the fact that, to enable autodiff, direct modifications of attributes are not allowed in [equinox](https://docs.kidger.site/equinox/), which AMICI utilizes under the hood. Consequently, attributes of instances like `JAXModel` or `JAXProblem` cannot be updated directly — this is the price we have to pay for autodiff.\n", + "\n", + "However, `JAXProblem` provides a convenient method called [update_parameters](https://amici.readthedocs.io/en/latest/generated/amici.jax.html#amici.jax.JAXProblem.update_parameters). The caveat is that this method creates a new JAXProblem instance instead of modifying the existing one." + ], + "id": "4cc3d595de4a4085" + }, + { + "metadata": {}, + "cell_type": "code", + "outputs": [], + "execution_count": null, + "source": [ + "# Update the parameters and create a new JAXProblem instance\n", + "jax_problem = jax_problem.update_parameters(jax_problem.parameters + noise)\n", + "\n", + "# Run simulations with the updated parameters\n", + "llh, results = run_simulations(jax_problem)\n", + "\n", + "# Plot the simulation results\n", + "plot_simulation(results)" + ], + "id": "e47748376059628b" + }, + { + "metadata": {}, + "cell_type": "markdown", + "source": [ + "## Computing Gradients\n", + "\n", + "Similar to updating attributes, computing gradients in the JAX ecosystem can feel a bit unconventional if you’re not familiar with the JAX ecosysmt. JAX offers [powerful automatic differentiation](https://jax.readthedocs.io/en/latest/automatic-differentiation.html) through the `jax.grad` function. However, to use `jax.grad` with `JAXProblem`, we need to specify which parts of the `JAXProblem` should be treated as static." + ], + "id": "660baf605a4e8339" + }, + { + "metadata": {}, + "cell_type": "code", + "outputs": [], + "execution_count": null, + "source": [ + "try:\n", + " # Attempt to compute the gradient of the run_simulations function\n", + " jax.grad(run_simulations, has_aux=True)(jax_problem)\n", + "except TypeError as e:\n", + " print(\"Error:\", e)" + ], + "id": "7033d09cc81b7f69" + }, + { + "metadata": {}, + "cell_type": "markdown", + "source": "Fortunately, `equinox` simplifies this process by offering [filter_grad](https://docs.kidger.site/equinox/api/transformations/#equinox.filter_grad), which enables autodiff functionality that is compatible with `JAXProblem` and, in theory, also with `JAXModel`.", + "id": "dc9bc07cde00a926" + }, + { + "metadata": {}, + "cell_type": "code", + "outputs": [], + "execution_count": null, + "source": [ + "import equinox as eqx\n", + "\n", + "# Compute the gradient using equinox's filter_grad, preserving auxiliary outputs\n", + "grad, _ = eqx.filter_grad(run_simulations, has_aux=True)(jax_problem)" + ], + "id": "a6704182200e6438" + }, + { + "metadata": {}, + "cell_type": "markdown", + "source": "Functions transformed by `filter_grad` return gradients that share the same structure as the first argument (unless specified otherwise). This allows us to access the gradient with respect to the parameters attribute directly `via grad.parameters`.", + "id": "851c3ec94cb5d086" + }, + { + "metadata": {}, + "cell_type": "code", + "outputs": [], + "execution_count": null, + "source": "grad.parameters", + "id": "c00c1581d7173d7a" + }, + { + "metadata": {}, + "cell_type": "markdown", + "source": "Attributes for which derivatives cannot be computed (typically anything that is not a [jax.numpy.array](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.array.html)) are automatically set to `None`.", + "id": "375b835fecc5a022" + }, + { + "metadata": {}, + "cell_type": "code", + "outputs": [], + "execution_count": null, + "source": "grad", + "id": "f7c17f7459d0151f" + }, + { + "metadata": {}, + "cell_type": "markdown", + "source": "Observant readers may notice that the gradient above appears to include numeric values for derivatives with respect to some measurements. However, `simulation_conditions` internally disables gradient computations using `jax.lax.stop_gradient`, resulting in these values being zeroed out.", + "id": "8eb7cc3db510c826" + }, + { + "metadata": {}, + "cell_type": "code", + "outputs": [], + "execution_count": null, + "source": "grad._measurements[simulation_condition]", + "id": "3badd4402cf6b8c6" + }, + { + "metadata": {}, + "cell_type": "markdown", + "source": "However, we can compute derivatives with respect to data elements using `JAXModel.simulate_condition`. In the example below, we differentiate the observables `y` (specified by passing `y` to the `ret` argument) with respect to the timepoints at which the model outputs are computed after the solving the differential equation. While this might not be particularly practical, it serves as an nice illustration of the power of automatic differentiation.", + "id": "58eb04393a1463d" + }, + { + "metadata": {}, + "cell_type": "code", + "outputs": [], + "execution_count": null, + "source": [ + "import jax.numpy as jnp\n", + "import diffrax\n", + "\n", + "# Define the simulation condition\n", + "simulation_condition = (\"model1_data1\",)\n", + "\n", + "# Load condition-specific data\n", + "ts_preeq, ts_dyn, ts_posteq, my, iys = jax_problem._measurements[\n", + " simulation_condition\n", + "]\n", + "\n", + "# Load parameters for the specified condition\n", + "p = jax_problem.load_parameters(simulation_condition[0])\n", + "# Disable preequilibration\n", + "p_preeq = jnp.array([])\n", + "\n", + "\n", + "# Define a function to compute the gradient with respect to dynamic timepoints\n", + "@eqx.filter_jacfwd\n", + "def grad_ts_dyn(tt):\n", + " return jax_problem.model.simulate_condition(\n", + " p=p,\n", + " p_preeq=p_preeq,\n", + " ts_preeq=ts_preeq,\n", + " ts_dyn=tt,\n", + " ts_posteq=ts_posteq,\n", + " my=jnp.array(my),\n", + " iys=jnp.array(iys),\n", + " solver=diffrax.Kvaerno5(),\n", + " controller=diffrax.PIDController(atol=1e-8, rtol=1e-8),\n", + " max_steps=2**10,\n", + " adjoint=diffrax.DirectAdjoint(),\n", + " ret=\"y\", # Return observables\n", + " )[0]\n", + "\n", + "\n", + "# Compute the gradient with respect to `ts_dyn`\n", + "g = grad_ts_dyn(ts_dyn)\n", + "g" + ], + "id": "1a91aff44b93157" + }, + { + "metadata": {}, + "cell_type": "markdown", + "source": [ + "## Compilation & Profiling\n", + "\n", + "To maximize performance with JAX, code should be just-in-time (JIT) compiled. This can be achieved using the `jax.jit` or `equinox.filter_jit` decorators. While JIT compilation introduces some overhead during the first function call, it significantly improves performance for subsequent calls. To demonstrate this, we will first clear the JIT cache and then profile the execution." + ], + "id": "9f870da7754e139c" + }, + { + "metadata": {}, + "cell_type": "code", + "outputs": [], + "execution_count": null, + "source": [ + "from time import time\n", + "\n", + "# Clear JAX caches to ensure a fresh start\n", + "jax.clear_caches()\n", + "\n", + "# Define a JIT-compiled gradient function with auxiliary outputs\n", + "gradfun = eqx.filter_jit(eqx.filter_grad(run_simulations, has_aux=True))" + ], + "id": "58ebdc110ea7457e" + }, + { + "metadata": {}, + "cell_type": "code", + "outputs": [], + "execution_count": null, + "source": [ + "# Measure the time taken for the first function call (including compilation)\n", + "start = time()\n", + "run_simulations(jax_problem)\n", + "print(f\"Function compilation time: {time() - start:.2f} seconds\")\n", + "\n", + "# Measure the time taken for the gradient computation (including compilation)\n", + "start = time()\n", + "gradfun(jax_problem)\n", + "print(f\"Gradient compilation time: {time() - start:.2f} seconds\")" + ], + "id": "e1242075f7e0faf" + }, + { + "metadata": {}, + "cell_type": "code", + "outputs": [], + "execution_count": null, + "source": [ + "%%timeit\n", + "run_simulations(\n", + " jax_problem,\n", + " controller=diffrax.PIDController(\n", + " rtol=1e-8, # same as amici default\n", + " atol=1e-16, # same as amici default\n", + " pcoeff=0.4, # recommended value for stiff systems\n", + " icoeff=0.3, # recommended value for stiff systems\n", + " dcoeff=0.0, # recommended value for stiff systems\n", + " ),\n", + ")" + ], + "id": "27181f367ccb1817" + }, + { + "metadata": {}, + "cell_type": "code", + "outputs": [], + "execution_count": null, + "source": [ + "%%timeit \n", + "gradfun(\n", + " jax_problem,\n", + " controller=diffrax.PIDController(\n", + " rtol=1e-8, # same as amici default\n", + " atol=1e-16, # same as amici default\n", + " pcoeff=0.4, # recommended value for stiff systems\n", + " icoeff=0.3, # recommended value for stiff systems\n", + " dcoeff=0.0, # recommended value for stiff systems\n", + " ),\n", + ")" + ], + "id": "5b8d3a6162a3ae55" + }, + { + "metadata": {}, + "cell_type": "code", + "outputs": [], + "execution_count": null, + "source": [ + "from amici.petab import simulate_petab\n", + "import amici\n", + "\n", + "# Import the PEtab problem as a standard AMICI model\n", + "amici_model = import_petab_problem(\n", + " petab_problem,\n", + " verbose=False,\n", + " compile_=True,\n", + " jax=False, # load the amici model this time\n", + ")\n", + "\n", + "# Configure the solver with appropriate tolerances\n", + "solver = amici_model.getSolver()\n", + "solver.setAbsoluteTolerance(1e-8)\n", + "solver.setRelativeTolerance(1e-8)\n", + "\n", + "# Prepare the parameters for the simulation\n", + "problem_parameters = dict(\n", + " zip(jax_problem.parameter_ids, jax_problem.parameters)\n", + ")" + ], + "id": "d733a450635a749b" + }, + { + "cell_type": "code", + "execution_count": 20, + "id": "413ed7c60b2cf4be", + "metadata": { + "ExecuteTime": { + "end_time": "2024-11-19T09:51:55.259985Z", + "start_time": "2024-11-19T09:51:55.257937Z" + } + }, + "outputs": [], + "source": [ + "# Profile simulation only\n", + "solver.setSensitivityOrder(amici.SensitivityOrder.none)" + ] + }, + { + "cell_type": "code", + "execution_count": 21, + "id": "768fa60e439ca8b4", + "metadata": { + "ExecuteTime": { + "end_time": "2024-11-19T09:51:57.417608Z", + "start_time": "2024-11-19T09:51:55.273367Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "26.1 ms ± 2.71 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)\n" + ] + } + ], + "source": [ + "%%timeit \n", + "simulate_petab(\n", + " petab_problem,\n", + " amici_model,\n", + " solver=solver,\n", + " problem_parameters=problem_parameters,\n", + " scaled_parameters=True,\n", + " scaled_gradients=True,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 22, + "id": "b8382b0b2b68f49e", + "metadata": { + "ExecuteTime": { + "end_time": "2024-11-19T09:51:57.497361Z", + "start_time": "2024-11-19T09:51:57.494502Z" + } + }, + "outputs": [], + "source": [ + "# Profile gradient computation using forward sensitivity analysis\n", + "solver.setSensitivityOrder(amici.SensitivityOrder.first)\n", + "solver.setSensitivityMethod(amici.SensitivityMethod.forward)" + ] + }, + { + "cell_type": "code", + "execution_count": 23, + "id": "3bae1fab8c416122", + "metadata": { + "ExecuteTime": { + "end_time": "2024-11-19T09:51:59.897459Z", + "start_time": "2024-11-19T09:51:57.511889Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "29.1 ms ± 1.82 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)\n" + ] + } + ], + "source": [ + "%%timeit \n", + "simulate_petab(\n", + " petab_problem,\n", + " amici_model,\n", + " solver=solver,\n", + " problem_parameters=problem_parameters,\n", + " scaled_parameters=True,\n", + " scaled_gradients=True,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 24, + "id": "71e0358227e1dc74", + "metadata": { + "ExecuteTime": { + "end_time": "2024-11-19T09:51:59.972149Z", + "start_time": "2024-11-19T09:51:59.969006Z" + } + }, + "outputs": [], + "source": [ + "# Profile gradient computation using adjoint sensitivity analysis\n", + "solver.setSensitivityOrder(amici.SensitivityOrder.first)\n", + "solver.setSensitivityMethod(amici.SensitivityMethod.adjoint)" + ] + }, + { + "cell_type": "code", + "execution_count": 25, + "id": "e3cc7971002b6d06", + "metadata": { + "ExecuteTime": { + "end_time": "2024-11-19T09:52:03.266074Z", + "start_time": "2024-11-19T09:51:59.992465Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "39.3 ms ± 1.6 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)\n" + ] + } + ], + "source": [ + "%%timeit \n", + "simulate_petab(\n", + " petab_problem,\n", + " amici_model,\n", + " solver=solver,\n", + " problem_parameters=problem_parameters,\n", + " scaled_parameters=True,\n", + " scaled_gradients=True,\n", + ")" + ] + } + ], + "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.13.0" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}