SciMLSensitivity.jl

SciMLSensitivity.jl is a component package within the SciML Scientific Machine Learning ecosystem. It provides specialized utilities for sensitivity analysis.

Note for users: If your primary goal is to perform sensitivity analysis in the context of solving differential equations, you should use DifferentialEquations.jl instead, as it integrates these capabilities for that specific purpose.

SciMLSensitivity.jl provides tools for obtaining derivatives of equation solvers (like ODE solvers). These derivatives are useful for local sensitivity analysis or computing gradients for model calibration and parameter estimation. While this guide focuses on differential equations, the interface applies to all SciML ecosystem solvers, including linear/nonlinear solvers and nonlinear optimization.

import OrdinaryDiffEq as ODE
import SciMLSensitivity as SMS

# Example setup for a Lotka-Volterra system
function lotka_volterra!(du, u, p, t)
    du[1] = dx = p[1] * u[1] - p[2] * u[1] * u[2]
    du[2] = dy = -p[3] * u[2] + p[4] * u[1] * u[2]
end
p = [1.5, 1.0, 3.0, 1.0];
u0 = [1.0; 1.0];
prob = ODE.ODEProblem(lotka_volterra!, u0, (0.0, 10.0), p)
sol = ODE.solve(prob, ODE.Tsit5(), reltol = 1e-6, abstol = 1e-6)

Forward Sensitivity Analysis computes local sensitivity by solving a sensitivity ODE simultaneously with the original ODE system. The sensitivity $S_j$ for a parameter $p_j$ is governed by:

$$\frac{d}{dt}\frac{\partial u}{\partial p_{j}}=\frac{\partial f}{\partial u}\frac{\partial u}{\partial p_{j}}+\frac{\partial f}{\partial p_{j}}=J\cdot S_{j}+F_{j}$$

Where:

• $J$ is the Jacobian of the system ($\frac{\partial f}{\partial u}$).
• $F_j$ are the parameter derivatives ($\frac{\partial f}{\partial p_j}$).
• $S_j$ is the vector of sensitivities ($\frac{\partial u}{\partial p_j}$).

To maintain efficiency, SciMLSensitivity.jl avoids explicitly forming the full Jacobian $J$ whenever possible. Instead, it computes the Jacobian-vector product $Jv$ using directional derivatives or dual numbers ($d = x + v \epsilon$), where $f(d) = f(x) + Jv \epsilon$.

When working with differential equations that include events (discontinuities/callbacks), sensitivity methods fall into two categories:

Continuous Adjoint Sensitivities

Methods like BacksolveAdjoint, InterpolatingAdjoint, and QuadratureAdjoint are compatible with events for ODEs.

• BacksolveAdjoint and InterpolatingAdjoint also support events for SDEs.
• Use BacksolveAdjoint if the event terminates the time evolution and several states are saved.
• Limitation: Currently, continuous adjoint sensitivities do not support multiple events per time point.

Discrete Sensitivity Analysis

All methods based on discrete sensitivity analysis via automatic differentiation are compatible with events. This applies to ODEs, SDEs, DAEs, and DDEs. Examples include:

• ReverseDiffAdjoint (commonly used with Enzyme)
• TrackerAdjoint
• ForwardDiffSensitivity

A general rule of thumb for choosing between differentiation modes is:

• Forward-mode: Use when differentiating a system of fewer than 100 equations.
• Reverse-mode: Use when differentiating systems with more than 100 equations.

More complex scenarios may require a more nuanced choice based on the specific problem structure.

To simulate the Brusselator PDE, discretize the unit square domain into an $N \times N$ grid. Store the state variables $U$ and $V$ in a single 3D array of shape (N, N, 2) to allow for efficient processing and easy extension to more variables.

N_GRID = 16
XYD = range(0.0f0, stop = 1.0f0, length = N_GRID)
# State tensor shape: (N_GRID, N_GRID, 2)
u0 = zeros(Float32, N_GRID, N_GRID, 2)

In sensitivity analysis, a discrete cost functional relies on a finite number of time points (e.g., points returned by solve or specified via saveat).

A continuous cost functional is defined as the integral of an instantaneous cost $g$ over the entire time interval $[t_0, T]$:

$$G(u,p) = \int_{t_{0}}^{T}g(u(t,p),p)dt$$

Because continuous functionals depend on the entire trajectory, they cannot be accurately computed using only discrete estimates of the state $u$. SciMLSensitivity.jl provides direct sensitivity analysis interfaces to evaluate these functionals and their gradients, which is a capability not available through standard automatic differentiation (AD) interfaces.