Help with getting a model working in RxInfer

[sdwfrost]

I'm trying to get a model working in RxInfer.jl, based on a previous example that I put together with your help. I get an error ERROR: MethodError: no method matching -(::Int64, ::GraphPPL.VariableRef{…}) from the below. The underlying model is quite simple (bar the beta priors), but I use a function to compute probabilities that I pass to a multinomial - the call to this function (simulate) is the one that fails. Do I need to wrap up the function call (which is deterministic) into a node?

using OrdinaryDiffEq
using DiffEqCallbacks
using Distributions
using Interpolations
using RxInfer
using ExponentialFamilyProjection
using Random
using Plots

# Utility function to turn an exponential rate into a proportion

@inline function rate_to_proportion(r, dt)
    1-exp(-r*dt)
end

# Define dynamical model

function sir_map!(du,u,p,t)
    (S,I,C,Y) = u # Susceptible, Infected, Cumulative infections, New infections
    (β,γ) = p
    infection = rate_to_proportion(β*I)*S
    recovery = rate_to_proportion(γ)*I
    @inbounds begin
        du[1] = S-infection
        du[2] = I+infection-recovery
        du[3] = C+infection
        du[4] = infection
    end
    nothing
end

# Wrapper function for the above

function simulate(β, γ, ρ, nsteps)
    tspan = (0, nsteps)
    u0 = [1-ρ, ρ, 0.0, 0.0] # Initial conditions: S, I, C, Y
    p = [β, γ]
    prob = DiscreteProblem(sir_map!, u0, tspan, p)
    sol = solve(prob, FunctionMap())
    τ = sol[end][3] # Final size
    f = sol[4,2:end] # New infections per time step
    f = clamp.(f,0.0,1.0) # Ensure stays in range
    f = f/τ # Normalize frequencies to sum to 1
    return sol, τ, f
end

function sample_infection_times(β, γ, ρ, nsteps, N)
    sol, τ, f = simulate(β, γ, ρ, nsteps)
    times = sol.t # Simulation times
    cdfτ = [sol[i][3]/τ for i in 1:length(sol.t)] # CDF evaluated at `times`, obtained by C/τ
    invcdfτ = ConstantInterpolation(cdfτ, times) # Inverse CDF for sampling
    K = rand(Binomial(N, τ)) # Number of infections during the epidemic
    ti = invcdfτ.(rand(K)) # Append infection times for K individuals
    return ti
end

function sample_recovery_times(γ, M, K)
    pγ = rate_to_proportion(γ, 1.0) # Convert recovery rate to proportion
    delta = 1.0 .+ rand(Geometric(pγ), M + K) # Recovery intervals for M+K infected individuals
    return delta
end

function time_frequencies(times, nsteps)
    freq = zeros(Int, nsteps)
    @inbounds for t in times
        1 ≤ t ≤ nsteps && (freq[t] += 1)
    end
    return freq
end

@model function sds_discrete_rxinfer(N, M, K, th, delta, nsteps)
    β ~ Beta(1.0, 1.0)
    γ ~ Beta(1.0, 1.0)
    ρ ~ Beta(1.0, 1.0)
    (_, τ, f) = simulate(β, γ, ρ, nsteps)
    K ~ Binomial(N, τ) # Number of infections during the epidemic
    th ~ Multinomial(K, f) # Infection times for K individuals
    pγ = rate_to_proportion(γ)
    for i in 1:(M+K)
        delta[i] ~ Geometric(pγ) # Recovery durations
    end
end

@constraints function sds_constraints()
    parameters = ProjectionParameters(
        strategy = ExponentialFamilyProjection.ControlVariateStrategy(nsamples = 200)
    )
    # In principle different parameters can be used for different states
    q(β) :: ProjectedTo(Beta; parameters = parameters)
    q(γ) :: ProjectedTo(Beta; parameters = parameters)
    q(ρ) :: ProjectedTo(Beta; parameters = parameters)

    # `MeanField` is required for `NodeFunctionRuleFallback`
    q(β,γ,ρ) = MeanField()
end

@initialization function sds_initialization()
    q(β) = Beta(1.0, 1.0)
    q(γ) = Beta(1.0, 1.0)
    q(ρ) = Beta(1.0, 1.0)
end

# Parameter values

nsteps = 40   # Number of time steps to simulate
tspan = (0, nsteps)   # Time span for the simulation
u0 = [0.99, 0.01, 0.0, 0.0] # Initial conditions: S, I, C, Y
β = 0.5 # Infection rate
γ = 0.25 # Recovery rate
ρ = 0.01 # Proportion of individuals initially infected
N = 99 # Number of susceptibles sampled at t=0
M = 1  # Number of initially infected individuals

ti = sample_infection_times(β, γ, ρ, nsteps, N) # Sample infection times
K = length(ti) # Number of infections during the epidemic
th = time_frequencies(ti, nsteps) # Time frequencies of infections
delta = sample_recovery_times(γ, M, K) # Sample recovery times

niter = 10
result = infer(
        model = sds_discrete_rxinfer(),
        data  = (N=N, M=M, K=K, th=th, delta=delta, nsteps=nsteps,),
        constraints = sds_constraints(),
        initialization = sds_initialization(),
        iterations = niter,
        showprogress = false,
        options = (
            # Read `https://reactivebayes.github.io/RxInfer.jl/stable/manuals/inference/undefinedrules/`
            rulefallback = NodeFunctionRuleFallback(),
        )
)

[wouterwln]

Hi @sdwfrost . Indeed the problem right now is in the simulate function. The problem is the following: Since RxInfer creates a factor graph representation of your generative model, the value of ρ is not known when it is being created in the factor graph. Now, under the hood, in your model macro, ρ now does not refer to a numerical value, but is the reference to the variable in the factor graph. When you then call (_, τ, f) = simulate(β, γ, ρ, nsteps), the actual simulate function is called with the provided arguments (which are the internal graph references). We made it like this because we want users to be able to execute arbitrary Julia code in the model macro with values that are known at model creation time. As a rule of thumb: if variables to be inferred enter a function, you should probably wrap that function in a factor node (with := probably). Afterwards, you will probably run into the next problem: RxInfer only accepts single random variables on the left hand side of the ~ or := operation, so even (_, τ, f) := simulate(β, γ, ρ, nsteps) will give you some trouble. This is because we cannot evaluate simulate at model creation time because the arguments are unknown, so we cannot properly generate a return type signature for this particular function without introducing ambiguity in the modeling language. Furthermore, if we were to accept multiple arguments on the left hand side of the ~ operator, we would have to support statements like m, y ~ NormalMeanVariance(v = v), which opens Pandora's box in terms of model definition. At this point, tinkering with your model for me would mean making changes in your model definition up to the point where I might lose track of your initial goal in terms of inference results. Would it be possible to redefine your model using these guidelines? Then we can look at running inference in your model afterwards.

[sdwfrost]

Thanks @wouterwln! I modified as you suggested, and made a few extra changes (from the examples, it seems as though I need to use BinomialPolya and MultinomialPolya, but now I'm running into the issue that the geometric distribution is not supported - can you suggest a workaround? Current code below:

using OrdinaryDiffEq
using DiffEqCallbacks
using Distributions
using Interpolations
using RxInfer
using ExponentialFamilyProjection
using Random
using Plots

# Utility function to turn an exponential rate into a proportion

@inline function rate_to_proportion(r)
    1-exp(-r)
end

# Define dynamical model

function sir_map!(du,u,p,t)
    (S,I,C,Y) = u # Susceptible, Infected, Cumulative infections, New infections
    (β,γ) = p
    infection = rate_to_proportion(β*I)*S
    recovery = rate_to_proportion(γ)*I
    @inbounds begin
        du[1] = S-infection
        du[2] = I+infection-recovery
        du[3] = C+infection
        du[4] = infection
    end
    nothing
end

# Wrapper functions for the above

function simulate(β, γ, ρ, nsteps)
    tspan = (0, nsteps)
    u0 = [1-ρ, ρ, 0.0, 0.0] # Initial conditions: S, I, C, Y
    p = [β, γ]
    prob = DiscreteProblem(sir_map!, u0, tspan, p)
    sol = solve(prob, FunctionMap())
    τ = sol[end][3] # Final size
    f = sol[4,2:end] # New infections per time step
    f = clamp.(f,0.0,1.0) # Ensure stays in range
    f = f/τ # Normalize frequencies to sum to 1
    return sol, τ, f
end

function pdf(β, γ, ρ, nsteps)
    (_, τ, f) = simulate(β, γ, ρ, nsteps)
    return f
end

function final_size(β, γ, ρ, nsteps)
    (_, τ, f) = simulate(β, γ, ρ, nsteps)
    return τ
end

function sample_infection_times(β, γ, ρ, nsteps, N)
    sol, τ, f = simulate(β, γ, ρ, nsteps)
    times = sol.t # Simulation times
    cdfτ = [sol[i][3]/τ for i in 1:length(sol.t)] # CDF evaluated at `times`, obtained by C/τ
    invcdfτ = ConstantInterpolation(cdfτ, times) # Inverse CDF for sampling
    K = rand(Binomial(N, τ)) # Number of infections during the epidemic
    ti = invcdfτ.(rand(K)) # Append infection times for K individuals
    return ti
end

function sample_recovery_times(γ, M, K)
    pγ = rate_to_proportion(γ, 1.0) # Convert recovery rate to proportion
    delta = 1.0 .+ rand(Geometric(pγ), M + K) # Recovery intervals for M+K infected individuals
    return delta
end

function time_frequencies(times, nsteps)
    freq = zeros(Int, nsteps)
    @inbounds for t in times
        1 ≤ t ≤ nsteps && (freq[t] += 1)
    end
    return freq
end

@model function sds_discrete_rxinfer(N, M, K, th, delta, nsteps)
    β ~ Beta(1.0, 1.0)
    γ ~ Beta(1.0, 1.0)
    ρ ~ Beta(1.0, 1.0)
    α ~ MvNormalWeightedMeanPrecision(0.0, 1.0)
    f := pdf(β, γ, ρ, nsteps)
    τ := final_size(β, γ, ρ, nsteps)
    K ~ BinomialPolya(N, τ, α) # Number of infections during the epidemic
    th ~ MultinomialPolya(K, f) # Infection times for K individuals
    pγ := rate_to_proportion(γ)
    for i in eachindex(delta)
        delta[i] ~ Geometric(pγ) # Recovery durations
    end
end

@constraints function sds_constraints()
    parameters = ProjectionParameters(
        strategy = ExponentialFamilyProjection.ControlVariateStrategy(nsamples = 200)
    )
    # In principle different parameters can be used for different states
    q(β) :: ProjectedTo(Beta; parameters = parameters)
    q(γ) :: ProjectedTo(Beta; parameters = parameters)
    q(ρ) :: ProjectedTo(Beta; parameters = parameters)

    # `MeanField` is required for `NodeFunctionRuleFallback`
    q(β,γ,ρ) = MeanField()
end

@initialization function sds_initialization()
    q(β) = Beta(1.0, 1.0)
    q(γ) = Beta(1.0, 1.0)
    q(ρ) = Beta(1.0, 1.0)
end

# Parameter values

nsteps = 40   # Number of time steps to simulate
tspan = (0, nsteps)   # Time span for the simulation
u0 = [0.99, 0.01, 0.0, 0.0] # Initial conditions: S, I, C, Y
β = 0.5 # Infection rate
γ = 0.25 # Recovery rate
ρ = 0.01 # Proportion of individuals initially infected
N = 99 # Number of susceptibles sampled at t=0
M = 1  # Number of initially infected individuals

ti = sample_infection_times(β, γ, ρ, nsteps, N) # Sample infection times
K = length(ti) # Number of infections during the epidemic
th = time_frequencies(ti, nsteps) # Time frequencies of infections
delta = sample_recovery_times(γ, M, K) # Sample recovery times

niter = 10
result = infer(
        model = sds_discrete_rxinfer(),
        data  = (N=N, M=M, K=K, th=th, delta=delta, nsteps=nsteps,),
        constraints = sds_constraints(),
        initialization = sds_initialization(),
        iterations = niter,
        showprogress = false,
        options = (
            # Read `https://reactivebayes.github.io/RxInfer.jl/stable/manuals/inference/undefinedrules/`
            rulefallback = NodeFunctionRuleFallback(),
        )
)

[sdwfrost]

I'm guessing that a geometric distribution would need to be added - in the meantime, I replaced this with a Poisson to debug downstream. My model now compiles, but the variables aren't updating, which means I need to expand the initialization - can you see what nodes need to be initialized?

using OrdinaryDiffEq
using DiffEqCallbacks
using Distributions
using Interpolations
using RxInfer
using ExponentialFamilyProjection
using Random
using Plots

# Utility function to turn an exponential rate into a proportion

@inline function rate_to_proportion(r)
    1-exp(-r)
end

# Define dynamical model

function sir_map!(du,u,p,t)
    (S,I,C,Y) = u # Susceptible, Infected, Cumulative infections, New infections
    (β,γ) = p
    infection = rate_to_proportion(β*I)*S
    recovery = rate_to_proportion(γ)*I
    @inbounds begin
        du[1] = S-infection
        du[2] = I+infection-recovery
        du[3] = C+infection
        du[4] = infection
    end
    nothing
end

# Wrapper functions for the above

function simulate(β, γ, ρ, nsteps)
    tspan = (0, nsteps)
    u0 = [1-ρ, ρ, 0.0, 0.0] # Initial conditions: S, I, C, Y
    p = [β, γ]
    prob = DiscreteProblem(sir_map!, u0, tspan, p)
    sol = solve(prob, FunctionMap())
    τ = sol[end][3] # Final size
    f = sol[4,2:end] # New infections per time step
    f = clamp.(f,0.0,1.0) # Ensure stays in range
    f = f/τ # Normalize frequencies to sum to 1
    return sol, τ, f
end

function pdf(β, γ, ρ, nsteps)
    (_, τ, f) = simulate(β, γ, ρ, nsteps)
    return f
end

function final_size(β, γ, ρ, nsteps)
    (_, τ, f) = simulate(β, γ, ρ, nsteps)
    return τ
end

function sample_infection_times(β, γ, ρ, nsteps, N)
    sol, τ, f = simulate(β, γ, ρ, nsteps)
    times = sol.t # Simulation times
    cdfτ = [sol[i][3]/τ for i in 1:length(sol.t)] # CDF evaluated at `times`, obtained by C/τ
    invcdfτ = ConstantInterpolation(cdfτ, times) # Inverse CDF for sampling
    K = rand(Binomial(N, τ)) # Number of infections during the epidemic
    ti = invcdfτ.(rand(K)) # Append infection times for K individuals
    return ti
end

function sample_recovery_times(γ, M, K)
    pγ = rate_to_proportion(γ, 1.0) # Convert recovery rate to proportion
    delta = 1.0 .+ rand(Geometric(pγ), M + K) # Recovery intervals for M+K infected individuals
    return delta
end

function time_frequencies(times, nsteps)
    freq = zeros(Int, nsteps)
    @inbounds for t in times
        1 ≤ t ≤ nsteps && (freq[t] += 1)
    end
    return freq
end

@model function sds_discrete_rxinfer(N, M, K, th, delta, nsteps)
    β ~ Beta(1.0, 1.0)
    γ ~ Beta(1.0, 1.0)
    ρ ~ Beta(1.0, 1.0)
    α ~ MvNormalWeightedMeanPrecision(zeros(1), diageye(1))
    f := pdf(β, γ, ρ, nsteps)
    τ := final_size(β, γ, ρ, nsteps)
    K ~ BinomialPolya(N, τ, α) where {
            dependencies = RequireMessageFunctionalDependencies(α = MvNormalWeightedMeanPrecision(zeros(1), diageye(1)))
        }
    th ~ MultinomialPolya(K, f) # Infection times for K individuals
    pγ := rate_to_proportion(γ)
    for i in eachindex(delta)
        #delta[i] ~ Geometric(pγ) # Recovery durations
        delta[i] ~ Poisson(pγ) # Recovery durations
    end
end

@constraints function sds_constraints()
    parameters = ProjectionParameters(
        strategy = ExponentialFamilyProjection.ControlVariateStrategy(nsamples = 200)
    )
    # In principle different parameters can be used for different states
    q(β) :: ProjectedTo(Beta; parameters = parameters)
    q(γ) :: ProjectedTo(Beta; parameters = parameters)
    q(ρ) :: ProjectedTo(Beta; parameters = parameters)

    # `MeanField` is required for `NodeFunctionRuleFallback`
    q(β,γ,ρ) = MeanField()
end

@initialization function sds_initialization()
    q(β) = Beta(1.0, 1.0)
    q(γ) = Beta(1.0, 1.0)
    q(ρ) = Beta(1.0, 1.0)
end

# Parameter values

nsteps = 40   # Number of time steps to simulate
tspan = (0, nsteps)   # Time span for the simulation
u0 = [0.99, 0.01, 0.0, 0.0] # Initial conditions: S, I, C, Y
β = 0.5 # Infection rate
γ = 0.25 # Recovery rate
ρ = 0.01 # Proportion of individuals initially infected
N = 99 # Number of susceptibles sampled at t=0
M = 1  # Number of initially infected individuals

ti = sample_infection_times(β, γ, ρ, nsteps, N) # Sample infection times
K = length(ti) # Number of infections during the epidemic
th = time_frequencies(ti, nsteps) # Time frequencies of infections
delta = sample_recovery_times(γ, M, K) # Sample recovery times

meta = @meta begin
    pdf(β, γ, ρ, nsteps) -> Linearization()
    final_size(β, γ, ρ, nsteps) -> Linearization()
    rate_to_proportion(γ) -> Linearization()
end

niter = 10
result = infer(
        model = sds_discrete_rxinfer(),
        meta = meta,
        data  = (N=N, M=M, K=K, th=th, delta=delta, nsteps=nsteps,),
        constraints = sds_constraints(),
        initialization = sds_initialization(),
        iterations = niter,
        showprogress = false,
        options = (
            # Read `https://reactivebayes.github.io/RxInfer.jl/stable/manuals/inference/undefinedrules/`
            rulefallback = NodeFunctionRuleFallback(),
        )
)

[wouterwln]

Thanks for the follow up! I think you're going in the right direction with your model. I can see a couple of issues with the way your model is defined, so I'll try and structure this story in the best way I can. If anything is unclear, feel free to ask!

General question

First of all, I see that you pass data M to the model, but it is never used. Is this obsolete or does it tie in with the rest of the model?

Inbound messages

I see that you read the documentation page on Bayesian Binomial Regression (kudos for that) and you have specified the RequireMessageFunctionalDependencies option. This is indeed correct, but there is a slight ambiguity there: The message you request and prespecified should be named according to the name of the interface relative to the node. This seems a bit weird, but at that point in creation the BinomialPolya node does not really know what α is (it is a variable in the model, and it is one of the neighboring edges in the factor graph, but consider the following:

@model function outer_model(data)
   α ~ something()
   data ~ inner(data = data)
end

@model function inner(data, β)
   α ~ something_else()
   data ~ my_likelihood(β, α) where {
      dependencies = RequireMessageFunctionalDependencies(α = my_initial_msg())
   }
end

We wouldn't really be able to resolve α because both neighbors will claim to be named α.) That's why you need to specify your inbound message dependency on β, which is the internal name of that interface for the BinomialPolya node ref. The same goes for the MultinomialPolya node, so that part of the model changes to:

  K ~ BinomialPolya(N, τ, α) where {
            dependencies = RequireMessageFunctionalDependencies(β = NormalMeanVariance(0.0, 1.0))
        }
    th ~ MultinomialPolya(K, f) where {
            dependencies = RequireMessageFunctionalDependencies(ψ = NormalMeanVariance(0.0, 1.0))
        }# Infection times for K individuals

Resolving the constraints

In their current form, the variational constraints would not allow for inference. The main problem for now is the usage of the BinomialPolya node. The problem is that the message towards the β interface (in your example α) defined here requires a PointMass variational posterior distribution on n (which is τ in your model). This means that τ should either have a PointMass variational posterior, or should be data. The first you can achieve probably by including q(τ, α)=q(τ)q(α) and projecting τ to PointMass, but a bigger problem is that we do not have inference rules defined towards τ for the BinomialPolya node and it is generally assumed to be data. While there probably is some workaround, I don't feel qualified enough to really advise you on this issue (as I've never touched this node before) and here the question really becomes if "an inference result" means "a good inference result". Anyway I'm sure we can come up with something.

Specifing meta

I see you have specified the Linearization meta, which, unfortunately ref does not work with non-gaussian distributions. I would consider using the CVIProjection meta, which should serve as a drop-in replacement.

Initialization problem

At this point, I've changed quite a bit. The main thing is: If I initialize the messages towards β, γ, ρ with

    μ(β) = Beta(1.0, 1.0)
    μ(γ) = Beta(1.0, 1.0)
    μ(ρ) = Beta(1.0, 1.0)

I at least get the RuleNotFoundError instead of the variable not updating error (if I disable the rule fallback option in the inference call). This is because both pdf and final_size connect to all three variables. So if pdf wants to send the message to β, it requests the message towards γ from final_size, which requests the message towards β from pdf and there you have your circular dependency. That should at least get you closer to a solution, and I will discuss with the rest of the RxInfer team if there might be a clever workaround of there not being a rule for the n interface of the BinomialPolya node, as I don't think the rule fallback option is the best route in this scenario.

That all being said, I remember a certain discussion that spawned a certain workshop entry ;) where the idea basically is to wrap the entire complex likelihood into a single node, and defining a logpdf around this node

[wouterwln]

One more thing: We can add the Geometric node with the following code:

using Distributions
using RxInfer

@node Geometric Stochastic [out, p]

@rule Geometric(:p, Marginalisation) (q_out::PointMass,) = begin
    return Beta(2, mean(q_out) + 1)
end

This is a minimal implementation that only has the rule you require. An important thing to remember is that in Julia, the Geometric distribution counts the number of failures before a success (so it can be 0). This will give you a backwards message towards p that is Beta distributed. I will probably add this node and rule to RxInfer over the weekend, but for now you can also use it like this.