Parameter Estimation for Stochastic Differential Equations and Ensembles

We can use any DEProblem, which not only includes DAEProblem and DDEProblems, but also stochastic problems. In this case, let's use the generalized maximum likelihood to fit the parameters of an SDE's ensemble evaluation.

Let's use the same Lotka-Volterra equation as before, but this time add noise:

using DifferentialEquations, DiffEqParamEstim, Plots, Optimization, ForwardDiff,
      OptimizationOptimJL
using Random, Statistics
using StochasticDiffEq: SRIW1
using SciMLLogging: None

pf_func = function (du, u, p, t)
    du[1] = p[1] * u[1] - p[2] * u[1] * u[2]
    du[2] = -3 * u[2] + u[1] * u[2]
end

u0 = [1.0; 1.0]
tspan = (0.0, 10.0)
p = [1.5, 1.0]
pg_func = function (du, u, p, t)
    du[1] = 1e-6u[1]
    du[2] = 1e-6u[2]
end
prob = SDEProblem(pf_func, pg_func, u0, tspan, p)
sol = solve(prob, SRIW1())

retcode: Success
Interpolation: 1st order linear
t: 568-element Vector{Float64}:
  0.0
  0.007499293007375976
  0.008999151608851171
  0.010686492535510766
  0.01258475107800281
  0.01472029193830636
  0.017122775406147855
  0.019825569307469537
  0.022866212446456428
  0.02628693597781668
  ⋮
  9.83071659821914
  9.851140215391485
  9.871892501523662
  9.892989212703517
  9.914446719704191
  9.93628188508557
  9.958511944177802
  9.98115430879565
 10.0
u: 568-element Vector{Vector{Float64}}:
 [1.0, 1.0]
 [1.0038129654972978, 0.9851277750640475]
 [1.0045906943595535, 0.9821833287100272]
 [1.0054716085457767, 0.9788827012009388]
 [1.0064702042798757, 0.9751844685472584]
 [1.007603293004617, 0.971042872210253]
 [1.0088902558501005, 0.9664073746408981]
 [1.0103533737503163, 0.9612224985077436]
 [1.0120188103055885, 0.9554272767861642]
 [1.0139170892225018, 0.9489552219857249]
 ⋮
 [0.9547257083612182, 1.2819367476811905]
 [0.9595110248985509, 1.229560112378899]
 [0.9654317656552531, 1.1786634503399018]
 [0.972516327365775, 1.129236508104995]
 [0.980799373519557, 1.0812687843292341]
 [0.9903220730903995, 1.0347496582943934]
 [1.0011305055713684, 0.9896683352838453]
 [1.0132775085799297, 0.9460154213759546]
 [1.024252260506165, 0.9113503993180596]

Now let's generate a compact dataset from repeated solutions of the SDE:

using RecursiveArrayTools # for VectorOfArray
t = collect(range(0, stop = 10, length = 200))
Random.seed!(1234)
function generate_data(t)
    sol = solve(prob, SRIW1())
    randomized = VectorOfArray([(sol(t[i]) + 0.01randn(2)) for i in 1:length(t)])
    data = convert(Array, randomized)
end
aggregate_data = convert(Array, VectorOfArray([generate_data(t) for i in 1:10]))
data = dropdims(mean(aggregate_data; dims = 3); dims = 3)

2×200 Matrix{Float64}:
 1.00276   1.02722   1.06243   1.10364   …  0.968177  0.994783  1.02241
 0.999879  0.909718  0.825304  0.744897     1.11028   1.00695   0.91269

Now let's estimate the parameters. Instead of using single runs from the SDE, we will use a EnsembleProblem. This means that it will solve the SDE N times to come up with an approximate probability distribution at each time point and use that in the likelihood estimate.

monte_prob = EnsembleProblem(prob)
function seeded_prob_generator(_, parameters)
    Random.seed!(2468)
    candidate_prob = remake(prob; u0 = eltype(parameters).(u0), p = parameters)
    return EnsembleProblem(candidate_prob)
end

seeded_prob_generator (generic function with 1 method)

The objective uses common random numbers so that every candidate is evaluated against the same simulated noise paths. This makes the objective deterministic enough for a gradient-based optimizer while retaining an ensemble evaluation.

We use Optim.jl for optimization below

obj = build_loss_objective(monte_prob, SRIW1(), L2Loss(t, data),
    Optimization.AutoForwardDiff(),
    maxiters = 1000, verbose = None(), trajectories = 10,
    prob_generator = seeded_prob_generator)
lower = zeros(2)
upper = fill(2.0, 2)
optprob = Optimization.OptimizationProblem(obj, [1.4, 0.95], lb = lower, ub = upper)
result = solve(optprob, Optim.BFGS())
@assert isapprox(result.u, p; atol = 0.05)

This regular L2Loss fits the mean state values. First differencing offers an alternative that also fits the increments between observations, as described in First Differencing.

result.original

 * Status: success

 * Candidate solution
    Final objective value:     4.259237e-03

 * Found with
    Algorithm:     Fminbox with BFGS

 * Convergence measures
    |x - x'|               = 7.04e-11 ≰ 0.0e+00
    |x - x'|/|x'|          = 4.70e-11 ≰ 0.0e+00
    |f(x) - f(x')|         = 3.41e-12 ≰ 0.0e+00
    |f(x) - f(x')|/|f(x')| = 8.00e-10 ≰ 0.0e+00
    |g(x)|                 = 2.41e-09 ≤ 1.0e-08

 * Work counters
    Seconds run:   14  (vs limit Inf)
    Iterations:    4
    f(x) calls:    29
    ∇f(x) calls:   29
    ∇f(x)ᵀv calls: 0

We enable first differencing by assigning weights to the state and increment residuals.

obj = build_loss_objective(monte_prob, SRIW1(),
    L2Loss(t, data, differ_weight = 1.0,
        data_weight = 0.5), Optimization.AutoForwardDiff(),
    verbose = None(), trajectories = 10, maxiters = 1000,
    prob_generator = seeded_prob_generator)
optprob = Optimization.OptimizationProblem(obj, [1.4, 0.95], lb = lower, ub = upper)
result = solve(optprob, Optim.BFGS())
@assert isapprox(result.u, p; atol = 0.05)
result.original

 * Status: success

 * Candidate solution
    Final objective value:     1.059971e-02

 * Found with
    Algorithm:     Fminbox with BFGS

 * Convergence measures
    |x - x'|               = 6.40e-11 ≰ 0.0e+00
    |x - x'|/|x'|          = 4.27e-11 ≰ 0.0e+00
    |f(x) - f(x')|         = 8.31e-11 ≰ 0.0e+00
    |f(x) - f(x')|/|f(x')| = 7.84e-09 ≰ 0.0e+00
    |g(x)|                 = 8.16e-10 ≤ 1.0e-08

 * Work counters
    Seconds run:   13  (vs limit Inf)
    Iterations:    4
    f(x) calls:    32
    ∇f(x) calls:   32
    ∇f(x)ᵀv calls: 0

Both drift parameters are recovered from the synthetic data. Larger ensembles reduce the Monte Carlo error in the objective.

println(result.u)

[1.4998893369007587, 0.9995757680575369]