dynamicHMC with diffeq MWE

code.jl

using OrdinaryDiffEq, DynamicHMC, TransformVariables, LogDensityProblems, MCMCDiagnostics,
    Parameters, Distributions, LinearAlgebra

struct BayesianODEProblem{F, U, S, T, P, N, D}
    f::F                        # ODE
    u::U                        # initial value
    timespan::S                 # timespan as a tuple
    timepoints::T               # evaluated at these points
    logprior::P                 # callable, return the log prior of parameters
    noise::N                    # a noise distribution for observations
    data::D                     # observed data
end

function simulate_problem(f, u, timepoints, noise, parameters, logprior)
    timespan = extrema(timepoints) .+ (0, √eps()) # NOTE: extend timespan for robustness
    odeproblem = ODEProblem(f, u, timespan, parameters)
    solution = solve(odeproblem, Tsit5()) # NOTE: method hardcoded
    data = [solution(timepoint) .+ rand(noise) for timepoint in timepoints]
    BayesianODEProblem(f, u, timespan, timepoints, logprior, noise, data)
end

function (problem::BayesianODEProblem)(parameters)
    @unpack f, u, timespan, timepoints, logprior, noise, data = problem
    u_widened = typeof(parameters.a).(u) # NOTE: UGLY HACK
    odeproblem = ODEProblem(f, u_widened, timespan, parameters)
    solution = solve(odeproblem, Tsit5())
    loglikelihood = sum(logpdf(noise, d .- solution(timepoint))
                        for (d, timepoint) in zip(data, timepoints))
    loglikelihood + logprior(parameters)
end

function parameterized_lotkavolterra(du, u, p, t)
    @unpack a = p
    x, y = u
    du[1] = a*x - x*y
    du[2] = -3*y + x*y
end

function prior_lotkavolterra(parameters)
    @unpack a = parameters
    logpdf(Normal(0, 10), a)    # very vague but proper prior
end

a₀ = 1.5
θ₀ = (a = a₀, )
p = simulate_problem(parameterized_lotkavolterra,                 # problem
                     [1.0, 1.0],                                  # initial position
                     range(0, stop = 10, length = 10),            # timepoints
                     MvNormal(zeros(2), Diagonal(ones(2) * 0.1)), # noise
                     θ₀,                                          # initial parameters
                     prior_lotkavolterra)

p((a = a₀, ))                         # test
t = as((a = asℝ₊, ))                  # transformation: to positive
P = TransformedLogDensity(t, p)       # transformed to work on ℝⁿ
∇P = ADgradient(Val(:ForwardDiff), P) # gradient via AD

logdensity(LogDensityProblems.ValueGradient, ∇P, [1.0]) # test gradient

chain, nuts = NUTS_init_tune_mcmc(∇P, 1000); q = inverse(t, (a = a₀,))

NUTS_statistics(chain)          # very nice

mapslices(effective_sample_size, get_position_matrix(chain); dims = 1) # OK

posterior = t.(get_position.(chain))

using PGFPlotsX
agrid = range(1.45, stop = 1.5, length = 1000)
ℓgrid = [p((a = a, )) for a in agrid]
@pgf Axis({ xlabel = "a", ylabel = "log posterior"},  Plot({ no_marks}, Table(agrid, ℓgrid)))