Type-stable solving with non-numeric parameters?

[hersle]

I am solving ODEs that depend on non-numeric parameters. The concrete parameter type differs between problems (e.g. with/without AD), so I can only type it abstractly in the system definition (e.g. @parameters P::SomeAbstractType).

With all combinations of SII.set*_oop and remake I have tried, this makes the generated ODE function type-unstable in solve. This kills my performance.

Are there any tools I should look into to work around this? Or is this an insurmountable limitation for now?

Let me know if an example helps. Thanks.

[ChrisRackauckas]

Is the array of things to set all matching in type? Or using a tuple?

[hersle]

Is the array of things to set all matching in type?

It is matching in the sense that all the updated concrete parameter types are subtypes of the declared abstract system parameter types. Here is a minimal example of what I'm trying to do, before any remake stuff (MTK v9.80.5):

using ModelingToolkit, OrdinaryDiffEq, DataInterpolations, SymbolicIndexingInterface, ProfileCanvas
using ModelingToolkit: t_nounits as t, D_nounits as D

@register_symbolic DataInterpolations._interpolate(s::LinearInterpolation, t)
@register_symbolic DataInterpolations.derivative(s::LinearInterpolation, t)

@variables x(t)
@parameters I::LinearInterpolation # abstract; don't know exact type user will pass
@mtkbuild sys = ODESystem([D(x) ~ DataInterpolations._interpolate(I, t) + DataInterpolations.derivative(I, t)], t)

Idummy = LinearInterpolation([0.0, 0.0], [0.0, 1.0])
prob = ODEProblem(sys, [sys.x => 0.0], (0.0, 1.0), [sys.I => Idummy])
println(typeof(prob.p)) # shows MTKParameters{..., Tuple{Vector{LinearInterpolation}}, ...}; why not typeof(Idummy)?
@profview sol = solve(prob, Tsit5(); adaptive = false, dt = 1e-5, save_everystep = false) # type-unstable

The @profview shows that solve is type-unstable.

Or using a tuple?

What exactly do you mean to make a tuple? Make the parameter Tuple{LinearInterpolation} instead of LinearInterpolation?

[hersle]

By accident I discovered this. If I remake the problem in this particular way, solve is type-stable:

newu0, newp = setsym_oop(prob, [sys.I])(prob, Any[Idummy])
newprob = remake(prob; u0 = newu0, p = newp, build_initializeprob = false)
@profview newsol = solve(newprob, Tsit5(); adaptive = false, dt = 1e-5, save_everystep = false)

But it becomes type-unstable again if I do any of these three things: remove the Any, or setsym with Idummy not inside [], or remove build_initializeprob = false.

[SebastianM-C]

Btw, you might be interested in the Interpolation / ParametrizedInterpolation blocks from MTKStdLib (full disclosure, I wrote them 😅 )

[hersle]

Thanks for the tip! Those look very relevant for me. Noted for a future refactor :)

[hersle]

The bad workaround I was using here no longer works. I think maybe #3957 changed how remake handles parameter types.

[hersle]

@AayushSabharwal Is there a nice way to accomplish what I want here? Here is the shortest example I can think of:

using ModelingToolkit, OrdinaryDiffEq
using ModelingToolkit: t_nounits as t, D_nounits as D

f(P) = P[1]
@register_symbolic f(P::Tuple)

@variables x(t)
@parameters P::Tuple
#@parameters P::Tuple{Int} # uncomment and everything is type-stable and fast
@mtkbuild M = System([D(x) ~ f(P)], t)

prob = ODEProblem(M, [x => 0.0, P => (1,)], (0.0, 1.0))
@profview sol = solve(prob, Tsit5(); adaptive = false, dt = 1e-5, dense = false, save_everystep = false)

I want the solve to be type-stable and fast. But it is not, because it cannot infer the concrete type of P:

The reason is of course that I only declare P::Tuple with an abstract type. But I don't know its concrete type until after creating the problem. Is there a way to remake the problem so that it specializes on the concrete type of P, e.g. something like this?

newprob = remake(prob; p = [P => (2,)], specialize = true) # for example
solve(newprob) # I want this to be type-stable and fast

Here P was a Tuple{Int}, but I wish for it to be a non-numeric parameter of any type.

[ChrisRackauckas]

This is what the respecialize function function does

[hersle]

Thanks, this is very very close to what I want! The problem I see is that it operates on the System, and it would be very expensive in my application to transform the System and recreate the ODEProblem from scratch. Is it possible to do it directly on the ODEProblem or MTKParameters object?

[AayushSabharwal]

There isn't. The problem is that if you change types, then the indices of the parameters in MTKParameters changes. The system needs to be updated accordingly, otherwise you won't be able to symbolically index the ODEProblem, and the generated code will be incorrect.

[hersle]

I see :/ Do you mean that the indexing of prob.p.nonnumeric[i][j] in general changes, since it groups nonnumeric parameters of equal types?