{ "cells": [ { "cell_type": "markdown", "id": "5f668bef-04b2-4ecd-a195-1f107ef83c07", "metadata": {}, "source": [ "# Getting Started" ] }, { "cell_type": "markdown", "id": "c1fbd084-6c5b-46c5-add0-7210f175aeb3", "metadata": {}, "source": [ "## Overview\n", "\n", "This getting started tutorial outlines how to set up and simulate a PEtab SciML problem using [AMICI](https://amici.readthedocs.io/en/latest/index.html). Some familiarity with the PEtab format is assumed, see [PEtab](https://petab.readthedocs.io/en/latest/v1/documentation_data_format.html) for a refresher.\n", "The environment and example PEtab files to run this notebook are provided in the PEtab SciML repo." ] }, { "cell_type": "markdown", "id": "6a17e84a-04b6-449c-81e7-9a686723a062", "metadata": {}, "source": [ "As an example, we will use the Lotka-Voltera system:\n", "\n", "$$\\frac{\\mathrm{d} \\text{prey}}{\\mathrm{d} t} = \\alpha \\cdot \\text{prey} - \\beta \\cdot \\text{prey} \\cdot \\text{predator}$$\n", "\n", "$$\\frac{\\mathrm{d} \\text{predator}}{\\mathrm{d} t} = \\gamma \\cdot \\text{prey} \\cdot \\text{predator} - \\delta \\cdot \\text{predator}$$\n", "\n", "We will replace the two interactive terms ($\\beta$ and $\\gamma$) with outputs from a neural network. \n", "The ``prey`` and ``predator`` variables in the model will be used as the inputs to this network.\n", "This example follows the Universal Differential Equation (UDE) case in [Rackauckas et al. 2020](https://arxiv.org/pdf/2001.04385).\n", "\n", "$$\\frac{\\mathrm{d} \\text{prey}}{\\mathrm{d} t} = \\alpha \\cdot \\text{prey} - \\text{NN}(\\text{prey}, \\text{predator})[0]$$\n", "\n", "$$\\frac{\\mathrm{d} \\text{predator}}{\\mathrm{d} t} = \\text{NN}(\\text{prey}, \\text{predator})[1] - \\delta \\cdot \\text{predator}$$" ] }, { "cell_type": "markdown", "id": "6a58aaaf-aae5-46c8-b657-aa3a257288e3", "metadata": {}, "source": [ "## Environment\n", "\n", "Support for PEtab SciML models is under active development. A ``requirements.txt`` file is provided with these docs which checks out branches of the external dependencies where support has been implemented." ] }, { "cell_type": "markdown", "id": "c00cbb82-588e-433c-9b2b-5b8240ee252b", "metadata": {}, "source": [ "## Loading the PEtab problem\n", "\n", "We will start by loading the PEtab problem using the PEtab Python library, so that we can interactively explore the files that are used to build and simulate the problem." ] }, { "cell_type": "code", "execution_count": null, "id": "54ca5094-de98-4eb7-9340-d82212266dad", "metadata": {}, "outputs": [], "source": [ "from petab.v2 import Problem\n", "\n", "petab_problem = Problem.from_yaml(\"problem.yaml\")" ] }, { "cell_type": "markdown", "id": "8201b3ab-6c90-4a03-807f-91aa079e76f9", "metadata": {}, "source": [ "The files needed to define a PEtab SciML problem are:\n", "- model file\n", "- neural network file\n", "- mapping table\n", "- hybridization table\n", "- parameters table\n", "- network array inputs\n", "- measurements table\n", "- observables table\n", "- conditions table\n", "- YAML file\n", "\n", "The hybridization, network YAML and network input files are new and specific to PEtab SciML problems. The other files already exist in the PEtab format but have been expanded in order to support PEtab SciML problems." ] }, { "cell_type": "markdown", "id": "89d0f00c-daef-47cb-a5db-7d2b2566f5da", "metadata": {}, "source": [ "### Model File\n", "\n", "PEtab SciML importers accept models in various exchangeable standard formats (e.g., SBML, CellML, BioNetGen). For our example we use an SBML model because this standard is widely supported across PEtab importers. The SBML model will need modifying because outputs from the neural network can only be assigned to parameters in the PEtab problem. This means the ``prey`` and ``predator`` species will need to be removed from the expressions we want to parameterise, i.e.\n", "\n", "$$\\frac{\\mathrm{d} \\text{prey}}{\\mathrm{d} t} = \\alpha \\cdot \\text{prey} - \\beta$$\n", "\n", "$$\\frac{\\mathrm{d} \\text{predator}}{\\mathrm{d} t} = \\gamma - \\delta \\cdot \\text{predator}$$\n", "\n", "Whichever model format is used, the idea is that parameters in the model are replaced by network elements. There are a large number of tools available for creating SBML model files like these, such as GUI based [COPASI](https://copasi.org/). " ] }, { "cell_type": "markdown", "id": "6c6bea9d-0f3d-43bc-ad9b-f1fb7a2e18d8", "metadata": {}, "source": [ "### Neural Network File\n", "\n", "This file defines the network architecture, and should be provided in the PEtab SciML associated network YAML format. You can generate a file like this by defining a neural network in PyTorch and exporting it to the SciML YAML format. The provided example defines a network with three ``Linear`` layers and a ``tanh`` activation function. See the page on supported layers and activation functions for a complete list." ] }, { "cell_type": "code", "execution_count": 4, "id": "f1a5301d-ee1e-4091-9ec7-755f6c36c0e8", "metadata": {}, "outputs": [], "source": [ "from petab_sciml.standard.nn_model import Input, NNModel, NNModelStandard\n", "import torch\n", "from torch import nn\n", "import torch.nn.functional as F\n", "\n", "class NeuralNetwork(nn.Module):\n", " def __init__(self):\n", " super().__init__()\n", " self.layer1 = torch.nn.Linear(2, 5)\n", " self.layer2 = torch.nn.Linear(5, 5)\n", " self.layer3 = torch.nn.Linear(5, 2)\n", "\n", " def forward(self, net_input):\n", " x = self.layer1(net_input)\n", " x = F.tanh(x)\n", " x = self.layer2(x)\n", " x = F.tanh(x)\n", " x = self.layer3(x)\n", " return x\n", "\n", "net1 = NeuralNetwork()\n", "nn_model1 = NNModel.from_pytorch_module(\n", " module=net1, nn_model_id=\"net1\", inputs=[Input(input_id=\"input0\")]\n", ")\n", "NNModelStandard.save_data(\n", " data=nn_model1, filename=\"net1_from_pytorch.yaml\"\n", ")" ] }, { "cell_type": "markdown", "id": "a580250b-0eb9-4502-a9bb-5b2fc514cbc1", "metadata": {}, "source": [ "Note that where any specific network layers or parameters are referenced in the ``mapping.tsv``, it should refer to them by the layer ids in this file." ] }, { "cell_type": "markdown", "id": "c94ad071-2c03-4be7-9a1a-f9bcefab7d4d", "metadata": {}, "source": [ "### Mapping Table\n", "\n", "The purpose of the mapping table is to define PEtab compatible identifiers for model entities that would otherwise not be valid PEtab. Inputs, outputs and parameters of our neural network are all examples of model entities that would not be valid PEtab identifiers, so we define PEtab ids for these in this file.\n", "\n", "Take the first row for example. The valid PEtab identifier ``net1_input1`` is mapped to the model id ``net1.inputs[0][0]``. This model id denotes the first index of the first input to the neural network (note that zero based indexing is used).\n", "\n", "``net1.parameters`` refers to all the trainable parameters in the network i.e. weights and biases of all the layers." ] }, { "cell_type": "code", "execution_count": 5, "id": "a6631bb7-f198-4ac7-8c28-8c076d2fdc74", "metadata": {}, "outputs": [ { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
modelEntityId
petabEntityId
net1_input1net1.inputs[0][0]
net1_input2net1.inputs[0][1]
net1_output1net1.outputs[0][0]
net1_output2net1.outputs[0][1]
net1_psnet1.parameters
\n", "
" ], "text/plain": [ " modelEntityId\n", "petabEntityId \n", "net1_input1 net1.inputs[0][0]\n", "net1_input2 net1.inputs[0][1]\n", "net1_output1 net1.outputs[0][0]\n", "net1_output2 net1.outputs[0][1]\n", "net1_ps net1.parameters" ] }, "execution_count": 5, "metadata": {}, "output_type": "execute_result" } ], "source": [ "petab_problem.mapping_df" ] }, { "cell_type": "markdown", "id": "31a5da7b-cb0b-4f50-9975-4c15f57a84cd", "metadata": {}, "source": [ "### Hybridization Table\n", "\n", "The hybridization table defines where our neural network inputs and outputs get used in the model. As these hybridization elements are integral to the construction of the model, the supporting tools add hybridization information when the model is defined. At this point we will therefore build the model from the PEtab problem in order to demonstrate how the neural network inputs and outputs are inserted into the ODE system." ] }, { "cell_type": "code", "execution_count": null, "id": "c344a871-395e-4aa8-b053-d1afc7916707", "metadata": {}, "outputs": [], "source": [ "from amici.petab import import_petab_problem\n", "from amici.jax import (\n", " JAXProblem,\n", " run_simulations,\n", ")\n", "\n", "# Create AMICI model for the petab problem\n", "jax_model = import_petab_problem(\n", " petab_problem,\n", " model_output_dir=\"model\",\n", " compile_=True,\n", " jax=True\n", ")\n", "\n", "# Create the JAXProblem - wrapper for the AMICI model\n", "jax_problem = JAXProblem(jax_model, petab_problem)" ] }, { "cell_type": "code", "execution_count": 7, "id": "4169bb35-28dc-4a6e-b23b-d188775a1d3d", "metadata": {}, "outputs": [ { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
targetValue
targetId
net1_input1prey
net1_input2predator
betanet1_output1
gammanet1_output2
\n", "
" ], "text/plain": [ " targetValue\n", "targetId \n", "net1_input1 prey\n", "net1_input2 predator\n", "beta net1_output1\n", "gamma net1_output2" ] }, "execution_count": 7, "metadata": {}, "output_type": "execute_result" } ], "source": [ "jax_problem._petab_problem.hybridization_df" ] }, { "cell_type": "markdown", "id": "b4b478b7-2bfd-401a-9454-522d51e6583b", "metadata": {}, "source": [ "In the example, the inputs to the network are the ``prey`` and ``predator`` values and the outputs of the network provide $\\beta$ and $\\gamma$. These substitutions hold for all simulation conditions. Note that we refer to the inputs and outputs of the network by their PEtab identifiers, as defined in the mapping file." ] }, { "cell_type": "markdown", "id": "3ebb97b5-878b-4d31-87ea-6d508801078a", "metadata": {}, "source": [ "### Parameters Table\n", "\n", "This file contains configuration information on all model parameters, which in our case, includes parameters of the neural network. Actual values of the network parameters should be provided in a separate HDF5 array file (see more details in the next section). The parameters file specifies the scale, bounds and nominal values of all the parameters. We want our neural network parameters to be freely optimised during training, so we set the bounds to ``[-inf, inf]`` and leave the nominal value blank (translates to a NaN when loaded into the PEtab problem)." ] }, { "cell_type": "code", "execution_count": 8, "id": "6171d8ed-a485-49a8-b822-068fd70d8559", "metadata": {}, "outputs": [ { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
parameterScalelowerBoundupperBoundnominalValueestimate
parameterId
alphalin0.015.01.31
deltalin0.015.01.81
net1_pslin-infinfNaN1
\n", "
" ], "text/plain": [ " parameterScale lowerBound upperBound nominalValue estimate\n", "parameterId \n", "alpha lin 0.0 15.0 1.3 1\n", "delta lin 0.0 15.0 1.8 1\n", "net1_ps lin -inf inf NaN 1" ] }, "execution_count": 8, "metadata": {}, "output_type": "execute_result" } ], "source": [ "petab_problem.parameter_df" ] }, { "cell_type": "markdown", "id": "376d7687-24dc-484b-820e-44a7a28a0a6e", "metadata": { "jp-MarkdownHeadingCollapsed": true }, "source": [ "### Array inputs\n", "\n", "In this example, the values of parameters for the neural network are provided in HDF5 format.\n", "The structure inside the HDF5 file is as follows.\n", "\n", "```\n", " arrays.hdf5\n", " ├── metadata\n", " │ └── perm \n", " └── parameters\n", " └── net1\n", " ├── layer1 \n", " │ ├── weight\n", " │ └── bias\n", " ├── layer2\n", " │ ├── weight\n", " │ └── bias\n", " └── layer3\n", " ├── weight\n", " └── bias\n", "```\n", "\n", "Note the network id ``net1`` matches the identifier in the network YAML file and the ``modelEntityId`` in the ``mapping.tsv``. Likewise the layer names (``layer1``, ``layer2``, ``layer3``) in the HDF5 are identifiers and must match the layer names in the network YAML and mapping files." ] }, { "cell_type": "markdown", "id": "b8097112-5f96-4c2a-822a-5c56647ccf14", "metadata": {}, "source": [ "### Measurements, Observables and Conditions Tables\n", "\n", "These tables are unchanged from a regular PEtab probem. For completeness, the measurement table contains the measurement data, which experimental condition they where collected, and which model entities each measurement maps to." ] }, { "cell_type": "code", "execution_count": 9, "id": "e119d605-964e-4603-b120-47ce367c3825", "metadata": {}, "outputs": [ { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
observableIdsimulationConditionIdmeasurementtime
0prey_ocond10.1730171.0
1prey_ocond10.4891772.0
2prey_ocond11.6439963.0
3prey_ocond15.4519634.0
4prey_ocond12.9775225.0
5prey_ocond10.1816636.0
6prey_ocond10.3481127.0
7prey_ocond10.9379198.0
8prey_ocond13.1132409.0
9prey_ocond18.86393310.0
10predator_ocond10.8474161.0
11predator_ocond10.2111352.0
12predator_ocond1-0.0250543.0
13predator_ocond10.1250104.0
14predator_ocond16.7004555.0
15predator_ocond12.0071586.0
16predator_ocond10.4200927.0
17predator_ocond10.0480328.0
18predator_ocond10.1286699.0
19predator_ocond11.19278410.0
\n", "
" ], "text/plain": [ " observableId simulationConditionId measurement time\n", "0 prey_o cond1 0.173017 1.0\n", "1 prey_o cond1 0.489177 2.0\n", "2 prey_o cond1 1.643996 3.0\n", "3 prey_o cond1 5.451963 4.0\n", "4 prey_o cond1 2.977522 5.0\n", "5 prey_o cond1 0.181663 6.0\n", "6 prey_o cond1 0.348112 7.0\n", "7 prey_o cond1 0.937919 8.0\n", "8 prey_o cond1 3.113240 9.0\n", "9 prey_o cond1 8.863933 10.0\n", "10 predator_o cond1 0.847416 1.0\n", "11 predator_o cond1 0.211135 2.0\n", "12 predator_o cond1 -0.025054 3.0\n", "13 predator_o cond1 0.125010 4.0\n", "14 predator_o cond1 6.700455 5.0\n", "15 predator_o cond1 2.007158 6.0\n", "16 predator_o cond1 0.420092 7.0\n", "17 predator_o cond1 0.048032 8.0\n", "18 predator_o cond1 0.128669 9.0\n", "19 predator_o cond1 1.192784 10.0" ] }, "execution_count": 9, "metadata": {}, "output_type": "execute_result" } ], "source": [ "petab_problem.measurement_df" ] }, { "cell_type": "markdown", "id": "fac45ce4-45ee-4b3b-82fa-f5f403039439", "metadata": {}, "source": [ "The observables table links the model entities to the measured values." ] }, { "cell_type": "code", "execution_count": 10, "id": "4984f33e-2373-4614-9198-e03398276f95", "metadata": {}, "outputs": [ { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
observableFormulanoiseFormulaobservableTransformationnoiseDistribution
observableId
prey_oprey0.05linnormal
predator_opredator0.05linnormal
\n", "
" ], "text/plain": [ " observableFormula noiseFormula observableTransformation \\\n", "observableId \n", "prey_o prey 0.05 lin \n", "predator_o predator 0.05 lin \n", "\n", " noiseDistribution \n", "observableId \n", "prey_o normal \n", "predator_o normal " ] }, "execution_count": 10, "metadata": {}, "output_type": "execute_result" } ], "source": [ "petab_problem.observable_df" ] }, { "cell_type": "markdown", "id": "613e0af7-41e0-4d74-9f55-579b31872f55", "metadata": {}, "source": [ "The conditions table in intended to fully describe the experimental conditions the measurements were collected under. It can have any number of columns, but in this example there is only one experimental condition and no additional information about it." ] }, { "cell_type": "code", "execution_count": 11, "id": "ef2136e1-9bdb-4c20-a6d2-9388f33a0fa0", "metadata": {}, "outputs": [ { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
conditionId
cond1
\n", "
" ], "text/plain": [ "Empty DataFrame\n", "Columns: []\n", "Index: [cond1]" ] }, "execution_count": 11, "metadata": {}, "output_type": "execute_result" } ], "source": [ "petab_problem.condition_df" ] }, { "cell_type": "markdown", "id": "b60ce1a4-87a0-4be6-92f6-dbd74fd39e2f", "metadata": {}, "source": [ "### YAML File\n", "\n", "Finally, the ``problem.yaml`` file brings everything together by describing how the problem should be constructed from the PEtab files. The ``extensions`` section contains the details of the SciML elements of our problem.\n", "\n", "```yaml\n", " format_version: 2\n", " parameter_file: \"parameters.tsv\"\n", " problems:\n", " - model_files:\n", " lv:\n", " location: \"lv.xml\"\n", " language: \"sbml\"\n", " measurement_files:\n", " - \"measurements.tsv\"\n", " observable_files:\n", " - \"observables.tsv\"\n", " condition_files:\n", " - \"conditions.tsv\"\n", " mapping_files:\n", " - \"mapping.tsv\"\n", " extensions:\n", " sciml:\n", " array_files:\n", " - \"net1_ps.hdf5\"\n", " hybridization_files:\n", " - \"hybridization.tsv\"\n", " neural_nets:\n", " net1:\n", " location: \"net1.yaml\"\n", " static: false\n", " format: \"YAML\"\n", "```\n", "\n", "Where the network is referenced in the ``mapping.tsv`` file, the ``modelEntityId`` must match the network id listed in this YAML, ``net1`` in this case." ] }, { "cell_type": "markdown", "id": "e80d4679-3ac3-47e5-afb0-165c1eb3f96f", "metadata": {}, "source": [ "## Running Simulations\n", "\n", "We are now ready to run a simulation using the jax problem. AMICI's `run_simulations` will simulate all conditions provided. In this example though, there is only one. See the [AMICI docs](https://amici.readthedocs.io/en/latest/python_examples.html) for more examples such as training loops." ] }, { "cell_type": "code", "execution_count": 12, "id": "47eebce8-9cbe-4970-8ac6-50f07a5a8498", "metadata": {}, "outputs": [], "source": [ "import equinox as eqx\n", "\n", "# Run simulations - results in llh - metrics in r\n", "llh, r = run_simulations(jax_problem)\n", "\n", "# Run simulations in gradient mode - gradients given in sllh\n", "sllh, rgrad = eqx.filter_grad(run_simulations, has_aux=True)(jax_problem)" ] }, { "cell_type": "markdown", "id": "febc0864-5602-4339-9a23-c123ce126534", "metadata": {}, "source": [ "## Next Steps\n", "\n", "This tutorial has outlined how to set up a PEtab SciML problem where a neural network appears in the ODE equations (creating what is sometimes referred to as an UDE). \n", "The PEtab SciML format also supports others ways of formulating a problem, which you can find examples of in the how-to-guides. These include:\n", "- neural ODEs (all ODE equations are replaced with a neural network)\n", "- using networks in the observable function which links simulated values to measurements\n", "- having a neural network upstream of the ODE mapping high-dimensional inputs (e.g., images) to ODE model parameters.\n", "\n", "We implemented a simple network in this tutorial for demonstration purposes but the PEtab SciML format supports a range of standard neural network layers and activation functions based on the PyTorch framework. The dedicated docs page lists these layers and functions, along with the tools that support each.\n", "\n", "For a complete description of the PEtab SciML format visit the [format specification](https://petab.readthedocs.io/en/latest/v1/documentation_data_format.html) page and for more information about more complex problem specifications supported by the PEtab format (e.g., models with partial observabillity, multiple experimental/simulations conditions, events), see the [PEtab documentation](https://petab.readthedocs.io/en/latest/)." ] } ], "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 }