How to embed a CubicSpline in an RxInfer model

[joshualeond]

Hi there, I've been working with the RxInfer.jl package and am really excited to see this work! I'm attempting to embed a fitted spline within a model like the following:

@model function SoC_stream(y, I, B, αw, βw, αv, βv, spline)
   λw ~ Gamma(shape = αw, rate = βw)
   λv ~ Gamma(shape = αv, rate = βv)
   SoC_prev ~ NormalMeanVariance(0.0, (0.01)^2)
   for t in 1:length(y)
       SoC_next[t] ~ NormalMeanPrecision(SoC_prev + B*I[t], λw)
       y[t] ~ NormalMeanPrecision(spline[SoC_next[t]], λv) # <---- Error happens here
       SoC_prev = SoC_next[t]
   end
end

where the spline is fit like the following:

spline = CubicSpline(soc_grid, ocv_grid, extrapl=[1,], extrapr=[1,])

Indexing the fitted spline simply returns the result of y given input x:

spline[0.5] = 3.3978251023417814

The error I'm seeing when I attempt to run infer on the SoC_stream model shown above is the following:

ERROR: MethodError: no method matching getindex(::CubicSpline{Float64}, ::GraphPPL.NodeLabel)
The function `getindex` exists, but no method is defined for this combination of argument types.

Closest candidates are:
  getindex(::CubicSpline, ::Real)
   @ CubicSplines C:\Users\usrname\.julia\packages\CubicSplines\sDUOi\src\CubicSplines.jl:147
  getindex(::CubicSpline, ::AbstractVector{<:Real})
   @ CubicSplines C:\Users\usrname\.julia\packages\CubicSplines\sDUOi\src\CubicSplines.jl:189

Using a fitted spline in Turing.jl seemed to work fine for the most part but I'm new to RxInfer.jl and am not really sure how to resolve this. Any info would be appreciated.

Here's the entire reproducible example:

using RxInfer, CubicSplines, Distributions, Random, Plots

# Simulate some data
Random.seed!(42)
n = 100
I = randn(n)                    # Simulated current measurements
soc_true = cumsum(0.01 .+ 0.1 .* randn(n))
soc_true = clamp.(soc_true, 0, 1)
ocv_data = 3.0 .+ 0.8 .* soc_true .+ 0.02 .* randn(n)
y = ocv_data

# Fit the spline
soc_grid = 0:0.1:1
ocv_grid = 3.0 .+ 0.8 .* soc_grid .+ 0.01 .* randn(length(soc_grid))
spline = CubicSpline(soc_grid, ocv_grid, extrapl=[1,], extrapr=[1,])

# Show indexing works
@show spline[0.5]

# Model definition
@model function SoC_stream(y, I, B, αw, βw, αv, βv, spline)
   λw ~ Gamma(shape = αw, rate = βw)
   λv ~ Gamma(shape = αv, rate = βv)
   SoC_prev ~ NormalMeanVariance(0.0, (0.01)^2)
   for t in 1:length(y)
       SoC_next[t] ~ NormalMeanPrecision(SoC_prev + B*I[t], λw)
       y[t] ~ NormalMeanPrecision(spline[SoC_next[t]], λv) # <---- Error happens here
       SoC_prev = SoC_next[t]
   end
end

# Set parameters
B = 0.01
αw, βw = 2.0, 2.0
αv, βv = 2.0, 2.0

# Attempt inference
init = @initialization begin
   # seed the plate’s marginals (same name as in @model)
   q(SoC_next) = NormalMeanVariance(0, (0.01)^2)
   # seed the two precision priors
   q(λw) = GammaShapeRate(αw, βw)
   q(λv) = GammaShapeRate(αv, βv)
end

constraints = @constraints begin
   # Group the entire SoC chain vs. the two precisions
   q(SoC_prev, SoC_next, λw, λv) = q(SoC_prev, SoC_next)q(λw)q(λv)
end

model = SoC_stream(B = B, αw = αw, βw = βw, αv = αv, βv = βv, spline = spline)

result = infer(
   model = model,
   data = (y = y, I = I),
   returnvars     = ( SoC_next = KeepEach(),
                      λw       = KeepLast(),
                      λv       = KeepLast() ),
   initialization = init,
   constraints = constraints,
   options        = (limit_stack_depth = 300,)
)

# Review results
margs = vcat(result.posteriors[:SoC_next]...)  # flatten Vector{Marginal}
soc_mean = [ mean(m) for m in margs ]

plot(soc_true)
plot!(soc_mean)

[ismailsenoz]

Hi @joshualeond! Thanks for trying out RxInfer. Below is a functioning example. Not sure of the quality of estimates as that is for you to judge. The problem was the following. In RxInfer you can not use the random variables as indexes. Instead spline needs to be wrapped to a function spline_fn such that RxInfer can work with it. Because you are using a nonlinear transformation of a random variable you would need an approximation method. I defined some auxiliary variables in the graph as those make life easier to assign factorization constraints.

using RxInfer, CubicSplines, Distributions, Random, Plots

# Simulate some data
Random.seed!(42)
n = 100
I = randn(n)                    # Simulated current measurements
soc_true = cumsum(0.01 .+ 0.1 .* randn(n))
soc_true = clamp.(soc_true, 0, 1)
ocv_data = 3.0 .+ 0.8 .* soc_true .+ 0.02 .* randn(n)
y = ocv_data

# Fit the spline
soc_grid = 0:0.1:1
ocv_grid = 3.0 .+ 0.8 .* soc_grid .+ 0.01 .* randn(length(soc_grid))
spline = CubicSpline(soc_grid, ocv_grid, extrapl=[1,], extrapr=[1,])
spline_fn(x) = spline(x)
# Show indexing works
@show spline(0.5)

# Model definition
@model function SoC_stream(y, I, B, αw, βw, αv, βv, spline)
   λw ~ Gamma(shape = αw, rate = βw)
   λv ~ Gamma(shape = αv, rate = βv)
   SoC_prev ~ NormalMeanVariance(0.0, (0.01)^2)
   for t in 1:length(y)
       x[t] := SoC_prev + B*I[t]
       SoC_next[t] ~ NormalMeanPrecision(x[t], λw)
       z[t] := spline_fn(SoC_next[t])
       y[t] ~ NormalMeanPrecision(z[t], λv) # <---- Error happens here
       SoC_prev = SoC_next[t]
   end
end

# Set parameters
B = 0.01
αw, βw = 2.0, 2.0
αv, βv = 2.0, 2.0

# Attempt inference
init = @initialization begin
   μ(z) = NormalMeanVariance(0.0,1.0)
   μ(x) = NormalMeanVariance(0.0,1.0)
   # seed the plate’s marginals (same name as in @model)
   q(SoC_next) = NormalMeanVariance(0, (0.01)^2)
   # seed the two precision priors
   q(λw) = GammaShapeRate(αw, βw)
   q(λv) = GammaShapeRate(αv, βv)
end

constraints = @constraints begin
   # Group the entire SoC chain vs. the two precisions
   q(z,λv) = q(z)q(λv)
   q(x,λw,SoC_next) = q(x,SoC_next)q(λw)

end

meta = @meta begin
   spline_fn() -> Unscented()
end
model = SoC_stream(B = B, αw = αw, βw = βw, αv = αv, βv = βv, spline = spline)

result = infer(
   model = model,
   data = (y = y, I = I),
   returnvars     = ( SoC_next = KeepEach(),
                      λw       = KeepLast(),
                      λv       = KeepLast() ),
   initialization = init,
   constraints = constraints,
   iterations = 100,
   options        = (limit_stack_depth = 300,),
   meta = meta
)

# Review results
margs = vcat(result.posteriors[:SoC_next][end]...)  # flatten Vector{Marginal}
soc_mean = [ mean(m) for m in margs ]

plot(soc_true)
plot!(soc_mean)

[albertpod]

Hi @joshualeond!

Thanks for checking out rxinfer! A quick solution would be:

using RxInfer, CubicSplines, Distributions, Random, Plots

# Simulate some data
Random.seed!(42)
n = 100
I = randn(n)                    # Simulated current measurements
soc_true = cumsum(0.01 .+ 0.1 .* randn(n))
soc_true = clamp.(soc_true, 0, 1)
ocv_data = 3.0 .+ 0.8 .* soc_true .+ 0.02 .* randn(n)
y = ocv_data

# Fit the spline
soc_grid = 0:0.1:1
ocv_grid = 3.0 .+ 0.8 .* soc_grid .+ 0.01 .* randn(length(soc_grid))
spline = CubicSpline(soc_grid, ocv_grid, extrapl=[1,], extrapr=[1,])

my_spline(x) = spline(x)
@show my_spline(0.5)

# Model definition
@model function SoC_stream(y, I, B, αw, βw, αv, βv, spline)
   λw ~ Gamma(shape = αw, rate = βw)
   λv ~ Gamma(shape = αv, rate = βv)
   SoC_prev ~ NormalMeanVariance(0.0, (0.01)^2)
   for t in 1:length(y)
       SoC_next[t] ~ NormalMeanPrecision(SoC_prev + B*I[t], λw)
       tmp[t] := my_spline(SoC_next[t])
       y[t] ~ NormalMeanPrecision(tmp[t], λv) # <---- Error happens here
       SoC_prev = SoC_next[t]
   end
end

spline_meta = @meta begin 
    my_spline() -> Unscented()
end

# Set parameters
B = 0.01
αw, βw = 2.0, 2.0
αv, βv = 2.0, 2.0

# Attempt inference
init = @initialization begin
   # seed the plate’s marginals (same name as in @model)
   μ(SoC_next) = NormalMeanVariance(0, (0.01)^2)
   # seed the two precision priors
   q(λw) = GammaShapeRate(αw, βw)
   q(λv) = GammaShapeRate(αv, βv)
end

constraints = @constraints begin
   # Group the entire SoC chain vs. the two precisions
   q(SoC_prev, SoC_next, tmp, λw, λv) = q(SoC_prev, SoC_next)q(tmp)q(λw)q(λv)
end

model = SoC_stream(B = B, αw = αw, βw = βw, αv = αv, βv = βv, spline = spline)

result = infer(
   model = model,
   data = (y = y, I = I),
   returnvars     = ( SoC_next = KeepLast(),
                      tmp      = KeepLast(),
                      λw       = KeepLast(),
                      λv       = KeepLast() ),
   initialization = init,
   meta = spline_meta,
   constraints = constraints,
   iterations = 100,
   free_energy = true,
   options        = (limit_stack_depth = 300,)
)

plot(result.free_energy)

# Review results
margs = vcat(result.posteriors[:tmp]...)  # flatten Vector{Marginal}
soc_mean = [ mean(m) for m in margs ]
soc_std = [ std(m) for m in margs ]

scatter(y, label="observed SoC", xlabel="Time", ylabel="SoC", title="SoC Estimation")
plot!(3.0 .+ 0.8 .* soc_true, label="True SoC")
plot!(soc_mean, ribbon=soc_std, label="SoC mean", xlabel="Time", ylabel="SoC", title="SoC Estimation")

You need to do approximations around spline function.
I am sure there are more elegant solutions to what I proposed.

Please shoot any questions you might have.

UPD: Did few changes in the plotting and posterior extraction (changed to tmp)

[joshualeond]

Thanks so much for the help on this! RxInfer is awesome. Some questions if you don't mind:

1. What do you mean by "approximations around spline"? Is this what you did by adding a meta block?
2. You changed the initialization block to initialize the message for SoC_next rather than the marginal. Why is that?
3. Deterministic nodes (like tmp here) should be added to the constraints as well?
4. I haven't used iterations so far in my work with RxInfer. I'm assuming this has to do with variational inference and an optimization procedure. If the free energy plot seems to converge after 50 iterations or so then should I only pay attention to the last iteration in the result posteriors? Something more like: margs = vcat(result.posteriors[:SoC_next]...) |> x -> last(x, n) where n is the length of y in the script above?

[bvdmitri]

1. RxInfer’s model specification may resemble other probabilistic programming languages, where a statement like SoC_next[t] ~ ... might suggest you’re sampling a value (e.g., an Int or Float64) that you can use directly later on. However, in RxInfer, this is not the case. The expression SoC_next[t] ~ ... defines a random variable, represented internally as a RandomVariable object. This object doesn’t hold a concrete value but a distribution. If you need to pass such a variable to an arbitrary function, you’ll need an approximate value. This is where the meta block comes in—it allows you to specify approximation strategies suitable for your context. You can read more about this here.
2. The Unscented approximation method requires that messages be initialized before the iterative inference process can begin.
3. Isn’t strictly necessary in all cases, but in this context, it helps simplify and stabilize the inference procedure.
4. In the script, you’ll find a line like returnvars = (SoC_next = KeepLast(),). This is described in the documentation. In short, KeepLast instructs RxInfer to ignore intermediate variational results and return only the final (and typically best) approximation. This means you won’t need to do any additional post-processing on the result.posteriors no matter if you do 50 or 100 iterations you will always get the last result with this setting. This being said, its good to check the convergence via Free Energy and maybe reduce the amount of iterations because as soon as the algorithm converged there are only little changes in the result and mostly a waste of computational resources.