Getting Started#

This introductory tutorial shows how to set up a PEtab SciML problem using AMICI. It walks through the main PEtab SciML problem files and focuses on creating a problem where a neural network (NN) enters the model dynamics, which is often called a universal differential equation (UDE) (also referred to as a grey-box model or hybrid Neural ODE). Familiarity with the PEtab v2 format is assumed; see the PEtab tutorial.

The tutorial is provided as a notebook, available here along with a requirements.txt listing Python dependencies. The corresponding PEtab SciML problem files can be downloaded here.

Input problem#

As a working example, the Lotka–Volterra system is considered:

\[\begin{split}\begin{aligned} \frac{\mathrm{d}\,\text{prey}}{\mathrm{d}t} &= \text{alpha} \cdot \text{prey} - \text{beta} \cdot \text{prey} \cdot \text{predator}, \\ \frac{\mathrm{d}\,\text{predator}}{\mathrm{d}t} &= \text{gamma} \cdot \text{prey} \cdot \text{predator} - \text{delta} \cdot \text{predator}. \end{aligned}\end{split}\]

Following Rackauckas et al. [1], to create a UDE, the interaction terms are replaced by the outputs of a NN that takes the states prey and predator as input:

\[\begin{split}\begin{aligned} \frac{\mathrm{d}\,\text{prey}}{\mathrm{d}t} &= \text{alpha} \cdot \text{prey} - \mathrm{NN}(\text{prey}, \text{predator}; \text{theta})[0], \\ \frac{\mathrm{d}\,\text{predator}}{\mathrm{d}t} &= \mathrm{NN}(\text{prey}, \text{predator}; \text{theta})[1] - \text{delta} \cdot \text{predator}. \end{aligned}\end{split}\]

Here, \(\mathrm{NN}(\cdot; \text{theta})[i]\) denotes the i-th output of the neural network.

Given time-series measurements of the model states prey and predator, the goal of this tutorial is to set up a PEtab-SciML problem to estimate both mechanistic parameters (alpha, delta) and neural-network parameters (theta) from these measurements.

PEtab SciML problem files#

A PEtab SciML problem is defined by a set of tabular TSV files, a model file, NN file(s), optional array data files, and a problem YAML file that ties everything together. Compared to PEtab, PEtab SciML adds a small number of files to describe the ML models and how they interact with the mechanistic model. Files specific to PEtab SciML are:

neural network file(s)
hybridization table
array data files

Files inherited from PEtab:

model file
conditions table
experiment table
mapping table
measurements table
observables table
parameters table
YAML file

To make the format concrete, we will load the example problem (via AMICI) and then walk through the files, focusing on the PEtab SciML-specific additions.

[1]:

from petab.v2 import Problem

petab_problem = Problem.from_yaml("problem.yaml")

Model file for a UDE#

PEtab v2 (and, by extension, PEtab SciML) accepts mechanistic models in common exchange formats (e.g. SBML, CellML, BioNetGen). In this tutorial, an SBML model is used since it is widely supported across PEtab SciML importers.

In PEtab SciML, NN outputs are linked to the mechanistic model by assigning them to parameters in the model file. Therefore, the parts of the equations to be learned must be represented as parameters. In this example, the interaction terms are replaced by the parameters beta and gamma, which are later mapped to the network outputs to form a UDE, so the model file corresponds to:

\[\begin{split}\begin{aligned} \frac{\mathrm{d}\,\text{prey}}{\mathrm{d}t} &= \alpha \cdot \text{prey} - \text{beta}, \\ \frac{\mathrm{d}\,\text{predator}}{\mathrm{d}t} &= \gamma - \delta \cdot \text{predator}. \end{aligned}\end{split}\]

In the hybridization table (see below), beta and gamma are then assigned the outputs of a NN.

Neural network file#

PEtab SciML defines a programming-language-independent YAML format for storing NN models. This file is most easily generated by defining the network as a PyTorch module and exporting it to the PEtab SciML YAML representation.

For our working example, we will use a simple feed-forward network with a single hidden layer and a tanh activation function. It can be defined and exported as:

[2]:

from petab_sciml.standard.nn_model import Input, NNModel, NNModelStandard
import torch
import torch.nn.functional as F


class NeuralNetwork(torch.nn.Module):
    """
    Neural network for the Lotka–Volterra UDE example.

    Notes
    -----
    For PEtab SciML export, the model must:
    - define its layers in `__init__`
    - implement `forward` (to define the execution trace).
    """

    def __init__(self):
        """Define all layers of the network."""
        super().__init__()
        self.layer1 = torch.nn.Linear(2, 5)
        self.layer2 = torch.nn.Linear(5, 5)
        self.layer3 = torch.nn.Linear(5, 2)

    def forward(self, net_input):
        """Forward execution trace."""
        x = self.layer1(net_input)
        x = F.tanh(x)
        x = self.layer2(x)
        x = F.tanh(x)
        x = self.layer3(x)
        return x


# Export using NNModel from the petab_sciml library
net1 = NeuralNetwork()
nn_model1 = NNModel.from_pytorch_module(
    module=net1, nn_model_id="net1", inputs=[Input(input_id="input0")]
)
NNModelStandard.save_data(data=nn_model1, filename="net1_from_pytorch.yaml")

Here, nn_model_id is the unique NN model ID, which is used throughout the PEtab SciML problem to refer to this NN (e.g. in the mapping table and problem yaml file).

Mapping table#

The mapping table links NN entities (inputs, outputs, and parameters) to PEtab identifiers. These PEtab IDs can then be referenced in other PEtab SciML tables where appropriate. This step is needed because raw network entity names such as net1.inputs[0][0] are not valid PEtab IDs.

For the working example, the mapping table is:

[3]:

petab_problem.mapping_df

[3]:

	modelEntityId	name	index
petabEntityId
net1_input1	net1.inputs[0][0]	None	0
net1_input2	net1.inputs[0][1]	None	1
net1_output1	net1.outputs[0][0]	None	2
net1_output2	net1.outputs[0][1]	None	3
net1_ps	net1.parameters	None	4

For example, net1_input1 is mapped to net1.inputs[0][0], where the first index selects the input argument (the network forward function can have multiple input arguments), and the second index selects the entry within that input vector. The same indexing convention is used for outputs. Meanwhile, net1.parameters refers to all trainable parameters in the network (weights and biases across all layers). It is also possible to map subsets, for example net1.parameters[layer1], which can then be referenced (e.g. for priors) in the parameter table.

Hybridization table#

The hybridization table specifies where NN inputs and outputs are used in the mechanistic model (e.g. which parameters are replaced by network outputs). To make this concrete, let’s build the model from the PEtab problem and inspect how the NN inputs and outputs are hybridized:

[4]:

from amici.importers.petab import PetabImporter
from amici.sim.jax import run_simulations

# Create a simulator for the ODE and NN models.
petab_importer = PetabImporter(
    petab_problem, output_dir="model", compile_=True, jax=True
)

# The JAXProblem wraps the AMICI model together with the PEtab problem and
# handles additional simulation information (e.g. multiple experiments).
jax_problem = petab_importer.create_simulator()

The hybridization table can be inspected with:

[5]:

petab_problem.extensions.sciml.hybridization_df

[5]:

	targetId	targetValue
0	net1_input1	prey
1	net1_input2	predator
2	beta	net1_output1
3	gamma	net1_output2

From the output, we see the input variables defined in the mapping table are assigned the values of the model states prey and predator. The network outputs are mapped to the model parameters beta and gamma. Effectively, when importing the PEtab SciML problem the mechanistic model is constructed such that beta and gamma are provided by a NN evaluated on prey and predator, yielding the desired UDE.

Parameters table#

In PEtab SciML, the parameters table also lists NN parameters. For each network parameter ID defined in the mapping table, it specifies bounds and (optionally) nominal values. The actual numerical values used to initialize the network parameters are stored separately in an HDF5 array file (see the next section):

[6]:

petab_problem.parameter_df

[6]:

	lowerBound	upperBound	nominalValue	estimate	index
parameterId
alpha	0.0	15.0	1.3	true	0
delta	0.0	15.0	1.8	true	1
net1_ps	-inf	inf	array	true	2

In this example, NN parameters are left unconstrained by using bounds (-inf, inf). The nominal value is set to the reserved keyword array, meaning the values used to initialize the network parameters are taken from the network parameter array file (see the next section).

Array data#

PEtab TSV tables are not well-suited for storing large array-valued data (e.g. NN parameters). PEtab SciML therefore stores such parameter data in an HDF5 file with the following structure

arrays.hdf5
├── metadata
│   └── pytorch_format
└── parameters
    └── net1
        ├── layer1
        │   ├── weight
        │   └── bias
        ├── layer2
        │   ├── weight
        │   └── bias
        └── layer3
            ├── weight
            └── bias

The network ID (net1) must match the identifier used in the PEtab SciML YAML file. Likewise, layer names (layer1, layer2, layer3) must match the layer names in the NN YAML file.

An array file with the correct parameter structure can most easily be generated (with random initialization) from the NN YAML file using the petab_sciml library:

[7]:

from petab_sciml.standard.array_data import (
    ArrayData,
    ArrayDataStandard,
    extract_nn_yaml_parameters,
)

# Generates array file with parameters in the required structure
net_ps = extract_nn_yaml_parameters("net1_from_pytorch.yaml")
array_data = ArrayData.model_validate(net_ps)
ArrayDataStandard.save_data(array_data, filename="net1_ps_tmp.hdf5")

Measurements, observables, experiment and conditions tables#

These tables are unchanged from standard PEtab. Briefly, the measurements table contains the measurement data, the experiment under which each measurement was collected, and the observable (observableId) that links the data to model output:

[8]:

petab_problem.measurement_df

[8]:

	modelId	observableId	experimentId	time	measurement	index
0	None	prey_o	exp1	1.0	0.173017	0
1	None	prey_o	exp1	2.0	0.489177	1
2	None	prey_o	exp1	3.0	1.643996	2
3	None	prey_o	exp1	4.0	5.451963	3
4	None	prey_o	exp1	5.0	2.977522	4
5	None	prey_o	exp1	6.0	0.181663	5
6	None	prey_o	exp1	7.0	0.348112	6
7	None	prey_o	exp1	8.0	0.937919	7
8	None	prey_o	exp1	9.0	3.113240	8
9	None	prey_o	exp1	10.0	8.863933	9
10	None	predator_o	exp1	1.0	0.847416	10
11	None	predator_o	exp1	2.0	0.211135	11
12	None	predator_o	exp1	3.0	-0.025054	12
13	None	predator_o	exp1	4.0	0.125010	13
14	None	predator_o	exp1	5.0	6.700455	14
15	None	predator_o	exp1	6.0	2.007158	15
16	None	predator_o	exp1	7.0	0.420092	16
17	None	predator_o	exp1	8.0	0.048032	17
18	None	predator_o	exp1	9.0	0.128669	18
19	None	predator_o	exp1	10.0	1.192784	19

The observables table links the model entities to the measurements:

[9]:

petab_problem.observable_df

[9]:

	observableName	observableFormula	noiseFormula	noiseDistribution
observableId
prey_o	None	prey	0.0500000000000000	normal
predator_o	None	predator	0.0500000000000000	normal

The conditions table describes the experimental conditions the measurements were collected under, as a set of targetId/targetValue changes applied to the model. In this example there is a single condition, which sets the initial values of the two model states:

[10]:

petab_problem.condition_df

[10]:

	conditionId	targetId	targetValue
0	cond1	prey	0.442493
1	cond1	predator	4.628059

The experiments table assembles conditions into experiments, by listing the simulation periods (start time and conditionId) of each experiment. Measurements refer to an experiment via experimentId:

[11]:

petab_problem.experiment_df

[11]:

	experimentId	time	conditionId
0	exp1	0.0	cond1

YAML problem file#

Finally, problem.yaml ties everything together by listing the PEtab files used to build the problem. PEtab SciML-specific information is provided under extensions.sciml:

[1]:

with open("problem.yaml") as f:
    print(f.read())

format_version: 2.0.0
model_files:
  lv:
    location: "lv.xml"
    language: "sbml"
measurement_files:
  - "measurements.tsv"
observable_files:
  - "observables.tsv"
experiment_files:
  - "experiments.tsv"
condition_files:
  - "conditions.tsv"
mapping_files:
  - "mapping.tsv"
parameter_files:
  - "parameters.tsv"
extensions:
  sciml:
    version: "0.1.0"
    required: true
    array_files:
      - "net1_ps.hdf5"
    hybridization_files:
      - "hybridization.tsv"
    neural_networks:
      net1:
        location: "net1.yaml"
        pre_initialization: false
        format: "YAML"

In more detail, extensions.sciml must list all neural networks. Each network is given a unique ID (here net1), which is then referenced throughout the PEtab SciML files (e.g. in the mapping table and array files). For each network, the YAML specifies where the network definition is stored (location), its format, and whether it is evaluated pre-initialization (pre_initialization = true) or during simulation (false, as in this UDE example). The extensions.sciml section should also list the hybridization table (if provided) and any array files.

The problem is now ready to simulate with the JAX backend. AMICI’s run_simulations evaluates the model for all experiments defined in the PEtab problem:

[12]:

import equinox as eqx

# Run simulations - results in llh - metrics in r
llh, r = run_simulations(jax_problem)

# Run simulations in gradient mode - gradients given in sllh
sllh, rgrad = eqx.filter_grad(run_simulations, has_aux=True)(jax_problem)

For more examples (including model training), see the AMICI docs. Likewise, PEtab SciML problems can also be imported and trained in Julia using PEtab.jl.

Next steps#

This tutorial showed how to set up a PEtab SciML problem where a machine-learning (ML) model enters the ODE dynamics (a UDE). PEtab SciML also supports other hybridizations, covered in the how-to guides:

Neural ODEs: Replace all ODE equations with a neural network.
ML models in observable formulas: Use an ML model in the observable formula (e.g. to correct model misspecification).
Pre-initialization ML models: Use ML models that map input data (e.g. images) to ODE parameters and/or initial conditions prior to simulation.

A collection of real-data PEtab SciML problems is also available in the PEtab SciML benchmark collection.

By design, the PEtab SciML tutorials focus on problem specification. For downstream tasks such as simulation and model training, see the Software support page for tools supporting PEtab SciML.

Lastly, this tutorial only scratches the surface of what PEtab SciML supports. For a complete description of the format, see the format specification. While a simple NN architecture is used here, the PEtab SciML NN YAML format supports a wide range of standard layers and activation functions; see the layers overview. As PEtab SciML extends PEtab, it also supports PEtab features such as partial observability, multiple simulation conditions, and events.

References#

Rackauckas, Christopher, et al. “Universal differential equations for scientific machine learning.” arXiv preprint arXiv:2001.04385 (2020).