Multi-Backend ODE Solving¶

Learning Objectives:

• Use VectorBuilder for custom ODE solvers
• Compare Diffsol, Diffrax (JAX), and DifferentialEquations.jl
• Understand performance trade-offs
• Integrate external solvers with Diffid

Prerequisites: Optimisation Basics, ODE Fitting, basic JAX or Julia knowledge (optional)

Runtime: ~20 minutes

Introduction¶

Diffid's DiffsolBuilder provides high-performance ODE solving for most cases. But sometimes you need:

• JAX/Diffrax: Automatic differentiation, GPU acceleration
• Julia/DifferentialEquations.jl: Specialized solvers, stiff equations
• Custom simulators: Agent-based models, PDEs, hybrid systems

VectorBuilder lets you integrate any Python-callable forward model with Diffid's optimisers.

The Lotka-Volterra Model¶

The predator-prey equations:

$$\begin{aligned} \frac{dx}{dt} &= \alpha x - \beta x y \\ \frac{dy}{dt} &= \delta x y - \gamma y \end{aligned}$$

where:

• $x$ is prey population
• $y$ is predator population
• $\alpha, \beta, \gamma, \delta$ are interaction rates

This tutorial demonstrates parameter fitting with three different solver backends.

Coming Soon¶

This advanced tutorial is under development. It will cover:

Backend 1: Diffsol (Built-in)¶

builder = (
    diffid.DiffsolBuilder()
    .with_diffsl(lotka_volterra_dsl)
    .with_data(data)
    .with_parameter("alpha", 1.0)
    # ... more parameters
)

Backend 2: JAX/Diffrax¶

import jax
from diffrax import diffeqsolve, ODETerm, Tsit5

def diffrax_solver(params):
    # Your Diffrax integration
    return predictions

builder = (
    diffid.VectorBuilder()
    .with_objective(diffrax_solver)
    .with_data(data)
    .with_parameter("alpha", 1.0)
)

Backend 3: Julia/DifferentialEquations.jl¶

from diffeqpy import de

def diffeqpy_solver(params):
    # Your Julia integration
    return predictions

builder = (
    diffid.VectorBuilder()
    .with_objective(diffeqpy_solver)
    .with_data(data)
    .with_parameter("alpha", 1.0)
)

Performance Comparison¶

The tutorial will benchmark all three backends:

• Accuracy: Parameter recovery quality
• Speed: Time per function evaluation
• Ease of use: Setup complexity
• Special features: Gradients, GPU, stiff solvers

When to Use Each Backend¶

Example Data¶

The predator-prey examples directory contains:

• generate_data_diffrax.py: Creates synthetic data
• predator_prey_diffsol.py: Diffsol backend
• predator_prey_diffrax.py: JAX/Diffrax backend
• predator_prey_diffeqpy.py: Julia backend

Run these scripts directly to see the backends in action!

Installation¶

For JAX/Diffrax:

pip install jax diffrax

For Julia/DifferentialEquations.jl:

pip install diffeqpy
python -c "from diffeqpy import de; de.install()"

Key Takeaways¶

1. VectorBuilder integrates any Python callable
2. Diffsol is the default - fast and easy
3. JAX/Diffrax for gradients and GPU
4. Julia/DiffEq for specialized solvers
5. All backends work with Diffid's optimisers

Next Steps¶

• Custom Solvers Guide - Detailed integration guide
• VectorBuilder API
• Examples Gallery - More backend examples