Bayesian Methods

DiffEqBayes.jl

This repository is a set of extension functionality for estimating the parameters of differential equations using Bayesian methods. It allows the choice of using CmdStan.jl, Turing.jl, DynamicHMC.jl and ApproxBayes.jl to perform a Bayesian estimation of a differential equation problem specified via the DifferentialEquations.jl interface.

To begin you first need to add this repository using the following command.

Pkg.add("DiffEqBayes")
using DiffEqBayes

Tutorials and Documentation

For information on using the package, see the stable documentation. Use the in-development documentation for the version of the documentation, which contains the unreleased features.

Example

using ParameterizedFunctions, OrdinaryDiffEq, RecursiveArrayTools, Distributions
f1 = @ode_def LotkaVolterra begin
    dx = a * x - x * y
    dy = -3 * y + x * y
end a

p = [1.5]
u0 = [1.0, 1.0]
tspan = (0.0, 10.0)
prob1 = ODEProblem(f1, u0, tspan, p)

σ = 0.01                         # noise, fixed for now
t = collect(1.0:10.0)   # observation times
sol = solve(prob1, Tsit5())
priors = [Normal(1.5, 1)]
randomized = VectorOfArray([(sol(t[i]) + σ * randn(2)) for i in 1:length(t)])
data = convert(Array, randomized)

using CmdStan #required for using the Stan backend
bayesian_result_stan = stan_inference(prob1, t, data, priors)

bayesian_result_turing = turing_inference(prob1, Tsit5(), t, data, priors)

using DynamicHMC #required for DynamicHMC backend
bayesian_result_hmc = dynamichmc_inference(prob1, Tsit5(), t, data, priors)

bayesian_result_abc = abc_inference(prob1, Tsit5(), t, data, priors)

Using save_idxs to declare observables

You don't always have data for all of the variables of the model. In case of certain latent variables you can utilise the save_idxs kwarg to declare the observed variables and run the inference using any of the backends as shown below.

sol = solve(prob1, Tsit5(), save_idxs = [1])
randomized = VectorOfArray([(sol(t[i]) + σ * randn(1)) for i in 1:length(t)])
data = convert(Array, randomized)

using CmdStan #required for using the Stan backend
bayesian_result_stan = stan_inference(prob1, t, data, priors, save_idxs = [1])

bayesian_result_turing = turing_inference(prob1, Tsit5(), t, data, priors, save_idxs = [1])

using DynamicHMC #required for DynamicHMC backend
bayesian_result_hmc = dynamichmc_inference(prob1, Tsit5(), t, data, priors, save_idxs = [1])

bayesian_result_abc = abc_inference(prob1, Tsit5(), t, data, priors, save_idxs = [1])

StanODEData()

Marker used in the vars argument of stan_inference to include the simulated ODE data in the likelihood tuple passed to the Stan model. Other entries in vars are interpreted as prior distributions for likelihood parameters.

Examples

using Distributions: Normal

vars = (StanODEData(), Normal(0, 1))

StanResult(model, return_code, chains)

Result returned by stan_inference after a successful Stan run.

Fields

• model::M: the StanSample.SampleModel used for sampling.
• return_code::R: the result returned by StanSample.stan_sample.
• chains::C: the samples returned by StanSample.read_samples, using the requested output_format.

The result displays the contents of chains when printed with the text/plain MIME type.

Examples

result = StanResult(nothing, 0, (;))
result.chains

stan_inference(prob, alg, t, data, priors = nothing; kwargs...)

Run Bayesian parameter inference for a SciML problem with StanSample.jl.

When stanmodel is omitted, prob is translated to Stan's ODE interface and a Stan model is generated. The model is then sampled against data at the save times t. When stanmodel is supplied, it is sampled directly and diffeq_string can be used to provide the already-generated differential-equation function.

Arguments

• prob: an AbstractSciMLProblem containing the initial condition, parameters, and time span used by the generated Stan model.
• alg: one of :adams, :rk45, or :bdf when a model is generated. These select Stan's corresponding ODE integrators.
• t: the time points at which observations are available.
• data: an array whose columns correspond to the entries of t and whose rows correspond to the components of the problem state.
• priors: an iterable of prior distributions for the model parameters, or nothing to use Stan's normal(0, 1) prior for every parameter.

Keywords

• stanmodel: an existing StanSample.SampleModel, or nothing to generate one.
• likelihood: the likelihood distribution type or value understood by stan_string. The default is Normal.
• vars: a tuple describing the likelihood parameters. Use StanODEData() for the simulated data and distributions for additional likelihood hyperparameters.
• sample_u0: whether the selected initial-condition entries should be sampled as parameters. The default is false.
• solve_kwargs: a dictionary of Stan ODE options. Supported keys are :save_idxs, :reltol, :abstol, and :maxiter.
• diffeq_string: a pre-generated Stan differential-equation function, or nothing to generate it from prob with ModelingToolkit.
• sample_kwargs: a dictionary of Stan sampling options. Supported keys are :num_samples, :num_warmups, :num_cpp_chains, :num_chains, :num_threads, and :delta.
• output_format: the format passed to StanSample.read_samples. The default is :dataframe.
• print_summary: whether StanSample prints the sampling summary. The default is true.
• tmpdir: directory used for generated Stan files and build artifacts. The default is a new temporary directory.

Returns

Returns a StanResult containing the model, Stan return code, and sampled chains when sampling succeeds. If Stan reports a failure, the error object from the Stan return code is returned instead.

Throws

Throws an error if a generated model is requested with an unsupported alg.

Examples

using DiffEqBayes

result = stan_inference(prob, :rk45, times, observations, priors;
    vars = (StanODEData(),))
result.chains

turing_inference(prob, alg, t, data, priors; kwargs...)

Run Bayesian parameter inference for a SciML problem with Turing.jl. The problem prob is solved with alg at save times t, compared against data, and the unknown parameters are sampled from priors.

Arguments

• prob: an AbstractSciMLProblem to solve for each parameter draw.
• alg: the solver passed to solve.
• t: save times at which the solution is compared with data.
• data: observations corresponding to t.
• priors: an iterable of priors for the model parameters.

Keywords

• likelihood_dist_priors: priors for the likelihood scale parameters. The default is [InverseGamma(2, 3)].
• likelihood: callable returning the observation distribution from (u, p, t, σ).
• syms: Turing variable names for the sampled parameters. The default creates one name per entry in priors.
• sample_u0: whether the initial condition is included in the sampled parameters.
• progress: whether Turing displays sampling progress.
• solve_kwargs: dictionary of keyword arguments forwarded to solve; :save_idxs defaults to nothing.
• sample_args: named arguments controlling Turing's sampler, with defaults of NUTS(0.65), MCMCSerial(), 1000 samples, and one chain.
• sample_kwargs: additional keyword arguments forwarded to Turing.sample.

Returns

Returns the chain produced by Turing.sample. With current Turing releases this is a FlexiChains.VNChain of size (num_samples, n_chains) rather than the older MCMCChains.Chains. Draws are indexed using the VarName values in syms.

Examples

chain = turing_inference(prob, Tsit5(), t, data, priors; syms = [:a])
mean(chain[@varname(a)])

dynamichmc_inference(
    problem,
    algorithm,
    t,
    data,
    parameter_priors;
    ...
)
dynamichmc_inference(
    problem,
    algorithm,
    t,
    data,
    parameter_priors,
    parameter_transformations;
    σ_priors,
    sample_u0,
    rng,
    num_samples,
    AD_gradient_kind,
    save_idxs,
    solve_kwargs,
    mcmc_kwargs
)

Run MCMC for an ODE problem with DynamicHMC.jl.

The ODE is solved at the times in t for each sampled parameter vector. The resulting likelihood is combined with the parameter and noise-scale priors, transformed to an unconstrained space, and passed to DynamicHMC.mcmc_with_warmup.

Arguments

• problem: an AbstractSciMLProblem to solve.
• algorithm: the ODE algorithm passed to solve.
• t: the time values at which the solution is compared with data.
• data: a matrix with one column for each value in t.
• parameter_priors: an iterable of parameter priors with one entry for each sampled parameter.
• parameter_transformations: a TransformVariables transformation mapping an unconstrained real vector to the valid parameter space.

Keywords

• σ_priors: priors for the noise scale of each observed component. The default is Normal(0, 5) for every component.
• sample_u0: whether the selected initial-condition entries are sampled. The default is false.
• rng: random number generator used for MCMC. The default is Random.default_rng().
• num_samples: number of MCMC draws. The default is 1000.
• AD_gradient_kind: gradient implementation passed to LogDensityProblemsAD.ADgradient. The default is Val(:ForwardDiff); load the corresponding AD package when changing it.
• save_idxs: state indices used for observations and initial-condition sampling, or nothing to use every state component.
• solve_kwargs: keyword arguments forwarded to solve.
• mcmc_kwargs: keyword arguments forwarded to DynamicHMC.mcmc_with_warmup. Its initialization.q vector must have one entry for every sampled parameter and noise scale.

Returns

Returns a NamedTuple containing the fields returned by DynamicHMC.mcmc_with_warmup and an additional posterior field. posterior contains the sampled parameter vectors transformed back to the constrained parameter space.

Examples

using DiffEqBayes
using TransformVariables: as, asℝ₊

posterior = dynamichmc_inference(
    prob, Tsit5(), times, observations, priors,
    as(Vector, asℝ₊, length(priors)); num_samples = 100
)