Scientific Machine Learning | Victor Boussange

Developing an algebraic multigrid solver in JAX

Wed, 19 Nov 2025 00:00:00 +0000

Multigrid methods represent the state-of-the-art for solving large-scale linear systems arising from discretized partial differential equations, offering optimal computational complexity for many problem classes.

Established implementations such as pyAMG and AMG.jl provide robust solvers but lack two critical capabilities for modern scientific machine learning: GPU acceleration and automatic differentiation compatibility. These features are essential for scientific machine learning workflows where differentiable simulation components (e.g., neural networks embedded in physical models) require efficient iterative solves with gradient backpropagation for end-to-end optimization.

This project aims to develop an algebraic multigrid (AMG) solver in JAX that natively supports automatic differentiation and GPU acceleration. The work involves analyzing existing Python and Julia implementations to design an architecture compatible with JAX’s functional programming paradigm and just-in-time compilation model. A successful implementation could have substantial impact on the JAX ecosystem, from accelerating finite element packages to accelerating ecological connectivity analysis tools.

The project scope is flexible and can emphasize software engineering or algorithmic optimization depending on the student’s background and interests. Prior experience with JAX or advanced numerical linear algebra is beneficial but not required; students will gain deep expertise in iterative solvers, functional programming patterns, and best practices for scientific open-source software development.

A primer on mechanistic inference with differentiable process-based models in Julia

Thu, 02 Jan 2025 00:00:00 +0000

This tutorial covers techniques for inferring parameters of differentiable process-based models from observational data. These methods are fundamental to mechanistic inference, where we want to explain patterns in a system by understanding the processes that generate them, in contrast to purely statistical or empirical inference, which might identify patterns or correlations in data without necessarily understanding the causes. We’ll mostly focus on differential equation models. Make sure that you stick to the end, where we’ll see how we can not only infer parameter values but also the functional form of processes within the model, by parametrizing the relevant components with neural networks.

Preliminaries

Differentiable Models: Definition and Properties

One can usually write a model as a map ℳ mapping some parameters p, an initial state u₀ and a time t to a future state u_t

u_t = ℳ(u₀, t, p).

A model ℳ is differentiable if we can compute its partial derivatives with respect to parameters p or initial conditions u₀. The derivative $\frac{\partial \mathcal{M}}{\partial \theta}$ quantifies the sensitivity of model outputs to infinitesimal perturbations in parameter θ.

Recall your Calculus class!

$$\frac{df}{dx}(x) = \lim_{h \to 0} \frac{f(x + h) - f(x)}{h}$$

Let’s illustrate this concept with the logistic equation model. This model has an analytic formulation given by:

$$\mathcal{M}(u_0, p, t) = \frac{K}{1 + \big( \frac{K-u_0}{u_0} \big) e^{rt}}$$

Let’s code it

using UnPack
using Plots
using Random
using ComponentArrays
using BenchmarkTools
Random.seed!(0)

function mymodel(u0, p, t)
    T = eltype(u0)
    @unpack r, K = p

    @. K / (one(T) + (K - u0) / u0 * exp(-r * t))
end

p = ComponentArray(;r = 1., K = 1.)
u0 = 0.005

tsteps = range(0, 20, length=100)
y = mymodel(u0, p, tsteps)

plot(tsteps, y)

What is a ComponentArray?

A ComponentArray is a convenient Array type that allows to access array elements with symbols, similarly to a NamedTuple, while behaving like a standard array. For instance, you could do something like
cv = ComponentVector(;a = 1, b = 2)
cv .= [3, 4]
ComponentVector{Int64}(a = 3, b = 4)
This is useful, because you can only calculate a gradient w.r.t a Vector!

Now let’s try to calculate the gradient of this model. While you could in this case derive the gradient analytically, an analytic derivation is generally tricky with complex models. And what about models that can only be simulated numerically, with no analytic expressions? We need to find a more automatized way to calculate gradients.

How about the finite difference method?

Exercise: finite differences

Implement the function ∂mymodel_∂K(h, u0, p, t) which returns the model’s derivative with respect to K, calculated with a small h to be provided by the user.

Solution

function ∂mymodel_∂K(h, u0, p, t)
    phat = (; r = p.r, K= p.K + h)
    return (mymodel(u0, phat, t) - mymodel(u0, p, t)) / h
end
∂mymodel_∂K(1e-1, u0, p, 1.)

0.00010443404854589694

The gradient of the model is useful to understand how a parameter influences the output of the model. Let’s calculate the importance of the carrying capacity K on the model output:

dm_dp = ∂mymodel_∂K(1e-1, u0, p, tsteps)
plot(tsteps, dm_dp)

As you can observe, the carrying capacity has no effect at small t where population is small, and its influence on the dynamics grows as the population grows. We expect the reverse effect for r.

The Role of Gradients in Statistical Inference

The ability to calculate the derivative of a model is crucial when it comes to inference. Both within a full Bayesian inference context, where one wants to sample the posterior distribution of parameters θ given data u, p(θ|u), or when one wants to obtain a point estimate $\theta^\star = \text{argmax}_\theta (p(\theta | u))$ (frequentist or machine learning context), the model gradient proves very useful. In a full Bayesian inference context, they are used e.g. with Hamiltonian Markov Chains methods, such as the NUTS sampler, and in a machine learning context, they are used with gradient-based optimizer.

Gradient Descent Optimization

Gradient descent provides a fundamental algorithm for parameter estimation. The following figure illustrates the algorithm for the scalar parameter case.

Starting from an initial parameter estimate p₀, the algorithm iteratively updates parameters using the gradient $\frac{d \mathcal{M}}{dp}$ according to:

$$p_{n+1} = p_n - \eta \frac{d \mathcal{M}}{dp}(u_0, t, p) $$

where η denotes the learning rate (step size). Gradient-based optimization methods exhibit favorable scaling properties in high-dimensional parameter spaces, often achieving computational complexity advantages over derivative-free alternatives.

Automatic differentiation

Let’s go back to our method ∂mymodel_∂p. What is the optimal value of h to calculate the derivative? This is a tricky question, because a too small h can lead to round off errors (see more explanations here) while h too large also leads to a bad approximation of the asymptotic definition.

Fortunately, a bunch of techniques referred to as automatic differentiation (AD) allows to exactly differentiate any piece of numerical functions. In practice, your code must be exclusively written within an AD-backend, such as Torch, JAX or Tensorflow. Those libraries do not know how to differentiate code not written in their own language, such as normal Python code.

Fortunately, Julia is an AD-pervasive language! This means that any piece of Julia code is theoretically differentiable with AD.

using ForwardDiff

@btime ForwardDiff.gradient(p -> mymodel(u0, p, 1.), p);

  1.225 μs (12 allocations: 432 bytes)

This property makes Julia particularly well-suited for model calibration and inference: models written in native Julia are automatically compatible with AD-based inference frameworks.

For comprehensive coverage of AD in Julia, consult this tutorial and technical presentation.

Now let’s get started with inference.

Mechanistic inference

The Mechanistic Model and Synthetic Data Generation

We’ll use a simple dynamical community model, the Lotka Volterra model, to generate data. We’ll then contaminate this data with noise, and try to recover the parameters that have generated the data. The goal of the session will be to estimate those parameters from the data, using a bunch of different techniques.

So let’s first generate the data.

using OrdinaryDiffEq

# Define Lotka-Volterra model.
function lotka_volterra(du, u, p, t)
    # Model parameters.
    @unpack α, β, γ, δ = p
    # Current state.
    x, y = u

    # Evaluate differential equations.
    du[1] = (α - β * y) * x # prey
    du[2] = (δ * x - γ) * y # predator

    return nothing
end

# Define initial-value problem.
u0 = [2.0, 2.0]
p_true = (;α = 1.5, β = 1.0, γ = 3.0, δ = 1.0)
# tspan = (hudson_bay_data[1,:t], hudson_bay_data[end,:t])
tspan = (0., 5.)
tsteps = range(tspan[1], tspan[end], 51)
alg = Tsit5()

prob = ODEProblem(lotka_volterra, u0, tspan, p_true)

saveat = tsteps
sol_true = solve(prob, alg; saveat)
# Plot simulation.
plot(sol_true)

This is the true state of the system. Now let’s contaminate it with observational noise.

Exercise: Introducing observational noise

Create a data_mat array consisting of the ODE solution perturbed by lognormally-distributed multiplicative noise with standard deviation 0.3.

Note

We employ lognormal rather than Gaussian noise to ensure observations remain strictly positive, consistent with the physical constraint that population abundances cannot be negative.

Solution

data_mat = Array(sol_true) .* exp.(0.3 * randn(size(sol_true)))
# Plot simulation and noisy observations.
plot(sol_true; alpha=0.3)
scatter!(sol_true.t, data_mat'; color=[1 2], label="")

Now that we have our data, let’s do some inference!

Mechanistic Inference via Optimization

We’ll get started with a very crude approach to inference, where we’ll treat the calibration of our LV model similarly to a supervised machine learning task. To do so, we’ll write a loss function, defining a distance between our model and the data, and we’ll try to minimize this loss. The parameter minimizing this loss will be our best model parameter estimate.

function loss(p)
    predicted = solve(prob,
                        alg; 
                        p, 
                        saveat,
                        abstol=1e-6, 
                        reltol = 1e-6)

    l = 0.
    for i in 1:length(predicted)
        if all(predicted[i] .> 0)
            l += sum(abs2, log.(data_mat[:, i]) - log.(predicted[i]))
        end
    end
    return l, predicted
end

loss (generic function with 1 method)

Note

We explicitly verify that predictions remain positive, as the logarithm is undefined for non-positive values and would otherwise cause numerical errors.

Let’s define a helper function, that will plot how good does the model perform across different iterations.

losses = []
callback = function (p, l, pred; doplot=true)
    push!(losses, l)
    if length(losses)%100==1
        println("Current loss after $(length(losses)) iterations: $(losses[end])")
        if doplot
            plt = scatter(tsteps, data_mat',  color = [1 2], label=["Prey abundance data" "Predator abundance data"])
            plot!(plt, tsteps, pred', color = [1 2], label=["Inferred prey abundance" "Inferred predator abundance"])
            display(plot(plt, yaxis = :log10, title="it. : $(length(losses))"))
        end
    end
    return false
end

#13 (generic function with 1 method)

And let’s define a wrong initial guess for the parameters

pinit = ComponentArray(;α = 1., β = 1.5, γ = 1.0, δ = 0.5)

callback(pinit, loss(pinit)...; doplot = true)

Current loss after 1 iterations: 251.10349846646116

false

Our initial predictions are bad, but you’ll likely get even worse predictions in a real-case scenario!

We’ll use the library Optimization, which is a wrapper library around many optimization libraries in Julia. Optimization therefore provides us with many different types of optimizers to find parameters minimizing loss. We’ll specifically use the widely-adopted Adam optimizer, a stochastic gradient descent variant with adaptive learning rates.

using Optimization
using OptimizationOptimisers
using SciMLSensitivity

adtype = Optimization.AutoZygote()
optf = Optimization.OptimizationFunction((x, p) -> loss(x), adtype)
optprob = Optimization.OptimizationProblem(optf, pinit)

@time res_ada = Optimization.solve(optprob, Adam(0.1); callback, maxiters = 500)
res_ada.minimizer

Current loss after 101 iterations: 8.039887486778179

Current loss after 201 iterations: 7.9094080306025445

Current loss after 301 iterations: 7.806219868794404

Current loss after 401 iterations: 7.74345616951535

Current loss after 501 iterations: 7.712910946192632

 13.731183 seconds (49.62 M allocations: 3.145 GiB, 7.17% gc time, 93.45% compilation time: 8% of which was recompilation)

ComponentVector{Float64}(α = 1.5322556800023097, β = 1.0159023620691514, γ = 2.8926590524331766, δ = 0.9148575218436299)

The optimizer successfully converges to reasonable parameter estimates, demonstrating effective model calibration.

Exercise: Joint inference of initial conditions

The current implementation assumes knowledge of the true initial state u0, an unrealistic assumption in practical applications. In genuine inverse problems, initial conditions must also be inferred from data.

Modify the inference framework to simultaneously estimate both parameters and initial conditions.

Solution

function loss2(p)
    predicted = solve(prob,
                        alg; 
                        p,
                        u0 = p.u0,
                        saveat,
                        abstol=1e-6, 
                        reltol = 1e-6)
    l = 0.
    for i in 1:length(predicted)
        if all(predicted[i] .> 0)
            l += sum(abs2, log.(data_mat[:, i]) - log.(predicted[i]))
        end
    end
    return l, predicted
end
losses = []
pinit = ComponentArray(;α = 1., β = 1.5, γ = 1.0, δ = 0.5, u0 = data_mat[:,1])
adtype = Optimization.AutoZygote()
optf = Optimization.OptimizationFunction((x, p) -> loss2(x), adtype)
optprob = Optimization.OptimizationProblem(optf, pinit)
@time res_ada = Optimization.solve(optprob, Adam(0.1); callback, maxiters = 1000)
res_ada.minimizer

Current loss after 1 iterations: 416.2139476098838

Current loss after 101 iterations: 8.276915907364208

Current loss after 201 iterations: 7.932781156086005

Current loss after 301 iterations: 7.826220840461579

Current loss after 401 iterations: 7.742200328964401

Current loss after 501 iterations: 7.6847707674856744

Current loss after 601 iterations: 7.649835853033301

Current loss after 701 iterations: 7.6304539871467085

Current loss after 801 iterations: 7.620491408711084

Current loss after 901 iterations: 7.61570935872972

Current loss after 1001 iterations: 7.61357485440162

6.735983 seconds (36.27 M allocations: 2.207 GiB, 4.93% gc time, 72.28% compilation time) ComponentVector{Float64}(α = 1.4627582443041978, β = 0.9327814276650684, γ = 3.084479105946653, δ = 0.9916501731843601, u0 = [1.9639554456506427, 2.145084576010591])

Regularization Techniques

In supervised learning, it is common practice to regularize the model to prevent overfitting. Regularization can also help the model to converge. Regularization is done by adding a penalty term to the loss function:

Loss(θ) = Loss_data(θ) + λ Reg(θ)

Exercise: Implementing regularization

Incorporate a regularization term that penalizes solutions with negative initial conditions, enforcing the physical constraint of non-negative population abundances.

Multiple Shooting Methods

Multiple shooting is a numerical technique that can significantly improve optimization convergence for dynamical systems. Rather than integrating the entire trajectory from a single initial condition (single shooting), multiple shooting partitions the time domain and integrates shorter sub-intervals with independent initial conditions.

Exercise: Implementing multiple shooting

Reformulate the loss function to employ multiple shooting by dividing the observation interval into shorter segments.

Solution

function multiple_shooting_idx(N, length_interval = 10)
    K = N ÷ length_interval
    @assert N % K == 1 "`N - 1` is not a multiple of `length_interval`"
    interval_idxs = [k*length_interval+1:(k+1)*length_interval+1 for k in 0:(K-1)]
    return interval_idxs
end
function loss_multiple_shooting(p)
    interval_idxs = multiple_shooting_idx(length(tsteps))
    l = 0.
    for idx in interval_idxs
        saveat = tsteps[idx]
        # u0_i = sol_true.u[idx[1]] # here we are cheating, using true states for initial conditions!
        u0_i = data_mat[:, idx[1]] # this is not cheating, but it does not work very well
        predicted = solve(prob,
                        alg; 
                        u0 = u0_i,
                        p, 
                        saveat,
                        tspan=(saveat[1], saveat[end]),
                        abstol=1e-6, 
                        reltol = 1e-6)
        for i in 1:length(predicted)
            if all(predicted[i] .> 0)
                l += sum(abs2, log.(data_mat[:, idx[i]]) - log.(predicted[i]))
            end
        end
    end
    predicted = solve(prob,
                    alg; 
                    p,
                    saveat=tsteps,
                    abstol=1e-6, 
                    reltol = 1e-6)
    return l, predicted
end
losses = []
pinit = ComponentArray(;α = 1., β = 1.5, γ = 1.0, δ = 0.5)
adtype = Optimization.AutoZygote()
optf = Optimization.OptimizationFunction((x, p) -> loss_multiple_shooting(x), adtype)
optprob = Optimization.OptimizationProblem(optf, pinit)
@time res_ada = Optimization.solve(optprob, Adam(0.1); callback, maxiters = 500)
res_ada.minimizer

Current loss after 1 iterations: 57.64884717929634

Current loss after 101 iterations: 15.985881478253205

Current loss after 201 iterations: 15.984751300361513

Current loss after 301 iterations: 15.984751280519914

Current loss after 401 iterations: 15.98475128052433

Current loss after 501 iterations: 15.984751280410928

3.995846 seconds (16.45 M allocations: 989.683 MiB, 3.27% gc time, 69.20% compilation time) ComponentVector{Float64}(α = 2.0111356895351227, β = 1.3936359371191127, γ = 2.8416910236613444, δ = 1.031702000687222)

Sensitivity Analysis Methods

The SciMLSensitivity package and adtype = Optimization.AutoZygote() specification merit explanation, as they determine how gradients are computed for ODE-constrained optimization problems.

Automatic differentiation encompasses two primary paradigms: forward-mode and reverse-mode (adjoint) methods, with numerous algorithmic variants for each.

You can specify which ones Optimization.jl will use to differentiate loss with adtype, see available options here.

But when it comes to differentiating the solve function from OrdinaryDiffEq, you want to use AutoZygote(), because when trying to differentiate solve, a specific adjoint rule provided by the SciMLSensitivity package will be used.

What are adjoint rules?

Adjoint rules (also called custom derivatives or custom vjps) are algorithmic prescriptions that specify to an AD framework the optimal procedure for computing derivatives of specific functions. For technical details, consult the ChainRules.jl documentation.

Sensitivity algorithms are specified via the sensealg keyword argument to solve. Multiple specialized algorithms exist (reviewed in Ma et al. 2024). When sensealg is omitted, an adaptive polyalgorithm automatically selects an appropriate method based on problem characteristics.

Consult the documentation for guidance on algorithm selection.

Exercise: Benchmarking sensitivity algorithms

Compare the computational performance of ForwardDiffSensitivity() and ReverseDiffAdjoint() for the Lotka-Volterra inference problem.

Solution

using Zygote
function loss_sensealg(p, sensealg)
    predicted = solve(prob,
                        alg; 
                        sensealg,
                        p,
                        u0 = p.u0,
                        saveat,
                        abstol=1e-6, 
                        reltol = 1e-6)
    l = 0.
    for i in 1:length(predicted)
        if all(predicted[i] .> 0)
            l += sum(abs2, log.(data_mat[:, i]) - log.(predicted[i]))
        end
    end
    return l
end

loss_sensealg (generic function with 1 method)

pinit = ComponentArray(;α = 1., β = 1.5, γ = 1.0, δ = 0.5, u0 = data_mat[:,1])
@btime Zygote.gradient(p -> loss_sensealg(p, ForwardDiffSensitivity()), pinit);

  1.039 ms (14955 allocations: 896.28 KiB)

@btime Zygote.gradient(p -> loss_sensealg(p, ReverseDiffAdjoint()), pinit);

  4.904 ms (104797 allocations: 4.45 MiB)

Forward-mode methods typically exhibit superior performance for problems with few parameters, while reverse-mode (adjoint) methods scale more favorably as parameter dimensionality increases.

Well done! Now, let’s jump into the Bayesian world…

Bayesian Inference Framework

Julia has a very strong library for Bayesian inference: Turing.jl.

Let’s declare our first Turing model!

This is done with the @model macro, which allows the library to automatically construct the posterior distribution based on the definition of your model’s random variables.

Frequentist (supervised learning) vs. Bayesian approach

The main difference between a frequentist approach and a Bayesian approach is that the latter considers that parameters are random variables. Hence instead of trying to estimate a single value for the parameters, the Bayesian will try to estimate the posterior (joint) distribution of those parameters.

$$ P(\theta | \mathcal{D}) = \frac{P(\mathcal{D} | \theta) P(\theta)}{P(\mathcal{D})} $$

Random variables are defined with the ~ symbol.

Our first Turing model

using Turing
using LinearAlgebra

@model function fitlv(data, prob)
    # Prior distributions.
    σ ~ InverseGamma(3, 0.5)
    α ~ truncated(Normal(1.5, 0.5); lower=0.5, upper=2.5)
    β ~ truncated(Normal(1.2, 0.5); lower=0, upper=2)
    γ ~ truncated(Normal(3.0, 0.5); lower=1, upper=4)
    δ ~ truncated(Normal(1.0, 0.5); lower=0, upper=2)

    # Simulate Lotka-Volterra model. 
    p = (;α, β, γ, δ)
    predicted = solve(prob, alg; p, saveat)

    # Observations.
    for i in 1:length(predicted)
        if all(predicted[i] .> 0)
            data[:, i] ~ MvLogNormal(log.(predicted[i]), σ^2 * I)
        end
    end

    return nothing
end

fitlv (generic function with 2 methods)

We now instantiate the probabilistic model and perform posterior inference via Hamiltonian Monte Carlo sampling.

model = fitlv(data_mat, prob)

# Sample 3 independent chains with forward-mode automatic differentiation (the default).
chain = sample(model, NUTS(), MCMCThreads(), 1000, 3; progress=true)

Chains MCMC chain (1000×17×3 Array{Float64, 3}):

Iterations        = 501:1:1500
Number of chains  = 3
Samples per chain = 1000
Wall duration     = 26.64 seconds
Compute duration  = 25.3 seconds
parameters        = σ, α, β, γ, δ
internals         = lp, n_steps, is_accept, acceptance_rate, log_density, hamiltonian_energy, hamiltonian_energy_error, max_hamiltonian_energy_error, tree_depth, numerical_error, step_size, nom_step_size

Summary Statistics
  parameters      mean       std      mcse    ess_bulk    ess_tail      rhat   ⋯
      Symbol   Float64   Float64   Float64     Float64     Float64   Float64   ⋯

           σ    0.2796    0.0195    0.0005   1721.3574   1789.3390    1.0013   ⋯
           α    1.4928    0.1501    0.0052    841.9730    862.7118    1.0025   ⋯
           β    0.9902    0.1210    0.0040    907.2968    995.0953    1.0008   ⋯
           γ    2.9967    0.2656    0.0090    863.2374    992.8786    1.0045   ⋯
           δ    0.9592    0.1043    0.0034    939.4457   1141.4992    1.0034   ⋯
                                                                1 column omitted

Quantiles
  parameters      2.5%     25.0%     50.0%     75.0%     97.5% 
      Symbol   Float64   Float64   Float64   Float64   Float64 

           σ    0.2449    0.2656    0.2788    0.2922    0.3212
           α    1.2283    1.3875    1.4784    1.5863    1.8276
           β    0.7748    0.9077    0.9816    1.0661    1.2575
           γ    2.4786    2.8192    2.9973    3.1727    3.5358
           δ    0.7563    0.8869    0.9584    1.0274    1.1650

Threads

How many threads do you have running? Threads.nthreads() will tell you!

Let’s see if our chains have converged.

using StatsPlots
plot(chain)

Posterior Predictive Checking

Let’s now generate simulated data using samples from the posterior distribution, and compare to the original data.

function plot_predictions(chain, sol, data_mat)
    myplot = plot(; legend=false)
    posterior_samples = sample(chain[[:α, :β, :γ, :δ]], 300; replace=false)
    for parr in eachrow(Array(posterior_samples))
        p = NamedTuple([:α, :β, :γ, :δ] .=> parr)
        sol_p = solve(prob, Tsit5(); p, saveat)
        plot!(sol_p; alpha=0.1, color="#BBBBBB")
    end

    # Plot simulation and noisy observations.
    plot!(sol; color=[1 2], linewidth=1)
    scatter!(sol.t, data_mat'; color=[1 2])
    return myplot
end
plot_predictions(chain, sol_true, data_mat)

Exercise: Joint inference of initial conditions

The current implementation assumes known initial conditions u0. In realistic applications, initial states are typically unknown and must be inferred alongside parameters.

Extend the probabilistic model to include prior distributions over initial conditions.

Solution

@model function fitlv2(data, prob)
    # Prior distributions.
    σ ~ InverseGamma(2, 3)
    α ~ truncated(Normal(1.5, 0.5); lower=0.5, upper=2.5)
    β ~ truncated(Normal(1.2, 0.5); lower=0, upper=2)
    γ ~ truncated(Normal(3.0, 0.5); lower=1, upper=4)
    δ ~ truncated(Normal(1.0, 0.5); lower=0, upper=2)
    u0 ~ MvLogNormal(data[:,1], σ^2 * I)
    # Simulate Lotka-Volterra model but save only the second state of the system (predators).
    p = (;α, β, γ, δ)
    predicted = solve(prob, alg; p, u0, saveat)
    # Observations.
    for i in 2:length(predicted)
        if all(predicted[i] .> 0)
            data[:, i] ~ MvLogNormal(log.(predicted[i]), σ^2 * I)
        end
    end
    return nothing
end
model2 = fitlv2(data_mat, prob)
# Sample 3 independent chains.
chain2 = sample(model2, NUTS(), MCMCThreads(), 3000, 3; progress=true)
plot(chain2)

Here is a small utility function to visualize your results.

`plot_predictions2`

function plot_predictions2(chain, sol, data_mat)
    myplot = plot(; legend=false)
    posterior_samples = sample(chain, 300; replace=false)
    for i in 1:length(posterior_samples)
        ps = posterior_samples[i]
        p = get(ps, [:α, :β, :γ, :δ], flatten=true)
        u0 = get(ps, :u0, flatten = true)
        u0 = [u0[1][1], u0[2][1]]

        sol_p = solve(prob, Tsit5(); u0, p, saveat)
        plot!(sol_p; alpha=0.1, color="#BBBBBB")
    end

    # Plot simulation and noisy observations.
    plot!(sol; color=[1 2], linewidth=1)
    scatter!(sol.t, data_mat'; color=[1 2])
    return myplot
end

plot_predictions2(chain2, sol_true, data_mat)

Maximum A Posteriori Estimation

Turing allows you to find the maximum likelihood estimate (MLE) or maximum a posteriori estimate (MAP).

$$ \theta_{MLE} = \underset{\theta}{\text{argmax}} \ P(\mathcal{D} | \theta), \qquad \theta_{MAP} = \underset{\theta}{\text{argmax}} \ P(\theta | \mathcal{D}). $$

MAP and regularization in supervised learning

Although Bayesian inference seems very different from the supervised learning approach we developed in the first part, estimating the MAP, which can be still considered as Bayesian inference, transforms in an optimization problem that can be seen as a supervised task.

To see that, we can log-transform the posterior:

log P(θ|𝒟) = log P(𝒟|θ) + log P(θ) − log P(𝒟)

Since the evidence P(𝒟) is independent of θ, it can be ignored when maximizing with respect to θ. Therefore, the MAP estimate simplifies to:

$$ \theta_{MAP} = \underset{\theta}{\text{argmax}} \ \left[\log P(\mathcal{D} | \theta) + \log P(\theta)\right] $$

Here, log P(𝒟|θ) can be seen as our previous non-regularized loss and log P(θ) acts as a regularization term, penalizing unlikely parameter values based on our prior beliefs. Priors on parameters can be seen as regularization term.

Turing provides maximum_likelihood and maximum_a_posteriori functions for point estimation.

Random.seed!(0)
maximum_a_posteriori(model2, maxiters = 1000)

ModeResult with maximized lp of -104.88
[0.3545376205457767, 1.4695692517420373, 0.9162499950736273, 3.263944963496157, 1.0243607922108577, 2.150749205538098, 2.4795481828054595]

Since Turing uses under the hood the same Optimization.jl library, you can specify which optimizer youd’d like to use.

map_res = maximum_a_posteriori(model2, Adam(0.01), maxiters=2000)

ModeResult with maximized lp of -104.88
[0.35455374965749115, 1.4707686527453756, 0.9171941147556801, 3.2614628620071664, 1.0235193248242322, 2.1506473758409883, 2.4789084651090993]

We verify optimization convergence by examining the result object:

@show map_res.optim_result

map_res.optim_result = retcode: Default
u: [-1.036895323466616, -0.05847935462336067, -0.16599185850450063, 1.1190957778225292, 0.04704732580671333, 0.765768901858866, 0.9078183282523818]
Final objective value:     104.87762402604213

retcode: Default
u: 7-element Vector{Float64}:
 -1.036895323466616
 -0.05847935462336067
 -0.16599185850450063
  1.1190957778225292
  0.04704732580671333
  0.765768901858866
  0.9078183282523818

What’s very nice is that Turing.jl provides you with utility functions to analyse your mode estimation results.

using StatsBase
coeftable(map_res)

	Coef.	Std. Error	z	Pr(>	z	)
σ	0.354554	0.0250558	14.1506	1.85249e-45	0.305445	0.403662
α	1.47077	0.157711	9.32571	1.10241e-20	1.16166	1.77988
β	0.917194	0.125103	7.3315	2.27594e-13	0.671996	1.16239
γ	3.26146	0.335101	9.73279	2.18526e-22	2.60468	3.91825
δ	1.02352	0.124744	8.20497	2.30644e-16	0.779026	1.26801
u0[1]	2.15065	0.18462	11.6491	2.31953e-31	1.7888	2.51249
u0[2]	2.47891	0.247352	10.0218	1.22272e-23	1.99411	2.96371

Exercise: Partially observed state

Let’s assume the following situation: for some reason, you only have observation data for the predator. Could you still infer all parameters of your model, including those of the prey?

Could be! Because the signal of the variation in abundance of the predator contains information on the dynamics of the whole predator-prey system.

Do it!

You’ll need to assume so prior state for the prey. Just assume that it is the same as that of the predator.

Solution

@model function fitlv3(data::AbstractVector, prob)
    # Prior distributions.
    σ ~ InverseGamma(2, 3)
    α ~ truncated(Normal(1.5, 0.5); lower=0.5, upper=2.5)
    β ~ truncated(Normal(1.2, 0.5); lower=0, upper=2)
    γ ~ truncated(Normal(3.0, 0.5); lower=1, upper=4)
    δ ~ truncated(Normal(1.0, 0.5); lower=0, upper=2)
    u0 ~ MvLogNormal([data[1], data[1]], σ^2 * I)
    # Simulate Lotka-Volterra model but save only the second state of the system (predators).
    p = (;α, β, γ, δ)
    predicted = solve(prob, Tsit5(); p, u0, saveat, save_idxs=2)
    # Observations of the predators.
    for i in 2:length(predicted)
        if predicted[i] > 0
            data[i] ~ LogNormal(log.(predicted[i]), σ^2)
        end
    end
    return nothing
end
model3 = fitlv3(data_mat[2, :], prob)
# Sample 3 independent chains.
chain3 = sample(model3, NUTS(), MCMCThreads(), 3000, 3; progress=true)
plot(chain3)
p = plot_predictions2(chain3, sol_true, data_mat)
plot!(p, yaxis=:log10)

Now you need to realise that up to now, we had a relatively simple model. How would this model scale, should we have a much larger model? Let’s cook-up some idealised LV model. –>

Automatic Differentiation Backend Selection for MCMC

The NUTS sampler uses automatic differentiation under the hood.

By default, Turing.jl uses ForwardDiff.jl as an AD backend, meaning that the SciML sensitivity methods are not used when the solve function is called. However, you could change the AD backend to Zygote with adtype=AutoZygote().

chain2 = sample(model2, NUTS(), MCMCThreads(), adtype=AutoZygote(), 3000, 3; progress=true)

Chains MCMC chain (3000×19×3 Array{Float64, 3}):

Iterations        = 1001:1:4000
Number of chains  = 3
Samples per chain = 3000
Wall duration     = 57.41 seconds
Compute duration  = 56.94 seconds
parameters        = σ, α, β, γ, δ, u0[1], u0[2]
internals         = lp, n_steps, is_accept, acceptance_rate, log_density, hamiltonian_energy, hamiltonian_energy_error, max_hamiltonian_energy_error, tree_depth, numerical_error, step_size, nom_step_size

Summary Statistics
  parameters      mean       std      mcse    ess_bulk    ess_tail      rhat   ⋯
      Symbol   Float64   Float64   Float64     Float64     Float64   Float64   ⋯

           σ    0.3690    0.0267    0.0003   6634.3954   6080.5757    1.0000   ⋯
           α    1.5026    0.1596    0.0030   2825.0996   3566.4279    1.0004   ⋯
           β    0.9458    0.1303    0.0024   3055.7314   3652.9872    1.0015   ⋯
           γ    3.2448    0.3214    0.0060   2806.7123   2850.3476    1.0009   ⋯
           δ    1.0199    0.1212    0.0022   3167.1794   3679.8852    1.0008   ⋯
       u0[1]    2.1903    0.2017    0.0026   6066.7978   5252.2598    1.0001   ⋯
       u0[2]    2.4814    0.2547    0.0034   5638.9388   5057.7901    1.0007   ⋯
                                                                1 column omitted

Quantiles
  parameters      2.5%     25.0%     50.0%     75.0%     97.5% 
      Symbol   Float64   Float64   Float64   Float64   Float64 

           σ    0.3219    0.3507    0.3675    0.3854    0.4254
           α    1.2290    1.3890    1.4889    1.6003    1.8558
           β    0.7245    0.8536    0.9332    1.0234    1.2380
           γ    2.6101    3.0243    3.2492    3.4636    3.8759
           δ    0.7863    0.9365    1.0169    1.1007    1.2611
       u0[1]    1.8288    2.0508    2.1789    2.3149    2.6257
       u0[2]    2.0080    2.3053    2.4727    2.6454    3.0176

Doing so, you could specify within solve the adtype. It is usually a good idea to try a few different sensitivity algorithm.

See here for more information.

Exercise: Sensitivity algorithm performance comparison

Benchmark the computational efficiency of ForwardDiffSensitivity() versus ReverseDiffAdjoint() in the Bayesian inference context.

Variational Inference

Variational inference (VI) consists in approximating the true posterior distribution P(θ|𝒟) by an approximate distribution Q(θ; ϕ), where ϕ is a parameter vector defining the shape, location, and other characteristics of the approximate distribution Q, to be optimzed so that Q is as close as possible to P. This is achieved by minimizing the Kullback-Leibler (KL) divergence between the true posterior P(θ|𝒟) and the approximate distribution :

$$ \phi^* = \underset{\phi}{\text{argmin}} \ \text{KL}\left(Q(\theta; \phi) \||\ P(\theta | \mathcal{D})\right) $$

The advantage of VI over traditional MCMC sampling methods is that VI is generally faster and more scalable to large datasets, as it transforms the inference problem into an optimization problem.

Let’s do VI in Turing!

import Flux
using Turing: Variational
model = fitlv2(data_mat, prob)
q0 = Variational.meanfield(model)
advi = ADVI(10, 10_000) # first arg is the 

q = vi(model, advi, q0; optimizer=Flux.ADAM(1e-2))

function plot_predictions_vi(q, sol, data_mat)
    myplot = plot(; legend=false)
    z = rand(q, 300)
    for parr in eachcol(z)
        p = NamedTuple([:α, :β, :γ, :δ] .=> parr[2:5])
        u0 = parr[6:7]
        sol_p = solve(prob, Tsit5(); u0, p, saveat)
        plot!(sol_p; alpha=0.1, color="#BBBBBB")
    end

    # Plot simulation and noisy observations.
    plot!(sol; color=[1 2], linewidth=1)
    scatter!(sol.t, data_mat'; color=[1 2])
    return myplot
end

plot_predictions_vi(q, sol_true, data_mat)

A key advantage of VI is that the resulting approximate posterior q is an explicit distribution from which sampling is computationally trivial.

q isa MultivariateDistribution

true

rand(q)

7-element Vector{Float64}:
 0.3910702850249754
 1.8261988965103004
 1.169798842596696
 2.80613428184438
 0.8402820867003005
 2.313112765625009
 2.406422925384114

Inferring Functional Forms via Universal Differential Equations

Up to now, we have been infering the value of the model’s parameters, assuming that the structure of our model is correct. But this is very idealistic, specifically in ecology. As a general trend, we have little idea of how does e.g. the functional response of a species look like.

What if instead of inferring parameter values, we could infer functional forms, or components within our model for which we have little idea on how to express it mathematically?

In Julia, we can do that.

To illustrate this, we’ll assume that we do not know the functional response of both prey and predator, i.e. the terms β * y and δ * x. Instead, we will parametrize this component in our DE model by a neural network, which can be seen as a simple non-linear regressor dependent on some extra parameters p_nn.

We then simply have to optimize those parameters, along with the other model’s parameters!

Let’s get started. To make the neural network, we’ll use the deep learning library Lux.jl, which is similar to Flux.jl but where models are explicitly parametrized. This explicit parametrization makes it simpler to integrate with an ODE model.

To make things simpler, we will define a single layer neural network

using Lux
Random.seed!(2)
rng = Random.default_rng()
nn_init = Lux.Chain(Lux.Dense(2,2, relu))
p_nn_init, st_nn = Lux.setup(rng, nn_init)

nn = StatefulLuxLayer(nn_init, st_nn)

StatefulLuxLayer{true}(
    Dense(2 => 2, relu),                # 6 parameters
)         # Total: 6 parameters,
          #        plus 0 states.

We use a StatefulLuxLayer to not having to carry around st_nn, a struct containing states of a Lux model, which is essentially useless for a multi-layer perceptron.

st_nn

NamedTuple()

We can now evaluate our neural network model as follows:

nn(u0, p_nn_init)

2-element Vector{Float64}:
 0.0
 0.0

instead of

nn_init(u0, p_nn_init, st_nn)

([0.0, 0.0], NamedTuple())

Let’s define a new parameter vectors, which will consist of the ODE model parameters as well as the neural net parameters

pinit = ComponentArray(;σ = 0.3, α = 1., γ = 1.0, p_nn=p_nn_init)

ComponentVector{Float64}(σ = 0.3, α = 1.0, γ = 1.0, p_nn = (weight = [-1.0083649158477783 -0.7284937500953674; -1.219232201576233 0.4427390396595001], bias = [0.0; 0.0;;]))

Exercise: Implementing a universal differential equation

Formulate the UDE by replacing the interaction terms in the Lotka-Volterra system with neural network outputs.

Solution

function lotka_volterra_nn(du, u, p, t)
    # Model parameters.
    @unpack α, γ, p_nn = p
    # Current state.
    x, y = u
    û = nn(u, p_nn) # Network prediction
    # Evaluate differential equations.
    du[1] = (α - û[1]) * x # prey
    du[2] = (û[2] - γ) * y # predator
    return nothing
end

lotka_volterra_nn (generic function with 1 method)

Let’s check our initial model predictions:

prob_nn = ODEProblem(lotka_volterra_nn, u0, tspan, pinit)
init_sol = solve(prob_nn, alg; saveat)
# Plot simulation.
plot(init_sol)

Now we can define our Turing Model. We’ll need to use a utility function vector_to_parameters that reconstructs the neural network parameter type based on a sampled parameter vector (taken from this Turing tutorial). You do not need to worry about this. Note that we could have used a component vector, but for some reason this did not work at the time of the writing of this tutorial…

`vector_to_parameters`

using Functors # for the `fmap`
function vector_to_parameters(ps_new::AbstractVector, ps::NamedTuple)
    @assert length(ps_new) == Lux.parameterlength(ps)
    i = 1
    function get_ps(x)
        z = reshape(view(ps_new, i:(i + length(x) - 1)), size(x))
        i += length(x)
        return z
    end
    return fmap(get_ps, ps)
end

vector_to_parameters (generic function with 1 method)

# Create a regularization term and a Gaussian prior variance term.
sigma = 0.2

@model function fitlv_nn(data, prob)
    # Prior distributions.
    σ ~ InverseGamma(3, 0.5)
    α ~ truncated(Normal(1.5, 0.5); lower=0.5, upper=2.5)
    γ ~ truncated(Normal(3.0, 0.5); lower=1, upper=4)

    nparameters = Lux.parameterlength(nn)
    p_nn_vec ~ MvNormal(zeros(nparameters), sigma^2 * I)

    p_nn = vector_to_parameters(p_nn_vec, p_nn_init)

    # Simulate Lotka-Volterra model. 
    p = (;α, γ, p_nn)

    predicted = solve(prob, alg; p, saveat)

    # Observations.
    for i in 1:length(predicted)
        if all(predicted[i] .> 0)
            data[:, i] ~ MvLogNormal(log.(predicted[i]), σ^2 * I)
        end
    end

    return nothing
end


model = fitlv_nn(data_mat, prob_nn)

DynamicPPL.Model{typeof(fitlv_nn), (:data, :prob), (), (), Tuple{Matrix{Float64}, ODEProblem{Vector{Float64}, Tuple{Float64, Float64}, true, ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(σ = 1, α = 2, γ = 3, p_nn = ViewAxis(4:9, Axis(weight = ViewAxis(1:4, ShapedAxis((2, 2))), bias = ViewAxis(5:6, ShapedAxis((2, 1))))))}}}, ODEFunction{true, SciMLBase.AutoSpecialize, typeof(lotka_volterra_nn), UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardODEProblem}}, Tuple{}, DynamicPPL.DefaultContext}(fitlv_nn, (data = [1.8655845948955276 2.298199048573464 … 4.071164055293614 5.672667515002083; 2.651867857608795 3.2812317734519048 … 1.351784872962806 1.1243450946947573], prob = ODEProblem{Vector{Float64}, Tuple{Float64, Float64}, true, ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(σ = 1, α = 2, γ = 3, p_nn = ViewAxis(4:9, Axis(weight = ViewAxis(1:4, ShapedAxis((2, 2))), bias = ViewAxis(5:6, ShapedAxis((2, 1))))))}}}, ODEFunction{true, SciMLBase.AutoSpecialize, typeof(lotka_volterra_nn), UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardODEProblem}(ODEFunction{true, SciMLBase.AutoSpecialize, typeof(lotka_volterra_nn), UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}(lotka_volterra_nn, UniformScaling{Bool}(true), nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, SciMLBase.DEFAULT_OBSERVED, nothing, nothing, nothing, nothing), [2.0, 2.0], (0.0, 5.0), (σ = 0.3, α = 1.0, γ = 1.0, p_nn = (weight = [-1.0083649158477783 -0.7284937500953674; -1.219232201576233 0.4427390396595001], bias = [0.0; 0.0;;])), Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}(), SciMLBase.StandardODEProblem())), NamedTuple(), DynamicPPL.DefaultContext())

using Optimization, OptimizationOptimisers
@time map_res = maximum_a_posteriori(model, ADAM(0.05), maxiters=3000, initial_params=pinit)
pmap = ComponentArray(;σ=0, pinit...)
pmap .= map_res.values[:]
sol_map = solve(prob_nn, alg;p=pmap, saveat, tspan = (0, 10))
scatter(tsteps, data_mat',  color = [1 2], label=["Predator abundance data" "Prey abundance data"])
plot!(sol_map, color = [1 2], label=["Inferred predator abundance" "Inferred prey abundance"], yscale=:log10)

 12.103192 seconds (31.83 M allocations: 5.634 GiB, 3.78% gc time, 86.93% compilation time: <1% of which was recompilation)

Initial optimization struggles to converge, a common challenge in UDE inference due to the complex loss landscape introduced by neural network parameterization.

Exercise: Improving UDE optimization

What modifications might improve convergence? Consider techniques explored earlier in the tutorial.

Solution

sigma = 0.2
@model function fitlv_nn(data, prob)
    # Prior distributions.
    σ ~ InverseGamma(3, 0.5)
    α ~ truncated(Normal(1.5, 0.5); lower=0.5, upper=2.5)
    γ ~ truncated(Normal(3.0, 0.5); lower=1, upper=4)
    nparameters = Lux.parameterlength(nn)
    p_nn_vec ~ MvNormal(zeros(nparameters), sigma^2 * I)
    p_nn = vector_to_parameters(p_nn_vec, p_nn_init)
    # Simulate Lotka-Volterra model. 
    p = (;α, γ, p_nn)
    interval_idxs = multiple_shooting_idx(length(tsteps))
    for ts_idx in interval_idxs
        saveat = tsteps[ts_idx]
        u0 = sol_true.u[ts_idx[1]]
        predicted = solve(prob_nn,
                            alg; 
                            tspan = (saveat[1], saveat[end]),
                            u0,
                            p, 
                            saveat,
                            abstol=1e-6, 
                            reltol = 1e-6)
        # Observations.
        for i in 1:length(predicted)
            if all(predicted[i] .> 0)
                data[:, ts_idx[i]] ~ MvLogNormal(log.(predicted[i]), σ^2 * I)
            end
        end
    end
    return nothing
end
model = fitlv_nn(data_mat, prob_nn)
@time map_res = maximum_a_posteriori(model, Adam(0.1), maxiters=3000, initial_params=pinit)
pmap = ComponentArray(;σ=0, pinit...)
pmap .= map_res.values[:]
sol_map = solve(prob_nn, alg;p=pmap, saveat, tspan = (0, 10))
plot(sol_map, label=["Inferred predator abundance" "Inferred prey abundance"])
scatter!(sol_map.t, data_mat',  color = [1 2], label=["Predator abundance data" "Prey abundance data"], yscale=:log10)

  6.436015 seconds (53.74 M allocations: 23.768 GiB, 22.96% gc time, 14.78% compilation time: 75% of which was recompilation)

Happy with the convergence? Now let’s investigate what did the neural network learn!

`plot_func_resp`

function plot_func_resp(p, data)
    # plotting prediction of functional response
    u1 = range(minimum(data[1,:]), maximum(data[1,:]), length=100) 
    u2 = range(minimum(data[2,:]), maximum(data[2,:]), length=100) 
    u = hcat(u1,u2)

    func_resp = nn(u', p.p_nn)

    myplot1 = plot(u2,
                    - p_true.β .* u2; 
                    label="True functional form", 
                    xlabel="Predator abundance")
    plot!(myplot1,
                u2,
                - func_resp[1,:]; 
                color="#BBBBBB",
                label="Inferred functional form")

    myplot2 = plot(u1,
                    p_true.δ .* u2; 
                    legend=false, xlabel="Prey abundance")

    plot!(myplot2,
            u1,
            func_resp[2,:]; 
            color="#BBBBBB")

    myplot = plot(myplot1, myplot2)
    return myplot
end

plot_func_resp (generic function with 1 method)

plot_func_resp(pmap, data_mat)

The neural network successfully recovers the true linear functional responses, demonstrating that UDEs can discover mechanistic relationships directly from data.

Exercise: Probabilistic functional forms

Could you try to obtain a bayesian estimate of the functional forms with e.g. VI?

This concludes the tutorial. The methods presented—from gradient-based optimization through Bayesian inference to universal differential equations—provide a comprehensive framework for mechanistic inference in computational science. These techniques are readily applicable to a wide range of scientific domains where interpretability and mechanistic understanding are paramount.

Resources

On combining machine learning-based and theoretical ecosystem models

Fri, 31 Mar 2023 00:00:00 +0000

In this post, I explore the benefits and drawbacks of using empirical (ML)-based models versus mechanistic models for predicting ecosystem responses to perturbations. To evaluate these different modelling approaches, I use a mechanistic ecosystem model to generate a synthetic time series dataset.

By applying both modelling approaches to this dataset, I can evaluate their performance. While the ML-based approach yields accurate forecasts under unperturbed dynamics, it inevitably fails when it comes to predicting ecosystem response to perturbations. On the other hand, the mechanistic model, which is simplified version of the ground truth model to reflect a realistic scenario, is inaccurate and cannot forecast, but provides a more adequate approach to predict ecosystem response to unobserved scenarios.

To improve the accuracy of mechanistic models, I introduce inverse modelling, and in particular an approach that I have developed called piecewise inference. This approach allows to accurately calibrate complex mechanistic models, and doing so open doors to improve our understanding ecosystems by performing model selection.

Finally, I discuss how hybrid models, which incorporate both ML-based and mechanistic components, offer the potential to benefit from the strengths of both modelling approaches. By examining the strengths and limitations of these different modelling approaches, I hope to provide insights into how best to use them to advance our knowledge of ecological and evolutionary dynamics.

Notes

This post is under construction, and contain typos! If you find some, please contact me so that I can correct
For the sake of clarity, some pieces of code have voluntarily been hidden in external Julia files, which are loaded throughout the post. If you want to inspect them, check out those files in the corresponding GitHub repository

Generating a synthetic dataset

To generate the synthetic dataset, I consider a 3 species ecosystem model, composed of a resource, consumer and prey species. The resource growth rate depends on water availability. Here is a simplified version of the dynamics

$$\begin{aligned} \text{basal growth of } \text{🌱} &= f(\text{💧})\\ \text{per capita growth rate }\text{🌱} &= \text{basal growth} - \text{competition} - \text{grazing} - \text{death}\\ \text{per capita growth rate }\text{🦓} &= \text{grazing} - \text{predation} - \text{death}\\ \text{per capita growth rate }\text{🦁} &= \text{predation} - \text{death} \end{aligned}$$

Let’s implement that in Julia!

cd(@__DIR__)
import Pkg; Pkg.activate(".")
using PythonCall
nx = pyimport("networkx")
np = pyimport("numpy")
include("model.jl")
include("utils.jl");

To implement the model, I use the library EcoEvoModelZoo, which provides to the user ready-to-use mechanistic eco-evolutionary models. I use the model type SimpleEcosystemModel, cf. documentation of EcoEvoModelZoo. Let’s first construct the trophic network, and plot it

using EcoEvoModelZoo

N = 3 # number of species

pos = Dict(1 => [0, 0], 2 => [0.2, 1], 3 => [0, 2])
labs = Dict(1 => "Resource", 2 => "Consumer", 3 => "Prey")

foodweb = DiGraph(N)
add_edge!(foodweb, 2 => 1) # Consumer (node 2) feeds on Resource (node 1)
add_edge!(foodweb, 3 => 2) # Predator (node 3) fonds on consumer (node 2)
println(foodweb)

Graphs.SimpleGraphs.SimpleDiGraph{Int64}(2, [Int64[], [1], [2]], [[2], [3],
 Int64[]])

plot_foodweb(foodweb, pos, labs)

(<py Figure size 640x480 with 1 Axes>, <py Axes: >)

Then, I implement the processes that drive the dynamics of the ecoystem. Those include resource limitation for the resource species (e.g. limitation in nutrients), intraspecific competition for the resource species, reproduction, and feeding interactions (grazing and predation). To better understand this piece of code, you may want to refer to one of my previous blog post, “Inverse ecosystem modelling made easy with PiecewiseInference.jl”.

function carrying_capacity(p, t)
    @unpack K₁₁ = p
    K = vcat(K₁₁, ones(Float32, N - 1))
    return K
end

carrying_capacity (generic function with 1 method)

function competition(u, p, t)
    @unpack A₁₁ = p
    A = spdiagm(vcat(A₁₁, zeros(Float32, 2)))
    return A * u
end

competition (generic function with 1 method)

resource_conversion_efficiency(p, t) = ones(Float32, N)

resource_conversion_efficiency (generic function with 1 method)

Functional responses

The feeding processes implemented are based on a functional response of type II. The attack rates q define the slope of the functional response, while the handling times H define the saturation of this response.

W = adjacency_matrix(foodweb)
I, J, _ = findnz(W)

function feeding(u, p, t)
    @unpack H₂₁, H₃₂, q₂₁, q₃₂ = p

    # handling time
    H = sparse(I, J, vcat(H₂₁, H₃₂), N, N)

    # attack rates
    q = sparse(I, J, vcat(q₂₁, q₃₂), N, N)

    return q .* W ./ (one(eltype(u)) .+ q .* H .* (W * u))
end

feeding (generic function with 1 method)

Dependence of resource growth rate on water availability

We model a time-varying water availability, and a growth rate of the resource which depends on the water availability.

water_availability(t) = sin.(2 * pi / 600 * 5 * t)

growth_rate_resource(r, water) = r * exp(-0.5*(water)^2)

intinsic_growth_rate(p, t) = [growth_rate_resource(p.r[1], water_availability(t)); p.r[2:end]]

intinsic_growth_rate (generic function with 1 method)

Let’s plot what does the water availability looks like through time

ts = tspan[1]:1:tspan[end]
fig = figure()
plot(ts, water_availability.(ts))
xlabel("Time (days)")
ylabel("Water availability (normalized)")
fig.set_facecolor("None")

Python None

Now we define the numerical values for the parameter of the ecosystem model.

p_true = ComponentArray(H₂₁=Float32[1.24], # handling times
                        H₃₂=Float32[2.5],
                        q₂₁=Float32[4.98], # attack rates
                        q₃₂=Float32[0.8],
                        r=Float32[1.0, -0.4, -0.08], # growth rates
                        K₁₁=Float32[1.0], # carrying capacity for the resource
                        A₁₁=Float32[1.0]) # competition for the resource

ComponentVector{Float32}(H₂₁ = Float32[1.24], H₃₂ = Float32[2.5], q₂₁ = Flo
at32[4.98], q₃₂ = Float32[0.8], r = Float32[1.0, -0.4, -0.08], K₁₁ = Float3
2[1.0], A₁₁ = Float32[1.0])

And with that, we can plot how we implemented the dependence between the resource growth rate and the water availability. It is a gaussian dependence, where the resource growth rate is maximum at a certain value of water availability. The intuition is that too much or too little water is detrimental to the resource.

water_avail_range = reshape(sort!(water_availability.(tsteps)),length(tsteps))

fig = figure()
plot(water_avail_range, growth_rate_resource.(p_true.r[1], water_avail_range))
xlabel("Water availability (normalized)")
ylabel("Resource basal growth rate")
fig.set_facecolor("none")

Python None

Now let’s simulate the dynamics

u0_true = Float32[0.77, 0.060, 0.945]

mp = ModelParams(;p=p_true,
                tspan,
                u0=u0_true,
                solve_params...)

model = SimpleEcosystemModel(; 
                            mp, 
                            intinsic_growth_rate,
                            carrying_capacity,
                            competition,
                            resource_conversion_efficiency,
                            feeding)

`Model` SimpleEcosystemModel

fig, ax = plot_time_series(simulate(model));

And let’s generate a dataset by contaminating the model output with lognormally distributed noise.

data = simulate(model) |> Array

# contaminating raw data with noise
data = data .* exp.(1f-1*randn(Float32, size(data)))

# plotting
fig, ax = subplots(figsize=(7,4))

for i in 1:size(data,1)
    ax.scatter(tsteps, data[i, :], label=labels_sp[i], color = species_colors[i], s=10.)
end
# ax.set_yscale("log")
ax.set_ylabel("Species abundance")
ax.set_xlabel("Time (days)")
fig.set_facecolor("None")
ax.set_facecolor("None")
fig.legend()
display(fig)

Empirical modelling

Let’s build a ML-model, and train the model on the time series.

We’ll use a recurrent neural network model

More specifically, we use a Long Short Term Memory cell, connected to two dense layers with relu and a radial basis (rbf) activation functions.

using Flux # Julia deep learning library

include("rnn.jl") # Load some utility functions
args = ArgsEco() # Set up hyperparameters

rbf(x) = exp.(-(x.^2)) # custom activation function

hls = 64 # hidden layer size

# Definition of the RNN
# our model takes in the ecosystem state variables, together with current water availability
rnn_model = Flux.Chain(LSTM(N + 1, hls),
            Flux.Dense(hls, hls, relu),
                Flux.Dense(hls, N, rbf))

@time train_model!(rnn_model, data, args)   # Train and output model

75.555727 seconds (137.67 M allocations: 50.030 GiB, 6.62% gc time, 37.42%
 compilation time)

We can now simulate our trained model in an autoregressive manner, which allows us to forecast over time steps beyond the training dataset.

tsteps_forecast = tspan[end]+dt:dt:tspan[end]+100*dt

water_availability_range = water_availability.(vcat(tsteps, tsteps_forecast))
init_states = data[:,1:2]

pred_rnn = simulate(rnn_model, init_states, water_availability_range)

3×200 Matrix{Float32}:
 0.762119  0.808972   0.814921   0.789933   …  0.871319  0.816685  0.742632
 0.118674  0.0965346  0.0672258  0.0434324     0.119743  0.14609   0.185878
 0.530488  0.531594   0.518911   0.476789      0.301607  0.29676   0.31418

And let’s plot it. Plain lines are the model output for the training time span, and dashed lines are the model forecast.



fig, ax = subplots(1, figsize=(7,4))
for i in 1:N
    ax.scatter(tsteps[3:end], 
                data[i, 3:end], 
                label = labels_sp[i], 
                color = species_colors[i],
                s = 6.)
    ax.plot(tsteps[2:end], 
            pred_rnn[i,1:length(tsteps)-1], 
            c = species_colors[i])
    ax.plot(tsteps_forecast .+ dt, 
            pred_rnn[i,length(tsteps):end], 
            linestyle="--", 
            c = species_colors[i])
end
ax.set_ylabel("Species abundance")
ax.set_xlabel("Time (days)")
fig.set_facecolor("None")
ax.set_facecolor("None")
fig.legend()
display(fig)

Looks pretty good!

Now let’s see what happens if the resource species go extinct. Because this species is at the bottom of the trophic chain, we expect a collapse of the ecosystem.

water_availability_range = water_availability.(tsteps)
init_states = data[:,1:2]
init_states[1,:] .= 0.
pred_rnn = simulate(rnn_model, init_states, water_availability_range)

3×100 Matrix{Float32}:
 0.904891   0.767216   0.741785   0.774109  …  0.183886  0.0288432  0.04645
4
 0.0890969  0.0426226  0.0287545  0.019797     0.596416  0.293802   0.08636
74
 0.59085    0.542144   0.487364   0.441334     0.178284  0.282761   0.29444
7

fig, ax = subplots(1, figsize=(7,4))
for i in 1:N
    ax.plot( 
            pred_rnn[i,:], 
            c = species_colors[i])
end
ax.set_ylabel("Species abundance")
ax.set_xlabel("Time (days)")
fig.set_facecolor("None")
ax.set_facecolor("None")
display(fig)

There is clearly a problem! The RNN model outputs almost unchanged dynamics, predicting that magically, the resource would revive. This were to be expected: the ML model did not see this type of collapse dynamics, so it cannot invent it.

In summary, ML-based model

👍 Very good for interpolation
👍 Does not require any knowledge from the system, only data!
👎 Cannot extrapolate to unobserved scenarios

Mechanistic modelling

We now define a new mechanistic model, deliberatively falsifying the numerical value of the parameters compared to the baseline ecosystem model, and simplifying the resource dependence on water availbility by assuming that its growth rate is constant. The resulting model_mechanistic is as such an inaccurate representation of the baseline ecosystem model.

p_mech = ComponentArray(H₂₁=Float32[1.1],
                        H₃₂=Float32[2.3],
                        q₂₁=Float32[4.98],
                        q₃₂=Float32[0.8],
                        r = Float32[1.0, -0.4, -0.08],
                        K₁₁=Float32[0.8],
                        A₁₁=Float32[1.2])

u0_mech = Float32[0.77, 0.060, 0.945]

mp = ModelParams(; p=p_mech,
                u0=u0_mech,
                solve_params...)

growth_rate_mech(p, t) = p.r

model_mechanistic = SimpleEcosystemModel(; mp, intinsic_growth_rate = growth_rate_mech ,
                                        carrying_capacity,
                                        competition,
                                        resource_conversion_efficiency,
                                        feeding)

simul_data = simulate(model_mechanistic)

Let’s now plot its dynamics

plot_time_series(simul_data);

The dynamics looks quite different from the time series, and cannot provide any realistic forecast of the ecosystem state in the future. Yet, it does capture some of the dynamics, since it reproduces an oscillatory behavior. As such, we can assume that this ecosystem model captures some of the processes driving the baseline ecosystem dynamics.

We can now use this mechanistic model to again try to understand what happens if the resource species go extinct.

simul_data = simulate(model_mechanistic, 
                u0 = Float32[0., 0.060, 0.945], 
                saveat=0:dt:100.)

plot_time_series(simul_data);

That makes much more sense!!! The model does predict a collapse, and we can even estimate how long it will take for the system to fully collapse.

In summary, mechanistic models

👍 Very good for extrapolating to novel scenarios
👎 Hard to parametrise
👎 Inacurate

Now let’s see how we can improve this ecosystem model, by making better use of the dataset that we have at hand.

Inverse modelling

Use the observed data to infer the parameters of a model

There are two broad classes of methods to perform inverse modelling:

Bayesian inference
- provide uncertainties estimations
- does not scale well with the number of parameters to explore
Variational optimization
- Does not suffer from the curse of dimensionality
- Convergence to local minima
- Need for parameter sensitivity

To better understand the caveats of the variational optimization approach, let’s further explain it. This method consists in defining a certain loss function, which allows to measure a distance between the model output and the observations:

$$ \text{L} = \sum_i [\text{observations} - \text{simulations}]^2 $$

Variational optimization methods seek to minimize $L$ by using its gradient (sensitivity) with respect to the model parameters. This gradient indicates how to update the parameters to reduce the distance between the observations and the simulation outputs.

This can be done iteratively until finding the parameters that minimize the loss, as illustrated below.

This works very well! In theory.

In practice, the landscape is rugged, as in the picture below

By following the gradient in such a landscape, variational optimization methods tend to get stuck in wrong regions of the parameter space, and provide false estimate.

PiecewiseInference.jl

PiecewiseInference.jl is a julia package that I have authored, which implements a method to smoothen the loss landscape, and that permits to automatically obtain the model parameter sensitivity. It is detailed in the following preprint

~~Boussange, V., Vilimelis-Aceituno, P., Pellissier, L., Mini-batching ecological data to improve ecosystem models with machine learning. bioRxiv (2022), 46 pages.~~

⚠️ We will shortly rename this preprint in a revised version, as the term “mini batching” is confusing. We now prefer the term “partitioning”

The method works by training the model on small chunks of data

Let’s use it to tune the paramters of our mechanitic model. We’ll group data points in groups of 11, indicated by the argument group_size = 11

using PiecewiseInference
include("utils.jl"); # some utility functions defining `inference_problem_args` and `piecewise_MLE_args`

loss_likelihood(data, pred, rg) = sum((log.(data) .- log.(pred)).^2)

infprob = InferenceProblem(model_mechanistic, 
                            p_mech; 
                            loss_likelihood, 
                            inference_problem_args...);

@time res_mech = piecewise_MLE(infprob;
                        group_size = 11,
                        data = data,
                        tsteps = tsteps,
                        optimizers = [Adam(1e-2)],
                        epochs = [500],
                        piecewise_MLE_args...)

piecewise_MLE with 101 points and 10 groups.
Current loss after 50 iterations: 14.972400665283203
Current loss after 100 iterations: 12.230918884277344
Current loss after 150 iterations: 11.497014045715332
Current loss after 200 iterations: 11.264657020568848
Current loss after 250 iterations: 11.19422721862793
Current loss after 300 iterations: 11.16870403289795
Current loss after 350 iterations: 11.156006813049316
Current loss after 400 iterations: 11.14731216430664
Current loss after 450 iterations: 11.141851425170898
Current loss after 500 iterations: 11.138410568237305
166.285066 seconds (1.17 G allocations: 177.010 GiB, 10.44% gc time, 27.02%
 compilation time: 0% of which was recompilation)
`InferenceResult` with model SimpleEcosystemModel

Now let’s plot what does this trained model predicts

simul_mech_forecast = simulate(model_mechanistic, 
                                p = res_mech.p_trained, 
                                u0 = data[:,1],
                                saveat = vcat(tsteps, tsteps_forecast), 
                                tspan = (0, tsteps_forecast[end]));

fig, ax = subplots(1, figsize=(7,4))
for i in 1:N
    ax.scatter(tsteps, 
                data[i,:], 
                label = labels_sp[i], 
                color = species_colors[i],
                s = 6.)
    ax.plot(tsteps[1:end], 
            simul_mech_forecast[i,1:length(tsteps)], 
            c = species_colors[i])
    ax.plot(tsteps_forecast, 
            simul_mech_forecast[i,length(tsteps)+1:end], 
            linestyle="--", 
            c = species_colors[i])
end
ax.set_ylabel("Species abundance")
ax.set_xlabel("Time (days)")
fig.set_facecolor("None")
ax.set_facecolor("None")
fig.legend()
display(fig)

Looks much better! The model now captures doubling period oscillations.

In summary, mechanistic models

👍 Very good for extrapolating to novel scenarios
👎 ~~Hard to parametrise~~
👎 Inacurate

Model selection

To improve the accuracy, one can try to formulate different models corresponding to alternative hypotheses about the processes driving the ecosystem dynamics. For instance, we could compare the performance of this model to an alternative model, which would capture some sort of dependence between the water availability and the resource growth rate.

If we find that the refined model performs better than the model with constant growth rate, we have learnt that water availability is an important driver that affects the ecosystem dynamics.

Hybrid models

Assume that we have no idea on what the dependence of the resource growth rate on water availability may look like.

To proceed, we can define a very generic parametric function that can capture any sort of dependence, and then try to learn the parameters of this function from the data.

Neural networks are parametric functions that are highly suited for this sort of task. So we’ll build a hyrbid model, which contains the mechanistic components of the previous model, but where the resource growth rate is parametrized by a neural network.

🤖 + 🔬= 🤯

Below, we define the neural network, and the resource growth rate based on this neural net.

using Lux
rng = Random.default_rng()
# Multilayer FeedForward
hlsize = 5
neural_net = Lux.Chain(Lux.Dense(1,hlsize,rbf), 
                        Lux.Dense(hlsize,hlsize, rbf), 
                        Lux.Dense(hlsize,hlsize, rbf), 
                        Lux.Dense(hlsize, 1))
# Get the initial parameters and state variables of the model
p_nn, st = Lux.setup(rng, neural_net)
p_nn = p_nn |> ComponentArray

growth_rate_resource_nn(p_nn, water) = neural_net([water], p_nn, st)[1]

function hybrid_growth_rate(p, t)
    return [growth_rate_resource_nn(p.p_nn, water_availability(t)); p.r]
end

hybrid_growth_rate (generic function with 1 method)

Now we define our new hybrid model that implements this hybrid_growth_rate function

p_hybr = ComponentArray(H₂₁=p_mech.H₂₁,
                        H₃₂=p_mech.H₃₂,
                        q₂₁=p_mech.q₂₁,
                        q₃₂=p_mech.q₃₂,
                        r = p_mech.r[2:3],
                        K₁₁=p_mech.K₁₁,
                        A₁₁=p_mech.A₁₁,
                        p_nn = p_nn)

mp = ModelParams(;p=p_hybr,
                u0=u0_mech,
                solve_params...)

model_hybrid = SimpleEcosystemModel(; mp, 
                                    intinsic_growth_rate = hybrid_growth_rate,
                                    carrying_capacity,
                                    competition,
                                    resource_conversion_efficiency,
                                    feeding)

`Model` SimpleEcosystemModel

And we train it

infprob = InferenceProblem(model_hybrid, 
                            p_hybr; 
                            loss_likelihood, 
                            inference_problem_args...);

res_hybr = piecewise_MLE(infprob;
                    group_size = 11,
                    data = data,
                    tsteps = tsteps,
                    optimizers = [Adam(2e-2)],
                    epochs = [500],
                    piecewise_MLE_args...)

piecewise_MLE with 101 points and 10 groups.
Current loss after 50 iterations: 5.201906204223633
Current loss after 100 iterations: 3.6290767192840576
Current loss after 150 iterations: 3.519521713256836
Current loss after 200 iterations: 3.4955437183380127
Current loss after 250 iterations: 3.4787349700927734
Current loss after 300 iterations: 3.5343759059906006
Current loss after 350 iterations: 3.4698166847229004
Current loss after 400 iterations: 3.4563584327697754
Current loss after 450 iterations: 3.5980613231658936
Current loss after 500 iterations: 3.4506189823150635
`InferenceResult` with model SimpleEcosystemModel

The loss has signficiantly reduced. This means that this model better explains the dynamics, and as such, we have discovered that water avilability is indeed an important driver of ecosystem dynamics! Let’s plot the simulation output of this hyrbid model

simul_hybr_forecast = simulate(model_hybrid, 
                                p = res_hybr.p_trained, 
#                                 u0 = res_hybr.u0s_trained[1],
                                u0 = data[:,1],
                                saveat = vcat(tsteps, tsteps_forecast), 
                                tspan = (0, tsteps_forecast[end]));

fig, ax = subplots(1, figsize=(7,4))
for i in 1:N
    ax.scatter(tsteps, 
                data[i,:], 
                label = labels_sp[i], 
                color = species_colors[i],
                s = 6.)
    ax.plot(tsteps[1:end], 
        simul_hybr_forecast[i,1:length(tsteps)], 
            c = species_colors[i])
    ax.plot(tsteps_forecast, 
            simul_hybr_forecast[i,length(tsteps)+1:end], 
            linestyle="--", 
            c = species_colors[i])
end
ax.set_ylabel("Species abundance")
ax.set_xlabel("Time (days)")
fig.set_facecolor("None")
ax.set_facecolor("None")
fig.legend()
display(fig)

The cool thing is that although it contains a fully parametric component (the neural network), this hybrid can still extrapolate, because it is constrained by mechanistic processes.

plot_time_series(simulate(model_hybrid, 
                p = res_hybr.p_trained, 
                u0 = Float32[0., 0.060, 0.945], 
                saveat=0:dt:100.));

And what’s even cooler is that by interpreting the neural network, we can actually learn a new process: the shape of the dependence between the resource growth rate and the water availability!

water_avail = reshape(sort!(water_availability.(tsteps)),1,:)

p_nn_trained = res_hybr.p_trained.p_nn
gr = neural_net(water_avail, p_nn_trained, st)[1]


fig, ax = subplots(1)
ax.plot(water_avail[:], 
        gr[:], 
        label="Neural network",
        linestyle="--")

gr_true = model.mp.p[1] .* exp.(-0.5 .* water_avail.^2)
ax.plot(water_avail[:], gr_true[:], label="Ground truth")
ax.set_ylim(0,2)
ax.set_xlim(-2,2)
ax.legend()
xlabel("Water availability (normalized)")
ylabel("Resource basal growth rate")
display(fig)

Wrap-up

Hybrid approaches are the future! But they

Require
- domain specific knowledge
- expertise in mechanistic modelling
  - 👍 EcoEvoModelZoo.jl
- expertise in machine learning
  - 👍 Flux.jl model zoo
Requires differentiable programming and interoperability between scientific libraries
- What is best at
- But JAX is also very cool 🫠

Inverse ecosystem modeling made easy with PiecewiseInference.jl

Sun, 26 Mar 2023 00:00:00 +0000

Mechanistic ecosystem models permit to quantiatively describe how population, species or communities grow, interact and evolve. Yet calibrating them to fit real-world data is a daunting task. That’s why I’m excited to introduce PiecewiseInference.jl, a new Julia package that provides a user-friendly and efficient framework for inverse ecosystem modeling. In this blog post, I will guide you through the main features of PiecewiseInference.jl and provide a step-by-step tutorial on how to use it with a three-compartment ecosystem model. Whether you’re a quantitative ecologist or a curious data scientist, I hope this post will encourage you to join the effort and use and develop inverse ecosystem modelling methods to improve our understanding and predictions of ecosystems.

Preliminary steps

This tutorial relies on three packages that I have authored but are (yet) not registered on the official Julia registry. Those are

PiecewiseInference,
EcoEvoModelZoo: a package which provides access to a collection of ecosystem models,
ParametricModels: a wrapper package to manipulate dynamical models. Specifically, ParametricModels avoids the hassle of specifying, at each time you want to simulate an ODE model, boring details such as the algorithm to solve it, the time span, etc…

To easily install them on your machine, you’ll have to add my personal registry by doing the following:

using Pkg; Pkg.Registry.add(RegistrySpec(url = "https://github.com/vboussange/VBoussangeRegistry.git"))

Once this is done, let’s import those together with other necessary Julia packages for this tutorial.

using Graphs
using EcoEvoModelZoo
using ParametricModels
using LinearAlgebra
using UnPack
using OrdinaryDiffEq
using Statistics
using SparseArrays
using ComponentArrays
using PythonPlot

We use Graphs to create a directed graph to represent the food web to be considered The OrdinaryDiffEq package provides tools for solving ordinary differential equations, while the LinearAlgebra package is used for linear algebraic computations. The UnPack package provides a convenient way to extract fields from structures, and the ComponentArrays package is used to store and manipulate the model parameters conveniently. Finally, the PythonCall package is used to interface with Python’s Matplotlib library for visualization.

Definition of the forward model

Defining hyperparameters for the forward simulation of the model.

Next, we define the algorithm used for solving the ODE model. We also define the absolute tolerance (abstol) and relative tolerance (reltol) for the solver. tspan is a tuple representing the time range we will simulate the system for, and tsteps is a vector representing the times we want to output the simulated data.

alg = BS3()
abstol = 1e-6
reltol = 1e-6
tspan = (0.0, 600)
tsteps = range(300, tspan[end], length=100)

300.0:3.0303030303030303:600.0

Defining the foodweb structure

We’ll define a 3-compartment ecosystem as presented in McCann et al. (1994). We will use SimpleEcosystemModel from EcoEvoModeZoo.jl, which requires as input a foodweb structure. Let’s use a DiGraph to represent it.

N = 3 # number of compartment

foodweb = DiGraph(N)
add_edge!(foodweb, 2 => 1) # C to R
add_edge!(foodweb, 3 => 2) # P to C

true

The N variable specifies the number of compartments in the model. The add_edge! function is used to add edges to the graph, specifying the flow of resources between compartments.

For fun, let’s just plot the foodweb. Here we use the PythonCall and PythonPlot packages to visualize the food web as a directed graph using networkx and numpy. We create a color list for the different species, and then create a directed graph g_nx with networkx using the adjacency matrix of the food web. We also specify the position of each node in the graph, and use nx.draw to draw the graph with

using PythonCall
nx = pyimport("networkx")
np = pyimport("numpy")
species_colors = ["tab:red", "tab:green", "tab:blue"]

g_nx = nx.DiGraph(np.array(adjacency_matrix(foodweb)))
pos = Dict(0 => [0, 0], 1 => [0.2, 1], 2 => [0, 2])
labs = Dict(0 => "Resource", 1 => "Consumer", 2 => "Prey")

fig, ax = subplots(1)
nx.draw(g_nx, pos, ax=ax, node_color=species_colors, node_size=1000, labels=labs)
display(fig)

Defining the ecosystem model

Now that we have defined the foodweb structure, we can build the ecosystem model, which will be a SimpleEcosystemModel from EcoEvoModelZoo.

The next several functions are required by SimpleEcosystemModel and define the specific dynamics of the model. The intinsic_growth_rate function specifies the intrinsic growth rate of each compartment, while the carrying_capacity function specifies the carrying capacity of each compartment. The competition function specifies the competition between and within compartments, while the resource_conversion_efficiency function specifies the efficiency with which resources are converted into consumer biomass. The feeding function specifies the feeding interactions between compartments.

intinsic_growth_rate(p, t) = p.r

function carrying_capacity(p, t)
    @unpack K₁₁ = p
    K = vcat(K₁₁, ones(N - 1))
    return K
end

function competition(u, p, t)
    @unpack A₁₁ = p
    A = spdiagm(vcat(A₁₁, 0, 0))
    return A * u
end

resource_conversion_efficiency(p, t) = ones(N)

resource_conversion_efficiency (generic function with 1 method)

To define the feeding processes, we use adjacency_matrix to get the adjacency matrix of the food web. We then use findnz from SparseArrays to get the row and column indices of the non-zero entries in the adjacency matrix, which we store in I and J. Those are then used to generate sparse matrices required for defining the functional responses of each species considered. The sparse matrices’ non-zero coefficients are the model parameters to be fitted.

using SparseArrays
W = adjacency_matrix(foodweb)
I, J, _ = findnz(W)

([2, 3], [1, 2], [1, 1])

function feeding(u, p, t)
    @unpack H₂₁, H₃₂, q₂₁, q₃₂ = p

    # handling time
    H = sparse(I, J, vcat(H₂₁, H₃₂), N, N)

    # attack rates
    q = sparse(I, J, vcat(q₂₁, q₃₂), N, N)

    return q .* W ./ (one(eltype(u)) .+ q .* H .* (W * u))
end

feeding (generic function with 1 method)

We are done defining the ecological processes.

Defining the ecosystem model parameters for generating a dataset

The parameters for the ecosystem model are defined using a ComponentArray. The u0_true variable specifies the initial conditions for the simulation. The ModelParams type from the ParametricModels package is used to specify the model parameters and simulation settings. Finally, the SimpleEcosystemModel type from the EcoEvoModelZoo package is used to define the ecosystem model.

p_true = ComponentArray(H₂₁=[1.24],
                        H₃₂=[2.5],
                        q₂₁=[4.98],
                        q₃₂=[0.8],
                        r=[1.0, -0.4, -0.08],
                        K₁₁=[1.0],
                        A₁₁=[1.0])

u0_true = [0.77, 0.060, 0.945]

mp = ModelParams(; p=p_true,
    tspan,
    u0=u0_true,
    alg,
    reltol,
    abstol,
    saveat=tsteps,
    verbose=false, # suppresses warnings for maxiters
    maxiters=50_000
)
model = SimpleEcosystemModel(; mp, intinsic_growth_rate,
    carrying_capacity,
    competition,
    resource_conversion_efficiency,
    feeding)

`Model` SimpleEcosystemModel

Let’s run the model to generate a dataset! There is nothing more simple than that. Let’s also plot it, to get a sense of what it looks like.

data = simulate(model, u0=u0_true) |> Array

# plotting
using PythonPlot;
function plot_time_series(data)
    fig, ax = subplots()
    for i in 1:N
        ax.plot(data[i, :], label="Species $i", color = species_colors[i])
    end
    # ax.set_yscale("log")
    ax.set_ylabel("Species abundance")
    ax.set_xlabel("Time (days)")
    fig.set_facecolor("None")
    ax.set_facecolor("None")
    fig.legend()
    return fig
end

display(plot_time_series(data))

Let’s add a bit of noise to the data to simulate experimental errors. We proceed by adding log normally distributed noise, so that abundance are always positive (negative abundance would not make sense, but could happen when adding normally distributed noise!).

data = data .* exp.(0.1 * randn(size(data)))

display(plot_time_series(data))

Inversion with `PiecewiseInference.jl`

Now that we have set up our model and generated some data, we can proceed with the inverse modelling using PiecewiseInference.jl.

PiecewiseInference.jl allows to perform inversion based on a segmentation method that partitions the data into short time series (segments), each treated independently and matched against simulations of the model considered. The segmentation approach helps to avoid the ill-behaved loss functions that arise from the strong nonlinearities of ecosystem models, when formulation the inference problem. Note that during the inversion, not only the parameters are inferred, but also the initial conditions, which are necessary to simulate the ODE model.

Definition of the `InferenceProblem`

We first import the packages required for the inversion. PiecewiseInference is the main package used, but we also need OptimizationFlux for the Adam optimizer, and SciMLSensitivity to define the sensitivity method used to differentiate the forward model.

using PiecewiseInference
using OptimizationFlux
using SciMLSensitivity

To initialize the inversion, we set the initial values for the parameters in p_init to those of p_true but modify the H₂₁ parameter.

p_init = p_true
p_init.H₂₁ .= 2.0 #

1-element view(::Vector{Float64}, 1:1) with eltype Float64:
 2.0

Next, we define a loss function loss_likelihood that compares the observed data with the predicted data. Here, we use a simple mean-squared error loss function while log transforming the abundance, since the noise is log-normally distributed.

loss_likelihood(data, pred, rg) = sum((log.(data) .- log.(pred)) .^ 2)# loss_fn_lognormal_distrib(data, pred, noise_distrib)

loss_likelihood (generic function with 1 method)

We then define the InferenceProblem, which contains the forward model, the initial parameter values, and the loss function.

infprob = InferenceProblem(model, p_init; loss_likelihood);

It is also handy to use a callback function, that will be called after each iteration of the optimization routine, for visualizing the progress of the inference. Here, we use it to track the loss value and plot the data against the model predictions.

info_per_its = 50
include("cb.jl") # defines the `plotting_fit` function
function callback(p_trained, losses, pred, ranges)
    if length(losses) % info_per_its == 0
        plotting_fit(losses, pred, ranges, data, tsteps)
    end
end

callback (generic function with 1 method)

`piecewise_MLE` hyperparameters

To use piecewise_MLE, the main function of PiecewiseInference to estimate the parameters that fit the observed data, we need to decide on two critical hyperparameters

group_size: the number of data points that define an interval, or segment. This number is usually small, but should be decided upon the dynamics of the model: to more nonlinear is the model, the lower group_size should be. We set it here to 11
batch_size: the number of intervals, or segments, to consider on a single epoch. The higher the batch_size, the more computationally expensive a single iteration of piecewise_MLE, but the faster the convergence. Here, we set it to 5, but could increase it to 10, which is the total number of segments that we have.

Another critical parameter to be decided upon is the automatic differentiation backend used to differentiate the ODE model. Two are supported, Optimization.AutoForwardDiff() and Optimization.Autozygote(). Simply put, Optimization.AutoForwardDiff() is used for forward mode sensitivity analysis, while Optimization.Autozygote() is used for backward mode sensitivity analysis. For more information on those, please refer to the documentation of Optimization.jl.

Other parameters required by piecewise_MLE are

optimizers specifies the optimization algorithm to be used for each batch. We use the Adam optimizer, which is the go-to optimizer to train deep learning models. It has a learning rate parameter that controls the step size at each iteration. We have chosen a value of 1e-2 because it provides good convergence without causing numerical instability,
epochs specifies the number of epochs to be used for each batch. We chose a value of 500 because it is sufficient to achieve good convergence,
info_per_its specifies after how many iterations the callback function should be called
verbose_loss prints the value of the loss function during training,

@time res = piecewise_MLE(infprob;
                        adtype = Optimization.AutoZygote(),
                        group_size = 11,
                        batchsizes = [5],
                        data = data,
                        tsteps = tsteps,
                        optimizers = [Adam(1e-2)],
                        epochs = [500],
                        verbose_loss = true,
                        info_per_its = info_per_its,
                        multi_threading = false,
                        cb = callback)

piecewise_MLE with 100 points and 10 groups.
Current loss after 50 iterations: 36.745952018526495
Current loss after 100 iterations: 15.064711239454626
Current loss after 150 iterations: 9.993029255013324
Current loss after 200 iterations: 7.994491307947515
Current loss after 250 iterations: 6.500818892986831
Current loss after 300 iterations: 5.3892647156988565
Current loss after 350 iterations: 3.0351181646280514
Current loss after 400 iterations: 2.674445730720996
Current loss after 450 iterations: 3.1591980829795676
Current loss after 500 iterations: 2.4343376293865995
157.765049 seconds (1.70 G allocations: 154.827 GiB, 10.16% gc time, 31.31%
 compilation time: 1% of which was recompilation)
`InferenceResult` with model SimpleEcosystemModel

Finally, we can examine the results of the inversion. We can look at the final parameters, and the initial conditions inferred for each segement:

# Some more code
p_trained = res.p_trained
u0s_trained = res.u0s_trained

function print_param_values(p_trained, p_true)
    for k in keys(p_trained)
        println(string(k))
        println("trained value = "); display(p_trained[k])
        println("true value ="); display(p_true[k])
    end
end

print_param_values(p_trained, p_true)

H₂₁
trained value = 
1-element Vector{Float64}:
 1.4786844814887716
true value =
1-element Vector{Float64}:
 2.0
H₃₂
trained value = 
1-element Vector{Float64}:
 1.891238277791975
true value =
1-element Vector{Float64}:
 2.5
q₂₁
trained value = 
1-element Vector{Float64}:
 4.550896291686214
true value =
1-element Vector{Float64}:
 4.98
q₃₂
trained value = 
1-element Vector{Float64}:
 0.7250599871665505
true value =
1-element Vector{Float64}:
 0.8
r
trained value = 
3-element Vector{Float64}:
  0.8705446490535288
 -0.30124597815843823
 -0.08241879418666838
true value =
3-element Vector{Float64}:
  1.0
 -0.4
 -0.08
K₁₁
trained value = 
1-element Vector{Float64}:
 1.0397189294700315
true value =
1-element Vector{Float64}:
 1.0
A₁₁
trained value = 
1-element Vector{Float64}:
 0.972947353012839
true value =
1-element Vector{Float64}:
 1.0

Your turn to play!

You can try to change e.g. the batch_sizes and the group_size. How do those parameters influence the quality of the inversion?

Conclusion

PiecewiseInference.jl provides an efficient and flexible way to perform inference on complex ecological models, making use of automatic differentiation and optimizers traditionally used in Machine Learning. The segmentation method implemented in PiecewiseInference.jl regularizes the inference problem and enables inverse modelling of complex dynamical systems, for which standard methods would otherwise fail.

Furthermore, PiecewiseInference.jl together with EcoEvoModelZoo.jl offer a powerful toolkit for ecologists and evolutionary biologists to benchmark and validate models against data. The combination of theoretical modelling and data can provide new insights into complex ecological systems, helping us to better understand and predict the dynamics of biodiversity.

We invite users to explore these packages and contribute to their development, by adding new models to the EcoEvoModelZoo.jl and improve the features of PiecewiseInference.jl. With these tools, we can continue to push the boundaries of ecological modelling and make important strides towards a more sustainable future.

Appendix

You can find the corresponding tutorial as a .jmd file at https://github.com/vboussange/MyTutorials. Please contact me, if you have found a mistake, or if you have any comment or suggestion on how to improve this tutorial.

Processes analogous to ecological interactions and dispersal shape the dynamics of economic activities

Tue, 24 Jan 2023 00:00:00 +0000

Forward and inverse modelling of eco-evolutionary dynamics in ecological and economic systems

Wed, 05 Oct 2022 00:00:00 +0000

Mini-batching ecological data to improve ecosystem models with machine learning

Tue, 26 Jul 2022 00:00:00 +0000

Deep learning approximations for non-local nonlinear PDEs with Neumann boundary conditions

Thu, 05 May 2022 00:00:00 +0000

Ecology informed machine learning

Sun, 27 Jun 2021 00:00:00 +0000

As biodiversity declines and anthropogenic pressure increases, there is a crucial need for accurate ecosystem models that can capture the dynamics of real ecosystems in order to predict and mitigate collapses. Yet it remains a daunting task to obtain an agreement between current ecosystem models and real world ecosystems. In this project we aim at bridging machine learning techniques and mechanistic models in order to extrapolate beyond data. By embedding prior scientific knowledge in the structure of mechanistic models, this approach should generalise better, be more interpretable and require less data than current techniques.

Parameter Inference in dynamical systems

Sat, 09 Jan 2021 00:00:00 +0000

One of the challenges modellers face in biological sciences is to calibrate models in order to match as closely as possible observations and gain predictive power. This can be done via direct measurements through experimental design, but this process is often costly, time consuming and even sometimes not possible. Scientific machine learning addresses this problem by applying optimisation techniques originally developed within the field of machine learning to mechanistic models, allowing to infer parameters directly from observation data. In this blog post, I shall explain the basics of this approach, and how the Julia ecosystem has efficiently embedded such techniques into ready to use packages. This promises exciting perspectives for modellers in all areas of environmental sciences.

🚧 This is Work in progress 🚧

Dynamical systems are models that allow to reproduce, understand and forecast systems. They connect the time variation of the state of the system to the fundamental processes we believe driving it, that is

$$ \begin{equation*} \text{ time variation of } 🌍_t = \sum \text{processes acting on } 🌍_t \end{equation*} $$

where $🌍_t$ denotes the state of the system at time $t$. This translates mathematically into

$$ \begin{equation} \partial_t(🌍_t) = f_\theta( 🌍_t ) \end{equation}\tag{1} $$

where the function $f_\theta$ captures the ensembles of the processes considered, and depend on the parameters $\theta$.

Eq. (1) is a Differential Equation, that can be integrated with respect to time to obtain the state of the system at time $t$ given an initial state $🌍_{t_0}$.

$$ \begin{equation} 🌍_t = 🌍_{t_0} + \int_0^t f_\theta( 🌍_s ) ds \end{equation}\tag{2} $$

Dynamical systems have been used for hundreds of years and have successfully captured e.g. the motion of planets (second law of Kepler), the voltage in an electrical circuit, population dynamics (Lotka Volterra equations) and morphogenesis (Turing patterns)…

Such models can be used to forecast the state of the system in the future, or can be used in the sense of virtual laboratories. In both cases, one of the requirement is that they reproduce patterns - at least at a qualitative level. To do so, the modeler needs to find the true parameter combination $\theta$ that correspond to the system under consideration. And this is tricky! In this post we adress this challenge.

Model calibration

How to determine $\theta$ so that $\text{simulations} \approx \text{empirical data}$?

The best way to do that is to design an experiment!

When possible, measuring directly the parameters in a controlled experiment with e.g. physical devices is a great approach. This is a very powerful scientific method, used e.g. in global circulation models where scientists can measure the water viscosity, the change in water density with respect to temperature, etc… Unfortunately, such direct methods are often not possible considering other systems.

An opposite approach, known as inverse modelling, is to infer the parameters undirectly with the empirical data available.

Parameter exploration

One way to find right parameters is to perform parameter exploration, that is, slicing the parameter space and running the model for all parameter combinations chosen. Comparing the simulation results to the empirical data available, one can elect the combination with the higher explanatory power.

But as the parameter space becomes larger (higher number of parameters) this becomes tricky. Such problem is often refered to as the curse of dimensionality. Feels very much like being lost in a giant maze. We need more clever technique to get out!

A Machine Learning problem

In machine learning, people try to predict a variable $y$ from predictors $x$ by finding suitable parameters $\theta$ of a parametric function $F_\theta$ so that

$$ \begin{equation} y = F_\theta(x) \end{equation}\tag{3} $$

For example, in computer vision, this function might be designed for the specific task of labelling images, such as for instance

$F_\theta ($

$) \to \{\text{cat}, \text{dog}\}$

Usually people use neural networks so that $F_\theta \equiv NN_\theta$, as they are good approximators for high dimensional function (see the Universal approximation theorem). One should really see neural networks as functions ! For example, feed forward neural networks are mathematically described by a series of matrix multiplications and nonlinear operations, i.e. $NN_\theta (x) = \sigma_1 \circ f_1 \circ \dots \circ \sigma_n \circ f_n(x)$ where $\sigma_i$ is an activation function and $f_i$ is linear function $$ \begin{equation*} f_i (x) = A_i x + b_i . \end{equation*} $$ Notice that Eq. (2) is similar to Eq. (3)! Indeed one can think of $🌍_0$ as the analogous to $x$ - i.e. the predictor - and $🌍_t$ as the variable $y$ to predict:

$$ \begin{equation*} 🌍_t = F_\theta(🌍_{t_0}) \end{equation*} $$

where $$F_\theta (🌍_{t_0}) \equiv 🌍_{t_0} + \int_0^t f_\theta( 🌍_s ) ds .$$

With this perspective in mind, techniques developed within the field of Machine Learning - to find suitable parameters $\theta$ that best predict $y$ - become readily available to reach our specific needs: model calibration!

Parameter inference

The general strategy to find a suitable neural network that can perform the tasks required is to “train” it, that is, to find the parameters $\theta$ so that its predictions are accurate.

In order to train it, one “scores” how good a combination of parameter $\theta$ performs. A way to do so is to introduce a “Loss function”

$$ \begin{equation*} L(\theta) = (F_\theta(x) - y_{\text{empirical}})^2 \end{equation*} $$

One can then use an optimisation method to find a local minima (and in the best scenario, the global minima) for $L$.

Gradient descent

You ready?

Gradient descent and stochastic gradient descent are “iterative optimisation methods that seek to find a local minimum of a differentiable function” (Wikipedia). Such methods have become widely used with the development of artifical intelligence.

Those methods are used to compute iteratively $\theta$ using the sensitivity of the loss function to changes in $\theta$, denoted by $\partial_\theta L(\theta)$

$$ \begin{equation*} \theta^{i+1} = \theta^{(i)} - \lambda \partial_\theta L(\theta) \end{equation*} $$

where $\lambda$ is called the learning rate.

In practice

The sensitivity with respect to the parameters $\partial_\theta L(\theta)$ is in practice obtained by differentiating the code (Automatic Differentiation).

For some programming languages this can be done automatically, with low computational cost. In particular, Flux.jl allows to efficiently obtain the gradient of any function written in the wonderful language Julia.

The library DiffEqFlux.jl based on Flux.jl implements differentiation rules (custom adjoints) to obtain even more efficiently the sensitivity of a loss function that depends on the numerical solution of a differential equation. That is, DiffEqFlux.jl precisely allows to do parameter inference in dynamical systems. Go and check it out!

Machine learning to solve highly dimensional non-local nonlinear PDEs

Sat, 27 Jun 2020 00:00:00 +0000

Partial Differential Equations (PDEs) are equations that arise in a variety of models in physics, engineering, finance and biology. I develop numerical methods based on machine learning techniques to simulate a special class of PDEs in high dimension, that can capture non-locality (cf generic form of the PDEs below). Such PDEs permit to e.g. model the evolution of biological populations in a realistic manner, by describing the dynamics of several traits characterising individuals. For plants, traits correspond for instance to flower colour, specific leaf area, seed mass, etc…. These characteristics define a high dimensional space that must be considered when modelling ecological and evolutionary processes in biological populations, as they determine the overall fitness of a population in a given environment.

Trait diversity in the genus Hemerocallis. Plants can be caracterised by many different traits, all of which can be assigned numerical values: Flower colour, specific leaf area, seed mass, Plant nitrogen fixation capacity, Leaf shape, Flower sex, plant woodiness. Source: H Cui et al. 2019

High dimensionality leads to numerical difficulties in simulating the models. The methods I develop overcome this so-called “curse of dimensionality”.

The equation below defines the class of PDEs I am interested in. These PDEs are also referred in the literature as non-local reaction diffusion equations.

$$ \begin{aligned} (\tfrac{\partial}{\partial t}u)(t,x) &= \int_{D} f\big(t,x,{\bf x}, u(t,x),u(t,{\bf x}), ( \nabla_x u )(t,x ),( \nabla_x u )(t,{\bf x} ) \big) \, \nu_x(d{\bf x}) \\ & \quad + \big\langle \mu(t,x), ( \nabla_x u )( t,x ) \big\rangle + \tfrac{ 1 }{ 2 } \text{Trace}\big( \sigma(t,x) [ \sigma(t,x) ]^* ( \text{Hess}_x u)( t,x ) \big). \end{aligned} $$

I have implemented those schemes in HighDimPDE.jl, a Julia package that should allow scientists to develop models that better capture the complexity of life. These techniques extend beyond biology and are also relevant for other fields, such as finance.

Scientific Machine Learning | Victor Boussange

Developing an algebraic multigrid solver in JAX

A primer on mechanistic inference with differentiable process-based models in Julia

Preliminaries

Differentiable Models: Definition and Properties

The Role of Gradients in Statistical Inference

Gradient Descent Optimization

Automatic differentiation

Mechanistic inference

The Mechanistic Model and Synthetic Data Generation

Mechanistic Inference via Optimization

Regularization Techniques

Multiple Shooting Methods

Sensitivity Analysis Methods

Bayesian Inference Framework

Our first Turing model

Posterior Predictive Checking

Maximum A Posteriori Estimation

Automatic Differentiation Backend Selection for MCMC

Variational Inference

Inferring Functional Forms via Universal Differential Equations

Resources

On combining machine learning-based and theoretical ecosystem models

Notes

Generating a synthetic dataset

Functional responses

Dependence of resource growth rate on water availability

Empirical modelling

Mechanistic modelling

Inverse modelling

PiecewiseInference.jl

Model selection

Hybrid models

🤖 + 🔬= 🤯

Wrap-up

Inverse ecosystem modeling made easy with PiecewiseInference.jl

Preliminary steps

Definition of the forward model

Defining hyperparameters for the forward simulation of the model.

Defining the foodweb structure

Defining the ecosystem model

Defining the ecosystem model parameters for generating a dataset

Inversion with PiecewiseInference.jl

Definition of the InferenceProblem

piecewise_MLE hyperparameters

Your turn to play!

Conclusion

Appendix

Processes analogous to ecological interactions and dispersal shape the dynamics of economic activities

Forward and inverse modelling of eco-evolutionary dynamics in ecological and economic systems

Mini-batching ecological data to improve ecosystem models with machine learning

Deep learning approximations for non-local nonlinear PDEs with Neumann boundary conditions

Ecology informed machine learning

Parameter Inference in dynamical systems

Model calibration

Parameter exploration

A Machine Learning problem

Parameter inference

Gradient descent

In practice

Machine learning to solve highly dimensional non-local nonlinear PDEs

Inversion with `PiecewiseInference.jl`

Definition of the `InferenceProblem`

`piecewise_MLE` hyperparameters