Indentify parameters in an ODEProblem

Specific Domains Modelling & Simulations
, ,
1

Hello,

I have specified a model in Turing.jl as follows:

using Turing


@model function LGDS(x::AbstractArray, prob, params::Dict, t_int)

    # prior distributions
    σ ~ InverseGamma(2, 3)
    lC ~ MvNormal(params[:μ_C], params[:Σ_C])
    lR ~ MvNormal(params[:μ_R], params[:Σ_R])
    lA ~ MvNormal(params[:μ_A], params[:Σ_A])

    p = [lC; lR; lA]
   
    # remake the problem with the sampled parameter values
    prob_samp = remake(prob; p = p)

    # solve the ODE
    T_ode = solve(prob_samp, Tsit5(); save_idxs=1, verbose=false)
    if !(SciMLBase.successful_retcode(T_ode.retcode))
       return
    end

    # time points corresponding to observations `x`
    x_est = T_ode(time)  
    x ~ MvNormal(x_est.u, σ^2 * I)    

    return nothing
end

# instantiate model
input_data = data[!, :T_i]
model = LGDS(input_data, prob, prior_params, Δt)

# sample independent chains
n_chains = 3
chain = sample(model, NUTS(0.65), MCMCSerial(), 30, n_chains; progress=true)

The output of Turing.jl looks like this:

point_est = mean(chain[:, :, 3])

Mean
  parameters      mean 
      Symbol   Float64 

           σ   11.8814
       lC[1]    2.1847
       lC[2]    1.0088
       lC[3]   -0.0872
       lR[1]   -1.4498
       lR[2]   -0.4706
       lR[3]    1.6218
       lA[1]    0.6406
       lA[2]    0.9099

The problem is now, how can I uniquely identify the parameters in the final p at the end of the estimation process?

It’s essentially similar to this problem.

I tried this:

prob_final = remake(prob; p=p)
prob_final.ps[sys.lC_i]

But that will give:

ArgumentError: invalid index: ModelingToolkit.ParameterIndex …

2

Use the getu/setu style

3

Do you mean setp/getp?

I tried using setp in the Turing.jl model but it doesn’t seem to propagate the dual numbers properly:

using Turing
using DifferentialEquations
using SymbolicIndexingInterface

# create setter
setter! = setp(sys, [
    sys.lC_i, 
    sys.lC_e, 
    sys.lC_h, 
    sys.lR_ie, 
    sys.lR_ea, 
    sys.lR_ih, 
    sys.lA_w, 
    sys.lA_e]
)

@model function LGDS(x::AbstractVector{T}, prob, params::Dict, t_int) where T

    # prior distributions
    σ ~ InverseGamma(2, 3)
    lC ~ MvNormal(params[:μ_C], params[:Σ_C])
    lR ~ MvNormal(params[:μ_R], params[:Σ_R])
    lA ~ MvNormal(params[:μ_A], params[:Σ_A])

    # params vector
    p = [lC; lR; lA]
   
    # set the sampled parameter values inside the ODE problem
    setter!(prob, p)

    # solve the ODE
    T_ode = solve(prob, Tsit5(); save_idxs=1, verbose=false)
    if !(SciMLBase.successful_retcode(T_ode.retcode))
       return
    end

    # time points corresponding to observations `x`
    x_est = T_ode(time)  
    x ~ MvNormal(x_est.u, σ^2 * I)    

    return nothing
end

# instantiate model
input_data = data[!, :T_i]
model = LGDS(input_data, prob, prior_params, Δt)

# sample independent chains
n_chains = 1
chain = sample(model, NUTS(0.4), MCMCSerial(), 10, n_chains; progress=true)
4

In the meantime, this works fine:

prob_samp = remake(prob; p = [sys.lC_i => lC[1], 
                            sys.lC_e => lC[2],
                            sys.lC_h => lC[3],
                            sys.lR_ie => lR[1],
                            sys.lR_ea => lR[2],
                            sys.lR_ih => lR[3],
                            sys.lA_w => lA[1],
                            sys.lA_e => lA[2]])

I can now at least identify the parameters from the Turing.jl chain again. But this is probably not the most optimal way to do it.

5

You need a out of place setter so it can build a new container for the type change. @cryptic.ax I thought that was in the interface, but I don’t see it on the doc page?

6

setp is in-place. remake is out-of-place. Assuming prob is created through ModelingToolkit.jl, there is no parameter order so p = [lC; lR; lA] doesn’t work. It ends up replacing the MTKParameters object with a Vector, which it of course can’t index the same was as MTKParameters. Providing a symbolic map to remake is the correct way to do this. The resulting problem can also be indexed symbolically. If you need to index the same variable from the same system in a loop, I recommend caching the function returned by getp and reusing it in the loop. This will improve performance and avoid having to AD through the construction of the function (if that is something you will end up doing).

7

Alternatively, you can do something like

# `copy` also works for `MTKParameters`, but `deepcopy` works for everything
newp = deepcopy(parameter_values(prob))
setp(sys, [syms...])(newp) # ideally cache the result of `setp(...)`
remake(prob, p = newp)

If the types of the variables don’t change.

8

There needs to be an out of place step in the interface here for first generating the type changed version though. I don’t think it’s a good answer to say it needs to use remake to do that. Why not have an out of place setp for this?

9

remake_buffer does that. It doesn’t return a function though, just does the update similar to remake.

10

Then we should probably expose it under a different name. For simplicity, it really just wants (prob, [syms...], newp) and return the prob of with the extended types, where newp is a vector matching the ordering of syms.

11

remake_buffer currently expects (indp::IndexProvider, valp::ValueProvider, newvals::Dict) where the newvals maps symbols to their new values. It returns a copy of valp with updated values and type. However, this doesn’t operate at the prob level, it operates at the MTKParameters level (valp::MTKParameters). Operating at the prob level is what remake does.

I do like your API better, though. We can make remake_buffer just wrap the new function for now.

12

It is possible to make remake_buffer (or whatever we call the new function) to operate on the problem. We’d need additional API which wraps @set! prob.p = newp (similarly for prob.u0) since SII doesn’t know which field the parameters are stored in. It would basically be equivalent to

newstates = remake_buffer(state_values(prob), syms, newvals)
newparams = remake_buffer(parameter_values(prob), syms, newvals)
# Both of the subsequent lines are actually function calls which wrap
# this behavior
@set! prob.p = newparams
@set! prob.u0 = newu0
13

Thanks, I’m trying to make this work.

How do I actually get the new parameters in here? Because you just seem to copy the old ones and then set those again in remake?

I am also a bit confused why the API requires me to specify syms. If I give a symbolic map to setp, then it should already be clear what I want to change.

remake_buffer doesn’t seem to have any docs, so I have no idea how to use that.

Also, what is IndexProvider? I can’t find any reference to that in the docs either.

14

remake_buffer doesn’t seem to have any docs, so I have no idea how to use that.

What version of SymbolicIndexingInterface.jl are you using? It’s documented here.

Also, what is IndexProvider? I can’t find any reference to that in the docs either.

It’s not a type. It’s terminology used by SII to be independent of the notion of systems/problems/integrators/solutions. These terms are defined in the Terminology section of the documentation.

I am also a bit confused why the API requires me to specify syms. If I give a symbolic map to setp, then it should already be clear what I want to change.

setp doesn’t change any values. setp(indp::IndexProvider, syms) returns a function which can perform an update to the parameters in syms for any object which uses indp as its index provider. For example:

# This is an index provider. It translates symbolic variables :x, :a, ...
# to numeric indices
indp = SymbolCache([:x, :y, :z], [:a, :b, :c], :t)
# `setter` is a function which takes
# `(valp::ValueProvider, [value_of_c, value_of_a])`
# and updates the value of `:c` in `valp` to `value_of_c`.
setter = setp(indp, [:c, :a])
# `ProblemState` is a simple value provider in SII
# It could have been named better
state = ProblemState(; u = [1.0, 2.0, 3.0], p = [4.0, 5.0, 6.0], t = 0.0)
# gets :c and :a from a value provider
getter = getp(indp, [:c, :a])
println(getter(state)) # prints [6.0, 4.0]
setter(state, [60.0, 40.0]) # update values
println(getter(state)) # prints [60.0, 40.0]

Note that the function returned by setp does an in-place update. If the types of variables have changed, you must use remake_buffer.

A more comprehensive example is provided here and here.

15

I was looking for an accompanying text/example in ModelingToolkit.jl. From that page I couldn’t figure how to apply it to my problem, it’s too abstract. That might be a shortcoming of the user but I’m just giving it as feedback because I am really struggling with this system if you couldn’t tell already.

Thanks, that clears up some stuff.

I missed that page. Since setp doesn’t work for my problem, I used remake_buffer like this:

# ordered symbols
syms = [sys.lC_i, sys.lC_e, sys.lC_h, sys.lR_ie, sys.lR_ea, sys.lR_ih, sys.lA_w, sys.lA_e]

@model function LGDS(x::AbstractArray, prob, params::Dict, t_int)

    # prior distributions
    σ ~ InverseGamma(2, 3)
    lC ~ MvNormal(params[:μ_C], params[:Σ_C])
    lR ~ MvNormal(params[:μ_R], params[:Σ_R])
    lA ~ MvNormal(params[:μ_A], params[:Σ_A])
   
    # remake the problem with the sampled parameter values and enforce order
    new_map = Dict(zip(syms, [lC; lR; lA]))
    newp = remake_buffer(sys, prob.p, new_map)
    new_prob = remake(prob, p=newp)

    # solve the ODE
    T_ode = solve(new_prob , Tsit5(); save_idxs=1, verbose=false)
    if !(SciMLBase.successful_retcode(T_ode.retcode))
       return
    end

    # time points corresponding to observations `x`
    x_est = T_ode(time)  
    x ~ MvNormal(x_est.u, σ^2 * I)    

    return nothing
end

Could you please confirm if this is the right way to do it? It’s sampling correctly now as far as I can tell.

Lastly, you mentioned caching the result of remake_buffer. Any recommendations on how best to do that? Turing.jl recommends Memoization.jl: Reuse Computations in Gibbs Sampling

16

I was looking for an accompanying text/example in ModelingToolkit.jl. From that page I couldn’t figure how to apply it to my problem, it’s too abstract. That might be a shortcoming of the user but I’m just giving it as feedback because I am really struggling with this system if you couldn’t tell already.

Thanks for the feedback, we’re definitely open to improving documentation wherever possible. This is probably the page you were looking for, since it describes a similar kind of process to what you are doing here. I’ll see where I can put it so it’s more visible instead of being hidden under a folding menu.

Could you please confirm if this is the right way to do it? It’s sampling correctly now as far as I can tell.

That looks correct

Lastly, you mentioned caching the result of remake_buffer.

The result of getp/setp can (and should) be cached since they return a function which doesn’t need to be repeatedly generated. You could actually cache the result of remake_buffer as well, since the returned parameter object has the correct type for all subsequent uses. PreallocationTools.jl might be of help here. It’s designed for AD, but the general idea applies here. Specifically this is probably what you want to use. The function you pass to GeneralLazyBufferCache takes the parameter values (as a vector) and does remake_buffer.

If you need help with PreallocationTools, please feel free to ask.

17

Indeed, I was following that page already. But I got a totally lost with all the different options available to modify ODE problems.

Apologies, I modified the code while you were typing a reply, i hope you caught that?

I implemented it like this:

# ordered symbols
syms = [sys.lC_i, sys.lC_e, sys.lC_h, sys.lR_ie, sys.lR_ea, sys.lR_ih, sys.lA_w, sys.lA_e]


@model function LGDS(x::AbstractArray, prob, params::Dict, t_int)

    # cache for remake(prob, p=newp)
    lbc = GeneralLazyBufferCache(function (p)
        remake_buffer(sys, prob.p, Dict(zip(syms, p)))    
    end)    

    # prior distributions
    σ ~ InverseGamma(2, 3)
    lC ~ MvNormal(params[:μ_C], params[:Σ_C])
    lR ~ MvNormal(params[:μ_R], params[:Σ_R])
    lA ~ MvNormal(params[:μ_A], params[:Σ_A])
   
    # remake the problem with the sampled parameter values and enforce order
    newp = lbc[[lC; lR; lA]]
    new_prob = remake(prob, p=newp)

    # solve the ODE
    T_ode = solve(new_prob, Tsit5(); save_idxs=1, verbose=false)
    if !(SciMLBase.successful_retcode(T_ode.retcode))
       return
    end

    # time points corresponding to observations `x`
    x_est = T_ode(time)  
    x ~ MvNormal(x_est.u, σ^2 * I)    

    return nothing
end

Is that appropriate?

I am bit skeptical of how much performance benefits it will actually yield, since each sample p would have to be cached and I don’t know enough about HMC or the NUTS sampler implementation to say whether they would be reused at all. Is there a way to check cache usage?

18

You should create lbc outside the function, and use it inside (pass it as an argument, probably). Also, do the same for setter = setp(sys, syms) After that,

vals = [lC; lR; lA]
newp = lbc[vals]
setter(newp, vals)
new_prob = remake(prob, p = newp)

This works because GeneralLazyBufferCache reuses the value returned from the function if the type of p is the same. See the docs here. setter is also cached so that repeatedly translating syms to indices is not necessary.