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.