Support saving discrete parameter changes for non-symbolic ODEProblems using `SciMLStructures.Discrete`

[ivborissov]

Is your feature request related to a problem? Please describe.

I would like to save parameter changes made by callbacks in an ODEProblem defined directly, without ModelingToolkit.

The parameter object implements the SciMLStructures interface and exposes parameters modified by callbacks through the Discrete portion. Conceptually, the goal is to support a custom parameter container similar to MTKParameters:

struct MyParameters{T,D}
    tunable::T
    discrete::D
end

and have changes to p.discrete saved by SciMLBase.save_discretes!, so that the correct historical parameter values can be used when interpolating the solution after the simulation.

Defining the SciMLStructures.Discrete portion is not sufficient to enable discrete parameter saving for a plain ODEProblem.

The discrete parameter time-series storage in the solution is only created when the ODEFunction has a sys. Without one, sol.discretes is nothing, so save_discretes! has nowhere to store the values.

A simplified example is:

using OrdinaryDiffEqTsit5
import SciMLStructures as SS

struct MyParameters{T,D}
    tunable::T
    discrete::D
end

SS.isscimlstructure(::MyParameters) = true
SS.ismutablescimlstructure(::MyParameters) = true


SS.hasportion(::SS.Tunable, ::MyParameters) = true
SS.hasportion(::SS.Discrete, ::MyParameters) = true
SS.hasportion(::SS.AbstractPortion, ::MyParameters) = false

function SS.canonicalize(::SS.Tunable, p::MyParameters)
    arr = p.tunable
    repack = let p = p
        function (new_val)
            return SS.replace(SS.Tunable(), p, new_val)
        end
    end
    return arr, repack, true
end

SS.replace(::SS.Tunable, p::MyParameters, values) =
    MyParameters(values, p.discrete)

function SS.replace!(::SS.Tunable, p::MyParameters, values)
    copyto!(p.tunable, values)
    return nothing
end

function SS.canonicalize(::SS.Discrete, p::MyParameters)
    arr = p.discrete
    repack = let p = p
        function (new_val)
            return SS.replace(SS.Discrete(), p, new_val)
        end
    end
    return arr, repack, true
end

SS.replace(::SS.Discrete, p::MyParameters, values) =
    MyParameters(p.tunable, values)

function SS.replace!(::SS.Discrete, p::MyParameters, values)
    copyto!(p.discrete, values)
    return nothing
end

function f!(du, u, p, t)
    du[1] = -p.tunable[1] * p.discrete[1] * u[1]
end

condition(u, t, integrator) = t - 0.5

function affect!(integrator)
    integrator.p.discrete[1] = 2.0
end

p = MyParameters([1.0], [1.0])
callback = ContinuousCallback(condition, affect!)

prob = ODEProblem(f!, [1.0], (0.0, 1.0), p)
sol = solve(prob, Tsit5(); callback)

@show sol.discretes  # nothing

Describe the solution you’d like
It would be useful to clarify the intended public interface for this use case.

Is providing a custom sys the intended way to enable discrete parameter time-series storage for a non-symbolic ODEProblem?

Alternatively, could SciMLBase provide a default implementation based on the SciMLStructures.Discrete portion of p when:

SciMLStructures.hasportion(
    SciMLStructures.Discrete(),
    p,
)

is true?

Describe alternatives you’ve considered

Current workaround implies introducing some dummy sys and adding methods to SciMLBase.create_parameter_timeseries_collection and SciMLBase.get_saveable_values

import SymbolicIndexingInterface as SII

struct MySystem end

Base.copy(p::MyParameters) =
    MyParameters(copy(p.tunable), copy(p.discrete))

function SciMLBase.create_parameter_timeseries_collection(
    ::MySystem,
    p::MyParameters,
    tspan,
)
    buffer = SciMLBase.DiffEqArray(
        typeof(copy(p.discrete))[],
        eltype(tspan)[],
        (1,),
    )

    return SII.ParameterTimeseriesCollection(
        (buffer,),
        copy(p),
    )
end

function SciMLBase.get_saveable_values(
    ::MySystem,
    p::MyParameters,
    timeseries_idx,
)

    return copy(p.discrete)

end

ode_function = ODEFunction(f!; sys = MySystem())
prob2 = ODEProblem(
    ode_function,
    [1.0],
    (0.0, 1.0),
    p,
)
sol2 = solve(prob2, Tsit5(); callback)

sol2.discretes[1].t # 0.5
sol2.discretes[1].u # [2.0]

This works, but it requires introducing a synthetic system object and extending APIs that appear primarily intended for symbolic systems and internal usage.