SciML starter tutorial: Universal Differential Equations (UDEs)

Hybrid scientific machine learning (SciML) combines mechanistic (ODE) models with machine learning (ML) models [19]. PEtab.jl supports three SciML problem types, and any combination of them:

1. ML model inside the ODE dynamics. This covers both Universal Differential Equations (UDEs) and Neural ODEs.

2. ML model in the observable formula which links ODE output to measurement data.

3. Pre-simulation ML model, where the ML model is evaluated before simulation to map input data (e.g. high-dimensional images) to ODE parameters and/or initial conditions.

This tutorial introduces SciML functionality in PEtab.jl using the UDE case. Other tutorials in this section cover the remaining cases. Familiarity with setting up a PEtab problem for mechanistic models is assumed; covered in the PEtab.jl starting tutorial.

While this tutorial focuses on a simple use-case, note that SciML support is implemented on top of the functionality for mechanistic models. Thus, features for mechanistic models (e.g. simulation conditions and events) are also supported for SciML problems.

Input problem

As a running example, we use the following two-state ODE model:

$$\begin{array}{rl}\frac{\mathrm{d}X}{\mathrm{d}t}& =\frac{v{Y}^n}{Y^n+K^n}-dX\\ \frac{\mathrm{d}Y}{\mathrm{d}t}& =X-dY\end{array}$$

with initial conditions

$$X(t_0)=2.0,Y(t_0)=0.1.$$

To turn this into a UDE [19], we assume the production term for $X$ is unknown and replace it by a neural network:

$$\begin{array}{rl}\frac{\mathrm{d}X}{\mathrm{d}t}& ={\mathrm{NN}}_1(Y)-dX\\ \frac{\mathrm{d}Y}{\mathrm{d}t}& =X-dY\end{array}$$

Where ${\mathrm{NN}}_1$ is a feed-forward neural network with input $Y$.

To estimate model parameters, measurements of both $X$ and $Y$ are assumed. The goal of this tutorial is to set up a PEtab parameter estimation problem and then estimate both the mechanistic parameter d and the parameters of NN1.

Creating a PEtab SciML problem

A PEtab SciML parameter estimation problem (PEtabODEProblem) is created largely the same way as for a mechanistic model. The main difference is that ML models are provided by (1) defining one or more Lux.jl neural networks and (2) wrapping them as MLModels to specify how they interact with the ODE model.

Defining the ML model

PEtab.jl only supports Lux.jl ML models. To be compatible, the Lux model must define a set of layers with unique identifiers together with a forward pass. This is often most easily done using Lux.Chain:

using Lux
lux_model1 = Lux.Chain(
    Lux.Dense(1 => 3, Lux.softplus, use_bias = false),
    Lux.Dense(3 => 3, Lux.softplus, use_bias = false),
    Lux.Dense(3 => 1, Lux.softplus, use_bias = false),
)

Chain(
    layer_1 = Dense(1 => 3, softplus, use_bias=false),  # 3 parameters
    layer_2 = Dense(3 => 3, softplus, use_bias=false),  # 9 parameters
    layer_3 = Dense(3 => 1, softplus, use_bias=false),  # 3 parameters
)         # Total: 15 parameters,
          #        plus 0 states.

Here, softplus is used to keep the output positive. For more control over the forward pass, the model can also be defined with @compact:

lux_model2 = @compact(
    layer1 = Lux.Dense(1 => 3, Lux.softplus, use_bias = false),
    layer2 = Lux.Dense(3 => 3, Lux.softplus, use_bias = false),
    layer3 = Lux.Dense(3 => 1, Lux.softplus, use_bias = false),
) do x # forward pass
    h = layer1(x)
    h = layer2(h)
    out = layer3(h)
    @return out
end

Regardless of how the Lux model is defined, it must be wrapped as an MLModel and given a unique ID:

using PEtab
ml_model = MLModel(:net1, lux_model1, false)

MLModel net1
  mode: simulation
  parameters: 15
  hint: see model structure in `ml_model.lux_model`

The third argument specifies whether the model is evaluated pre-simulation. In this tutorial, the ML model is evaluated during simulation (it enters the ODE dynamics), so false is used. Each Lux model must be wrapped as an MLModel so PEtab.jl can keep track of how the ML model interacts with the ODE model.

Defining the dynamic (UDE) model

A UDE is created by embedding one or more MLModels into the ODE dynamics. Currently, UDE models must be provided as an ODEProblem (integration with ModelingToolkitNeuralNets is in progress). The first step is to define the ODE right-hand side. Compared to a standard in-place ODE function, an extra ml_models argument is provided. It stores the declared ML models and can be indexed by their IDs:

function ude_f!(du, u, p, t, ml_models)
    X, Y = u
    net1 = ml_models[:net1]
    nn, _ = net1.lux_model([Y], p.net1, net1.st)
    du[1] = nn[1][1] - p.d * X
    du[2] = X - p.d * Y
end

The parameter vector p is expected to be a ComponentVector, so mechanistic parameters are accessed as p.id (e.g. p.d) and ML parameters as p.ml_id (e.g. p.net1).

Given the right-hand side, the ODEProblem can be constructed using the UDEProblem helper function:

using ComponentArrays
p_mechanistic = ComponentArray(d = 1.0)
u0 = ComponentArray(X = 2.0, Y = 0.1)
ude_prob = UDEProblem(ude_f!, u0, (0.0, 10.0), p_mechanistic, ml_model)

When the dynamics are provided as an ODEProblem/UDEProblem, mechanistic parameters (p_mechanistic) and initial values (u0) must be ComponentArrays or NamedTuples so PEtab.jl can track parameter and state IDs; see Supported model systems for details.

Defining ML parameters to estimate

A PEtabMLParameter must be declared to specify whether its parameters are estimated for each MLModel. Thereby, when specifying the parameters-to-estimate vector for a SciML problem, both mechanistic parameters (PEtabParameter) and ML parameters (PEtabMLParameter) must be included in the parameter vector:

pest = [
    PEtabMLParameter(:net1),
    PEtabParameter(:d; scale = :log10)
]

2-element Vector{Any}:
 PEtabMLParameter net1: estimate
 PEtabParameter d: estimate (scale = log10, bounds = [1.0e-03, 1.0e+03])

A PEtabMLParameter can optionally be given an initial value, priors, and be fixed settings. Note parameter d is estimated on log10 scale to enforce positivity.

Measurements and observables

Measurements and observables are defined as for a mechanistic PEtabODEProblem. For the running example, we here use simulated data:

using OrdinaryDiffEqTsit5, DataFrames, Random

function f_true!(du, u, p, t)
    X, Y = u
    v, K, n, d = p
    du[1] = (v * Y^n) / (Y^n + K^n) - d * X
    du[2] = X - d * Y
end

u0_true = [2.0, 0.1]
p_true = [1.1, 2.0, 3.0, 0.5]
tend = 44.0
ode_true = ODEProblem(f_true!, u0_true, (0.0, tend), p_true)
sol = solve(
    ode_true, Tsit5(); abstol = 1e-8, reltol = 1e-8, saveat = 0:2:tend
)

data_X = sol[1, :] .+ randn(length(sol.t)) .* 0.5
data_Y = sol[2, :] .+ randn(length(sol.t)) .* 0.7
df1 = DataFrame(obs_id = "obs_X", time = sol.t, measurement = data_X)
df2 = DataFrame(obs_id = "obs_Y", time = sol.t, measurement = data_Y)
measurements = vcat(df1, df2)

Since both X and Y are assumed measured, the observables are:

observables = [
    PEtabObservable(:obs_X, :X, 0.5),
    PEtabObservable(:obs_Y, :Y, 0.7),
]

2-element Vector{PEtabObservable}:
 PEtabObservable obs_X: data ~ Normal(μ=X, σ=0.5)
 PEtabObservable obs_Y: data ~ Normal(μ=Y, σ=0.7)

Bringing it all together

Given the dynamic model (ude_prob), ML models, measurements, observables, and parameters to estimate, a PEtabModel is created largely the same way as for a mechanistic problem. The difference is that the ML models must be provided:

model_ude = PEtabModel(
    ude_prob, observables, measurements, pest; ml_models = ml_model
)

Given a PEtabModel, a PEtabODEProblem can then be constructed:

petab_prob = PEtabODEProblem(model_ude)
describe(petab_prob)

PEtabODEProblem UDEProblemModel
Problem statistics
  Parameters to estimate: 16
  ODE: 2 states, 16 parameters
  Observables: 2
  Simulation conditions: 1

ML models
  net1: (mode=simulation, parameters=15)

Configuration
  Gradient method: ForwardDiff
  Hessian method: ForwardDiff
  ODE solver (nllh): Tsit5 (abstol=1.0e-08, reltol=1.0e-08, maxiters=1e+04)
  ODE solver (grad): Tsit5 (abstol=1.0e-08, reltol=1.0e-08, maxiters=1e+04)

As seen from the problem summary, the problem includes both mechanistic and ML parameters to estimate. It is further seen the non-stiff ODE solver Tsit5() is used and gradients are computed with ForwardDiff. These options can be changed, and the same PEtabODEProblem options are supported as for mechanistic models.

Parameter estimation (model training)

SciML problems can be trained using the same approaches as mechanistic models (e.g. multi-start local optimization with a BFGS-based method via calibrate). In practice, SciML training often benefits from optimizers commonly used for ML [20], such as the Adam optimizer [21].

Here, we will set up a training loop with Adam from Optimisers.jl. Like most optimizers, Adam requires a start guess, which can be generated with:

using StableRNGs
rng = StableRNG(1) # for reproducibility
x0 = get_startguesses(rng, petab_prob, 1)

ComponentVector{Float64}(log10_d = 0.5111678532745114, net1 = (layer_1 = (weight = [-0.12361419945955276; 1.0493848323822021; -0.5047034025192261;;]), layer_2 = (weight = [0.24596905708312988 -0.13716888427734375 -0.7991292476654053; -0.9744536876678467 -0.43548011779785156 -0.7964973449707031; -0.9340267181396484 0.16604280471801758 -0.434096097946167]), layer_3 = (weight = [0.8350656032562256 -0.20964527130126953 0.1564321517944336])))

The start guess includes both mechanistic parameters (e.g. x0.d) and ML parameters (e.g. x0.net1). By default, get_startguesses initializes ML parameters with the initializers in the Lux model. While not needed here, SciML training often improve by initializing ML parameters to smaller values [22].

With a start guess available, a training loop can be written leveraging that the objective function and its gradient can be computed with petab_prob.nllh(x) and petab_prob.grad(x) respectively. For example, to train for 5000 epochs with Adam:

using Optimisers
n_epochs = 5000
x = deepcopy(x0)
learning_rate = 1e-3

state = Optimisers.setup(Adam(learning_rate), x)
for epoch in 1:n_epochs
    g = petab_prob.grad(x)
    state, x = Optimisers.update(state, x, g)

    # Stop if the objective cannot be evaluated (e.g. simulation failure)
    if !isfinite(petab_prob.nllh(x))
        break
    end
end

From plotting the fitted model, we see the model fits the data well:

using Plots
plot(x, petab_prob)

The training loop above is intentionally minimal. In practice, Optimisers.jl and related packages can be used to build training loops with learning-rate schedules, gradient clipping, early stopping, logging, etc.

Plain Adam is often inefficient for SciML problems, and it is prudent to run parameter estimation from multiple start guesses to reduce sensitivity to local minima. More efficient strategies include curriculum training, multiple shooting, or combinations thereof, all of which are supported by PEtab.jl.

Next steps

This tutorial showed how to define a PEtabODEProblem for a SciML problem. For all available options when building SciML problems (e.g. PEtabMLParameter settings), see the API documentation. For additional features, including SciML problem types beyond UDEs, see the following tutorials:

• ML models in observables: define an ML model in the observable formula of a PEtabObservable (e.g. to correct model misspecification).

• Pre-simulation ML models: define ML models that map input data (e.g. high-dimensional images) to ODE parameters or initial conditions prior to model simulation.

• Importing PEtab SciML: import problems in the PEtab-SciML standard format.

Lastly, as for mechanistic models, PEtabODEProblem has many configurable options for SciML problems. A discussion of defaults and recommendations is available in Default PEtabODEProblem options.

Copy pasteable example

using ComponentArrays, Lux, PEtab

# MLModel
lux_model1 = Lux.Chain(
    Lux.Dense(1 => 3, Lux.softplus, use_bias = false),
    Lux.Dense(3 => 3, Lux.softplus, use_bias = false),
    Lux.Dense(3 => 1, Lux.softplus, use_bias = false),
)
ml_model = MLModel(:net1, lux_model1, false)

# Create UDE-problem
function ude_f!(du, u, p, t, ml_models)
    X, Y = u
    net1 = ml_models[:net1]
    nn, _ = net1.lux_model([Y], p.net1, net1.st)
    du[1] = nn[1][1] - p.d * X
    du[2] = X - p.d * Y
end
p_mechanistic = ComponentArray(d = 1.0)
u0 = ComponentArray(X = 2.0, Y = 0.1)
ude_prob = UDEProblem(ude_f!, u0, (0.0, 10.0), p_mechanistic, ml_model)

# Parameters to estimate
pest = [
    PEtabMLParameter(:net1),
    PEtabParameter(:d; scale = :log10)
]

# Observables linking model output to data
observables = [
    PEtabObservable(:obs_X, :X, 0.5),
    PEtabObservable(:obs_Y, :Y, 0.7),
]

# Simulate data
using OrdinaryDiffEqTsit5, DataFrames
function f_true!(du, u, p, t)
    X, Y = u
    v, K, n, d = p
    du[1] = (v * Y^n) / (Y^n + K^n) - d * X
    du[2] = X - d * Y
end
u0_true = [2.0, 0.1]
p_true = [1.1, 2.0, 3.0, 0.5]
tend = 44.0
ode_true = ODEProblem(f_true!, u0_true, (0.0, tend), p_true)
sol = solve(
    ode_true, Tsit5(); abstol = 1e-8, reltol = 1e-8, saveat = 0:2:tend
)
data_X = sol[1, :] .+ randn(length(sol.t)) .* 0.5
data_Y = sol[2, :] .+ randn(length(sol.t)) .* 0.7
df1 = DataFrame(obs_id = "obs_X", time = sol.t, measurement = data_X)
df2 = DataFrame(obs_id = "obs_Y", time = sol.t, measurement = data_Y)
measurements = vcat(df1, df2)

# Create parameter estimation problem
model_ude = PEtabModel(
    ude_prob, observables, measurements, pest; ml_models = ml_model
)
petab_prob = PEtabODEProblem(model_ude)

# Get random start-guess for model training
using StableRNGs
rng = StableRNG(1) # for reproducibility
x0 = get_startguesses(rng, petab_prob, 1)

# Simple Adam training loop
using Optimisers, Plots
n_epochs = 5000
x = deepcopy(x0)
learning_rate = 1e-3
state = Optimisers.setup(Adam(learning_rate), x)
for epoch in 1:n_epochs
    g = petab_prob.grad(x)
    state, x = Optimisers.update(state, x, g)

    # Stop if the objective cannot be evaluated (e.g. simulation failure)
    if !isfinite(petab_prob.nllh(x))
        break
    end
end
plot(x, petab_prob)

References

1. C. Rackauckas, Y. Ma, J. Martensen, C. Warner, K. Zubov, R. Supekar, D. Skinner, A. Ramadhan and A. Edelman. Universal differential equations for scientific machine learning, arXiv preprint arXiv:2001.04385 (2020).

2. M. Philipps, N. Schmid and J. Hasenauer. Current state and open problems in universal differential equations for systems biology, npj Systems Biology and Applications 11, 101 (2025).

3. D. P. Kingma and J. Ba. Adam: A method for stochastic optimization, arXiv preprint arXiv:1412.6980 (2014).

4. P. Kidger. On neural differential equations, arXiv preprint arXiv:2202.02435 (2022).