DifferentialEquation example?

[ChrisRackauckas]

I was curious if there was a differentialequations example in here and if not, what's the technical difficulty involved?

[ismailsenoz]

We do not have an example that utilizes DifferentialEquations if that is what you were asking. We have a state prediction and estimation demo for a Lorentz attractor using LSTM. We also have a demo for solving GP regression by discretizing the underlying SDE.

Regarding technical difficulties. It is hard to predict, tbh. It depends on what you want to do with RxInfer. If you're going to use RxInfer to learn the parameters of a differential equation, that should not be challenging. However, solving the underlying differential equation using message-passing with RxInfer might be difficult because RxInfer is not aimed at solving differential equations. It would be pretty interesting to use RxInfer for solving differential equations.

If you have something more concrete in your mind, it would be nice to discuss it.

[ChrisRackauckas]

What's required for it to hook together? Does the message passing require a quasi-static compute graph? I don't see Cassette in here so I'm curious how it's implemented.

[Nimrais]

I just wanted to see if there some preferred techniques you were using; given that simulation and inference are fast, I wondered whether some sort of conformal prediction could be used to recalibrate the posteriors. I'll think about it some model refinement.

There is no generic answer to your question unfortunately.

There are two possible ways to address the issue in my head.

1. Invest more computations, be aware that I am using main branch not a release of https://github.com/ReactiveBayes/ExponentialFamilyProjection.jl, this version will be released soon

using RxInfer
using OrdinaryDiffEq
using ExponentialFamilyManifolds
using ExponentialFamilyProjection
using Optimisers
using StableRNGs
using StatsFuns
using StaticArrays
using Plots
using BenchmarkTools
import BayesBase;
using ProgressBars
using Base.Threads

function sir_ode(u, p, t)
    (S,I,R,C) = u
    (β,c,γ) = p
    N = S+I+R
    infection = β*c*I/N*S
    recovery = γ*I
    return @SArray([
         -infection,
        infection - recovery,
        recovery,
        infection
    ]) # S, I, R, C
end;


function sir_ode_solve(problem, l, i₀, β)
    I = i₀*1000.0
    S = 1000.0 - I
    u0 = @SArray([S, I, 0.0, 0.0])
    p = @SArray([β, 10.0, 0.25])
    prob = remake(problem; u0=u0, p=p)
    sol = solve(prob, Tsit5(), saveat = 1.0)
    sol_C = view(sol, 4, :) # Cumulative cases
    sol_X = Array{eltype(sol_C)}(undef, l)
    @inbounds @simd for i in 1:length(sol_X)
        sol_X[i] = sol_C[i + 1] - sol_C[i]
    end
    return sol_X
end;


function simulate_data(l, i₀, β)
    prob = ODEProblem(sir_ode, @SArray([990.0, 10.0, 0.0, 0.0]), (0.0, l), @SArray([β, 10.0, 0.25]))
    X = sir_ode_solve(prob, l, i₀, β)
    Y = rand.(Poisson.(X))
    return X, Y
end;

struct ODEFused{I, B, L, F} <: DiscreteMultivariateDistribution
    i₀::I
    β::B
    l::L
    problem::F
end;

function BayesBase.logpdf(ode::ODEFused, y::Vector)
    sol_X = sir_ode_solve(ode.problem, ode.l, ode.i₀, ode.β)
    # `sum` over individual entries of the result of the `ODE` solver
    sumlpdf = sum(zip(sol_X, y)) do (x_i, y_i)
        return logpdf(Poisson(abs(x_i)), y_i)
    end
    # `clamp` to avoid infinities in the beginning, where 
    # priors are completely off
    return clamp(sumlpdf, -100000, Inf)
end;

function BayesBase.insupport(ode::ODEFused, y::Vector)
    return true
end

function BayesBase.mean(p::PointMass{D}) where { D <: ODEProblem }
    return p.point
end;


@node ODEFused Stochastic [ y, i₀, β, l, problem ];

@model function bayes_sir(y)
    l = length(y)
    prob = ODEProblem(sir_ode, @SArray([990.0, 10.0, 0.0, 0.0]), (0.0, l), @SArray([0.05, 10.0, 0.25]))    
    i₀ ~ Beta(1.0, 1.0)
    β  ~ Beta(1.0, 1.0)
    y  ~ ODEFused(i₀, β, l, prob)
end;

@constraints function sir_constraints()
    parameters = ProjectionParameters(
        strategy = ExponentialFamilyProjection.ControlVariateStrategy(nsamples = 40),
        direction = ExponentialFamilyProjection.BoundedNormUpdateRule(10),
        niterations = 1000
    )

    # In principle different parameters can be used for different states
    q(i₀) :: ProjectedTo(Beta; parameters = parameters)
    q(β) :: ProjectedTo(Beta; parameters = parameters)

    # `MeanField` is required for `NodeFunctionRuleFallback`
    q(i₀, β) = MeanField()
end;


@initialization function sir_initialization()
    q(β)  = Beta(1, 1)
    q(i₀) = Beta(1, 1)
end;

β_true = 0.05
i₀_true = 0.01
l = 40
X, Y = simulate_data(l, i₀_true, β_true);

ts = 1.0:1.0:l
plot(ts, X, label="Deterministic mean", xlabel="Time", ylabel="Daily incidence")
scatter!(ts, Y, label="Simulated observations")

niter = 15
result = infer(
        model = bayes_sir(),
        data  = (y = Y, ),
        constraints = sir_constraints(),
        initialization = sir_initialization(),
        iterations = niter,
        showprogress = false,
        options = (
            # Read `https://reactivebayes.github.io/RxInfer.jl/stable/manuals/inference/undefinedrules/`
            rulefallback = NodeFunctionRuleFallback(),
        )
);

pl_β_mean_i = plot(1:15, [mean(x) for x in result.posteriors[:β]], label=false, xlabel="Iteration", ylabel="Mean", title="β")
pl_i₀_mean_i = plot(1:15, [mean(x) for x in result.posteriors[:i₀]], label=false, xlabel="Iteration", ylabel="Mean", title="i₀")
plot(pl_β_mean_i, pl_i₀_mean_i, layout=(1,2), plot_title="Mean of posterior by iteration")

posterior_i₀ = result.posteriors[:i₀][end]
posterior_β = result.posteriors[:β][end];

mean_var(posterior_i₀) # Should be 0.01
mean_var(posterior_β) # Should be 0.05


p1 = plot(0.0:0.0001:0.02, x -> pdf(posterior_i₀, x); label="q(i₀)")
vline!(p1, [i₀_true], label=false)
p2 = plot(0.04:0.0001:0.06, x -> pdf(posterior_β, x); label="q(β)")
vline!(p2, [β_true], label=false)
plot(p1, p2)

@benchmark infer(
        model = bayes_sir(),
        data  = (y = Y, ),
        constraints = sir_constraints(),
        initialization = sir_initialization(),
        iterations = niter,
        showprogress = false,
        options = (
            # Read `https://reactivebayes.github.io/RxInfer.jl/stable/manuals/inference/undefinedrules/`
            rulefallback = NodeFunctionRuleFallback(),
        )
)

nsims = 100
i₀_mean = Array{Float64}(undef, nsims)
β_mean = Array{Float64}(undef, nsims)
i₀_coverage = Array{Float64}(undef, nsims)
β_coverage = Array{Float64}(undef, nsims)
Threads.@threads for i in ProgressBar(1:nsims) 
    X_sim, Y_sim = simulate_data(l, i₀_true, β_true)
    r = infer(
              model = bayes_sir(),
              data  = (y = Y_sim, ),
              constraints = sir_constraints(),
              initialization = sir_initialization(),
              iterations = niter,
              showprogress = false,
              options = ( rulefallback = NodeFunctionRuleFallback(), ))
    i0 = r.posteriors[:i₀][end]
    i₀_mean[i] = mean(i0)
    i0_cov = cdf(i0, i₀_true)
    b = r.posteriors[:β][end]
    β_mean[i] = mean(b)
    b_cov = cdf(b, β_true)
    i₀_coverage[i] = i0_cov
    β_coverage[i] = b_cov
end;

# Convenience function to check if the true value is within the credible interval
function in_credible_interval(x, lwr=0.025, upr=0.975)
    return x >= lwr && x <= upr
end;

pl_β_coverage = histogram(β_coverage, bins=0:0.1:1.0, label=false, title="β", ylabel="Density", density=true, xrotation=45, xlim=(0.0,1.0))
pl_i₀_coverage = histogram(i₀_coverage, bins=0:0.1:1.0, label=false, title="i₀", ylabel="Density", density=true, xrotation=45, xlim=(0.0,1.0))
plot(pl_β_coverage, pl_i₀_coverage, layout=(1,2), plot_title="Distribution of CDF of true value")

sum(in_credible_interval.(β_coverage)) / nsims

sum(in_credible_interval.(i₀_coverage)) / nsims

pl_β_mean = histogram(β_mean, label=false, title="β", ylabel="Density", density=true, xrotation=45, xlim=(0.045, 0.055))
vline!([β_true], label="True value")
pl_i₀_mean = histogram(i₀_mean, label=false, title="i₀", ylabel="Density", density=true, xrotation=45, xlim=(0.0,0.02))
vline!([i₀_true], label="True value")
plot(pl_β_mean, pl_i₀_mean, layout=(1,2), plot_title="Distribution of posterior means")

The respective proportion inside the credible intervals are 1.0 and 0.86 repsectivly for $\beta$ and $i_{o}$.

The main difference here is this change

parameters = ProjectionParameters(
        strategy = ExponentialFamilyProjection.ControlVariateStrategy(nsamples = 40),
        direction = ExponentialFamilyProjection.BoundedNormUpdateRule(10),
        niterations = 1000
    )

The higher is the limit parameter of the ExponentialFamilyProjection.BoundedNormUpdateRule(10) the more aggressive is optimisation. And the number of iterations is the hard cap of maximum gradient steps, be default it is 100. So here I am doing more gradients steps on each mean field iteration, however each one is a cheaper then ones before.

2. Refine moded or/and refine constraints

I am not sure that Beta distributions are actually good form to represent the posterior over $\beta$ and $i_{o}$. I think more reasonable one will be Continuous Bernouill https://en.wikipedia.org/wiki/Continuous_Bernoulli_distribution, however CB distribution is not integrated with Projections yet.

Just to show a point I made a such model with joint prior over $\beta$ and $i_{o}$ to be Direchlet distribution.

using RxInfer
using OrdinaryDiffEq
using ExponentialFamilyProjection
using Optimisers
using StableRNGs
using StatsFuns
using StaticArrays
using Plots
using BenchmarkTools
import BayesBase;
using ProgressBars
using Base.Threads

function sir_ode(u, p, t)
    (S,I,R,C) = u
    (β,c,γ) = p
    N = S+I+R
    infection = β*c*I/N*S
    recovery = γ*I
    return @SArray([
         -infection,
        infection - recovery,
        recovery,
        infection
    ]) # S, I, R, C
end;


function sir_ode_solve(problem, l, i₀, β)
    I = i₀*1000.0
    S = 1000.0 - I
    u0 = @SArray([S, I, 0.0, 0.0])
    p = @SArray([β, 10.0, 0.25])
    prob = remake(problem; u0=u0, p=p)
    sol = solve(prob, Tsit5(), saveat = 1.0)
    sol_C = view(sol, 4, :) # Cumulative cases
    sol_X = Array{eltype(sol_C)}(undef, l)
    @inbounds @simd for i in 1:length(sol_X)
        sol_X[i] = sol_C[i + 1] - sol_C[i]
    end
    return sol_X
end;


function simulate_data(l, i₀, β)
    prob = ODEProblem(sir_ode, @SArray([990.0, 10.0, 0.0, 0.0]), (0.0, l), @SArray([β, 10.0, 0.25]))
    X = sir_ode_solve(prob, l, i₀, β)
    Y = rand.(Poisson.(X))
    return X, Y
end;

struct ODEFused{IB, L, F} <: DiscreteMultivariateDistribution
    i_β::IB
    l::L
    problem::F
end;

function BayesBase.logpdf(ode::ODEFused, y::Vector)
    sol_X = sir_ode_solve(ode.problem, ode.l, ode.i_β[1], ode.i_β[2])
    # `sum` over individual entries of the result of the `ODE` solver
    sumlpdf = sum(zip(sol_X, y)) do (x_i, y_i)
        return logpdf(Poisson(abs(x_i)), y_i)
    end
    # `clamp` to avoid infinities in the beginning, where 
    # priors are completely off
    return clamp(sumlpdf, -100000, Inf)
end;

function BayesBase.insupport(ode::ODEFused, y::Vector)
    return true
end

function BayesBase.mean(p::PointMass{D}) where { D <: ODEProblem }
    return p.point
end;


@node ODEFused Stochastic [ y, iβ, l, problem ];

@model function bayes_sir(y)
    l = length(y)
    prob = ODEProblem(sir_ode, @SArray([990.0, 10.0, 0.0, 0.0]), (0.0, l), @SArray([0.05, 10.0, 0.25]))    
    i_β ~ Dirichlet([1, 1, 0.0001])
    y  ~ ODEFused(i_β, l, prob)
end;

i₀_true = 0.01
β_true = 0.05
l = 40
X, Y = simulate_data(l, i₀_true, β_true);

function experiment(Y, nsamples = 20, niter=1)

    @constraints function sir_constraints()
        parameters = ProjectionParameters(
            strategy = ExponentialFamilyProjection.ControlVariateStrategy(nsamples = nsamples)
        )

        # In principle different parameters can be used for different states
        q(i_β) :: ProjectedTo(Dirichlet, 3; parameters = parameters)
    end;


    @initialization function sir_initialization()
        q(i_β)  = Dirichlet([1, 1, 0.001])
    end;

    result = infer(
            model = bayes_sir(),
            data  = (y = Y, ),
            constraints = sir_constraints(),
            initialization = sir_initialization(),
            iterations = niter,
            showprogress = false,
            options = (
                # Read `https://reactivebayes.github.io/RxInfer.jl/stable/manuals/inference/undefinedrules/`
                rulefallback = NodeFunctionRuleFallback(),
            ),
    );

    posterior = result.posteriors[:i_β][end]

    posterior_i = Beta(posterior.alpha[1], posterior.alpha[2]+posterior.alpha[3])
    posterior_β = Beta(posterior.alpha[2], posterior.alpha[1]+posterior.alpha[3])

    
    return cdf(posterior_i, i₀_true), cdf(posterior_β, β_true)
end

nsims = 1000
i₀_coverage = Array{Float64}(undef, nsims)
β_coverage = Array{Float64}(undef, nsims)

for i in ProgressBar(1:nsims)
    _, Y = simulate_data(l, i₀_true, β_true);
    creds = experiment(Y)
    i₀_coverage[i] = creds[1]
    β_coverage[i] = creds[2]
end

function in_credible_interval(x, lwr=0.025, upr=0.975)
    return x >= lwr && x <= upr
end;

sum(in_credible_interval.(β_coverage)) / nsims

sum(in_credible_interval.(i₀_coverage)) / nsims

This models gives 1.0 and 1.0 also being substantially faster, because here we do not need to run iterations and 2 projections inside them, but just one gradient scheme. But you need to be aware that this model is more simple one it inferring only 3 parameters for posteriors not 4 like one we have with 2 betas.

[pasq-cat]

Good idea! @Rockdeldiablo is actually currently working on conformalized Bayes for his GSoC project, so we can probably draw on that.

I'm on vacation at the moment for the next two weeks. Can you ping me again here or on Slack if I forget to get back to this? Cheers

yeah. i am stuck right now, tomorrow me and mojtaba will work on it together. the score function is ready but the fit method is still giving me some trouble.

[pat-alt]

Good idea! @Rockdeldiablo is actually currently working on conformalized Bayes for his GSoC project, so we can probably draw on that.

I'm on vacation at the moment for the next two weeks. Can you ping me again here or on Slack if I forget to get back to this? Cheers

yeah. i am stuck right now, tomorrow me and mojtaba will work on it together. the score function is ready but the fit method is still giving me some trouble.

Obviously no call for action!! I just meant to say that you're working on interesting things on our end that we can probably leverage in the future in this context 😄

[albertpod]

Hi @sdwfrost,

I've reviewed your gist and think your understanding is excellent! It's a great example! We would be happy to reference it in our documentation, or you are welcome to make a PR to RxInferExamples.jl, either works.

As for coverage, we’ll need to consider our options.

[sdwfrost]

Thanks everyone, you've all been super helpful. It's interesting @Nimrais that changing some config parameters affects the coverage (though a coverage of 1.0 in some ways is as problematic as a coverage of 0.9). Though it is a bit like comparing apples and oranges, I'm keen to see how RxInfer can mesh with conformal prediction concepts. I'm going to flag a reply as an answer to this thread, now that you have a DiffEq example, if that's OK.

[bvdmitri]

Addressed in https://examples.rxinfer.com/categories/problem_specific/ode_parameter_estimation/ 🥳

[ChrisRackauckas]

But this isn't solved or resolved though, not addressed in that tutorial. That shows that if you use a fixed time step method you can make it work, that I'm not surprised. What I wanted to see is something with a dynamic compute graph via adaptive time stepping, using something like an implicit method on a difficult equation like Rosenbrock.

[bvdmitri]

Thanks for the feedback, @ChrisRackauckas. In the context of this discussion, I would argue the issue was resolved—at least based on what was originally stated. Now, it seems like what you actually wanted to see is something quite different, but that wasn’t clear from the title or the initial description.

Most of the follow-up questions were about how RxInfer is implemented, not about this specific problem. If that had been the focus from the start, we could have approached it differently. For example, @sdwfrost was very clear in defining the problem, providing an example and some preliminary code, which made it much easier for us to engage and explore potential (even if imperfect) solutions.

I’d really encourage you to start a new discussion and be more explicit about what you’re looking for from the beginning. Shifting the context into a new direction after people have already spent time trying to help isn’t exactly helpful. Some code examples would be helpful as well.

[ChrisRackauckas]

Sorry, I thought it was clear from Simon's responses that the goal would be to showcase RXInfer in a way that is similar to Turing, DynamicHMC, Stan, etc. integrations, and ultimately add it to the Julia-wide benchmarks on inference that have been maintained for this testing.

• https://docs.sciml.ai/SciMLBenchmarksOutput/stable/BayesianInference/DiffEqBayesFitzHughNagumo/
• https://docs.sciml.ai/SciMLBenchmarksOutput/stable/BayesianInference/DiffEqBayesLorenz/
• https://docs.sciml.ai/SciMLBenchmarksOutput/stable/BayesianInference/DiffEqBayesLotkaVolterra/

The main issue with the demonstration in https://examples.rxinfer.com/categories/problem_specific/ode_parameter_estimation/ is it moves the goal post into a simpler domain of quasi-static compute graphs. I mentioned that before you got started right here https://github.com/orgs/ReactiveBayes/discussions/56#discussioncomment-4882428 that it would be the major point of the discussion. The issue of course is highlighted in many things around the Julia ecosystem, especially things that talk about AD like https://www.stochasticlifestyle.com/useful-algorithms-that-are-not-optimized-by-jax-pytorch-or-tensorflow/, where inference algorithms that depend on having a fixed compute graph do not easily translate to settings where the compute graph is not fixed. The issue then of course is that all of the algorithms that apply to real problems require more dynamic handling as the adaptivity is pretty essential to making stiff ODE solving robust, as without rejections it can be difficult to ensure the Newton iterations are convergent. While there is a comment that we can just swap out to real methods, that's a bit of leap of faith since we know that many inference algorithms fail to make that exact jump. That doc example would be much more compelling if it used an adaptive solver like Tsit5.

If it just works with adaptivity thrown in there though, great! But it's not possible to tell that can be done from the example as written, and that makes a massive difference to complexity the tutorial demonstrates. The real question there is, is the only compiler-level handling being done by RXInfer the use of AD are there other tricks that need to be done? Gen.jl in particular had some issues moving to the dynamic compute graph space because of how they did some graph analyses, and ultimately the question is whether RXInfer has the same issue, or whether it's as generic as something like Turing.jl where all you need is AD support.

If that question is answered then we can take it from there. If that's the case, then we'd start subjecting it to a barrage of tests that traditional Bayesian estimation methods have had difficulties with, such as:

• https://docs.sciml.ai/Catalyst/stable/inverse_problems/examples/ode_fitting_oscillation/

along with replicating https://docs.sciml.ai/Overview/stable/showcase/missing_physics/ or the Bayesian versions in https://arxiv.org/abs/2012.07244, but that can wait.

[bvdmitri]

Thanks for clarifications! It will take some time to go through all the examples and attached links but that is very helpful indeed and is more clear! We will discuss if we can improve the tutorial and add more compelling examples.

[ismailsenoz]

Hi @ChrisRackauckas, you can check out https://github.com/ReactiveBayes/RxInferExamples.jl/blob/lorenz_example/examples/Problem%20Specific/Lorenz%20Differential%20Equation/Lorenz%20Differential%20Equation.ipynb on how to integrate adaptive solvers with RxInfer. However, this functionality requires methods that are not open sourced yet.

RxInfer requires having access to one-time-step integral approximation functions. Such as:

function rkf45_step(f::Function, x, θ, t, dt; tol=1e-6, safety=0.9)
    # Butcher tableau coefficients for RKF45
    a21 = 1/4
    a31, a32 = 3/32, 9/32
    a41, a42, a43 = 1932/2197, -7200/2197, 7296/2197
    a51, a52, a53, a54 = 439/216, -8, 3680/513, -845/4104
    a61, a62, a63, a64, a65 = -8/27, 2, -3544/2565, 1859/4104, -11/40
    
    # Coefficients for 4th order solution
    b41, b42, b43, b44, b45 = 25/216, 0, 1408/2565, 2197/4104, -1/5
    
    # Coefficients for 5th order solution
    b51, b52, b53, b54, b55, b56 = 16/135, 0, 6656/12825, 28561/56430, -9/50, 2/55
    
    # Error estimation coefficients (difference between 5th and 4th order solutions)
    e1 = b51 - b41
    e2 = b52 - b42
    e3 = b53 - b43
    e4 = b54 - b44
    e5 = b55 - b45
    e6 = b56 - 0     # b46 is 0 since k6 is not used in the 4th order solution
    
    # Compute the six stages
    k1 = f(x, θ, t)
    k2 = f(x + dt*a21*k1, θ, t + dt/4)
    k3 = f(x + dt*(a31*k1 + a32*k2), θ, t + 3dt/8)
    k4 = f(x + dt*(a41*k1 + a42*k2 + a43*k3), θ, t + 12dt/13)
    k5 = f(x + dt*(a51*k1 + a52*k2 + a53*k3 + a54*k4), θ, t + dt)
    k6 = f(x + dt*(a61*k1 + a62*k2 + a63*k3 + a64*k4 + a65*k5), θ, t + dt/2)
    
    # 4th order solution
    x4 = x + dt*(b41*k1 + b43*k3 + b44*k4 + b45*k5)
    
    # 5th order solution
    x5 = x + dt*(b51*k1 + b53*k3 + b54*k4 + b55*k5 + b56*k6)
    
    # Error estimation
    err = norm(x5 - x4)
    err = max(err, eps())
    if err >= tol
        dt_next = dt * safety * (tol/err)^(1/5)
    else
        dt_next = dt * safety * (tol/err)^(1/4)
    end

    return clamp.(x5, -100, 100), dt_next
end

Then, we define the integrator as a non-linear deterministic node and do message passing around it. Since the nonlinear node also returns $\mathrm{d}t$, we can do inference over the discretization interval as well. I am sure there are better and improved versions of 'rkf45_step' functions in DifferentialEquations.jl. If you point us to how we can utilize these integrators from DifferentialEquations.jl then we can use the integrators with improved Butcher tableau. Perhaps we can even write@ode or @sde macros within @model macro to create the non-linear deterministic node directly from the solvers that are available in DifferentialEquations.jl ecosystem.