Type instability with parameters that are abstract at system-level, but concrete at problem-level

[hersle]

Is this what is meant by this documentation? I thought to open an issue to track this and have a MWE.

I think it is common to do simulations with parameters whose types can only be known abstractly at system-definition time, but that will be known concretely at problem-construction and solution time. Ideally, I would want the final problem to be type-stable and performant. But it seems to me that such setups currently suffer a performance hit due to type instabilities.

For example, consider a problem that calculates the work done on a particle following any data-given trajectory:

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

splvalue(s, t) = s(t)
splderiv(s, t, order) = DataInterpolations.derivative(s, t, order)
@register_symbolic splvalue(s::CubicSpline, t)
@register_symbolic splderiv(s::CubicSpline, t, order)

function create_system(; Tspl = CubicSpline, kwargs...)
    # system that calculates work done on a particle following any data-given trajectory
    @parameters m yspl::Tspl
    @variables y(t) v(t) a(t) F(t) W(t)
    return structural_simplify(ODESystem([
        y ~ splvalue(yspl, t) # position (height)
        v ~ splderiv(yspl, t, 1) # velocity
        a ~ splderiv(yspl, t, 2) # acceleration
        F ~ m * a # force
        D(W) ~ F * v # power = time derivative of work
    ], t; defaults = [W => 0.0, m => 1.0], kwargs...))
end

function compute_work(v0; g = 1.0, concrete = false)    
    y(t) = v0*t - 1/2*g*t^2 # anal(ytical) traj of part in const grav field
    ts = range(0.0, v0/g, step=0.01) # simulate from ground to highest point
    yspline = CubicSpline(y.(ts), ts)

    # HACK: inform system about abstract/concrete type
    @named sys = create_system(; Tspl = concrete ? typeof(yspline) : CubicSpline)
    prob = ODEProblem(sys, [], extrema(ts), [sys.yspl => yspline])
    sol = solve(prob, Rodas5P(); reltol = 1e-10) # chosen to amplify performance differences
    return only(sol[end]) # get work done (the only unknown) at final time
end

This has a workaround to inform create_system() of which (abstract/concrete) type to use for the spline. Benchmarks report that the concrete version is faster:

@btime compute_work(1e4; concrete = true) # 2.705 s (7243956 allocations: 695.90 MiB)
@btime compute_work(1e4; concrete = false) # 4.226 s (57211102 allocations: 1.48 GiB)

This seems to be because of type instabilities (red) in the abstract version:

@profview compute_work(1e4; concrete = false)

In contrast, the concrete version looks type stable (blue):

@profview compute_work(1e4; concrete = true)

The performance hit becomes more severe for larger ODEs. I would like to get rid of the workaround, in order to make the problem naturally compatible and performant with automatic differentiation (which enters the spline type), etc. Is this something that could be fixed by MTK in the future, or should one find workarounds?

[ChrisRackauckas]

The performance hit becomes more severe for larger ODEs. I would like to get rid of the workaround, in order to make the problem naturally compatible and performant with automatic differentiation (which enters the spline type), etc. Is this something that could be fixed by MTK in the future, or should one find workarounds?

There are things that we're doing with symbolic performance etc., but that should be orthogonal to this use case. When doing AD of a model, you can simplify remake the prob. You only ever have to symbolically regenerate the model if you are making structural changes to it, which shouldn't happen in an autodiff loop. So this:

    @named sys = create_system(; Tspl = concrete ? typeof(yspline) : CubicSpline)
    prob = ODEProblem(sys, [], extrema(ts), [sys.yspl => yspline])

shouldn't be in there at all. Instead, SII tooling like setsym_oop can be used to take an existing problem object and update the values. That's the kind of stuff that https://docs.sciml.ai/ModelingToolkit/stable/examples/remake/ demonstrates, though admittedly we need to update it as there are some new primitives that can be quite helpful.

[hersle]

Thanks, Chris, this is helpful! I think the docs explain well that one should not regenerate the model in hot loops, so I am aware of that. My example is overly contrived in order to have a simple switch that triggers the type stable/unstable behavior.

I was not aware of setsym_oop, but see it now in the SymbolicIndexingInterface docs. I am sure the MTK docs can be improved to help the discoverability of that and show example usage. Do I get it right that (in addition to remake()/replace()) one should use setp() for efficiently changing a parameter without changing its type, and setp_oop() if the new type is different?

Still, I am surprised by the behavior of the code above. At the end of the day, compute_work(1e4; concrete = false) calls ODEProblem() with a concrete value for the yspl parameter, so doesn't the code generation have all the information it needs to generate type-stable code for the solve(), at least in theory? Or must one both make and remake a problem to get it type stable?

My understanding is limited, though, and I know there are many practical details going on under the hood that make things much easier said than done. Thanks!

[hersle]

I just tried to reupdate the spline parameter with setsym_oop() and remake() the problem:

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

splvalue(s, t) = s(t)
splderiv(s, t, order) = DataInterpolations.derivative(s, t, order)
@register_symbolic splvalue(s::CubicSpline, t)
@register_symbolic splderiv(s::CubicSpline, t, order)

function create_system(; Tspl = CubicSpline, kwargs...)
    # system that calculates work done on a particle following any data-given trajectory
    @parameters m yspl::Tspl
    @variables y(t) v(t) a(t) F(t) W(t)
    return structural_simplify(ODESystem([
        y ~ splvalue(yspl, t) # position (height)
        v ~ splderiv(yspl, t, 1) # velocity
        a ~ splderiv(yspl, t, 2) # acceleration
        F ~ m * a # force
        D(W) ~ F * v # power = time derivative of work
    ], t; defaults = [W => 0.0, m => 1.0], kwargs...))
end

function compute_work(v0; g = 1.0)    
    y(t) = v0*t - 1/2*g*t^2 # anal(ytical) traj of part in const grav field
    ts = range(0.0, v0/g, step=0.01) # simulate from ground to highest point
    yspline = CubicSpline(y.(ts), ts)

    @named sys = create_system()
    prob = ODEProblem(sys, [], extrema(ts), [sys.yspl => yspline])

    # set spline parameter with setsym_oop and remake problem to make solve() type-stable
    setyspl = ModelingToolkit.setsym_oop(prob, sys.yspl)
    newu0, newp = setyspl(prob, yspline)
    prob = remake(prob, u0 = newu0, p = newp)

    sol = solve(prob, Rodas5P(); reltol = 1e-10) # chosen to amplify performance differences
    return only(sol[end]) # get work done (the only unknown) at final time
end

The type-stable solve() behavior (concrete = true above) is then restored:

@btime compute_work(1e4) # 1.739 s (7244194 allocations: 695.90 MiB)
@profview compute_work(1e4)

Thanks!