Performant way to remake ODEProblem with updated external input function?

[hersle]

I have an ODE for x(t). It depends on a parameter P not only explicitly, but also implicitly through an external input (forcing) function F(t). I want to efficiently solve the ODE repeatedly for many values of P, always updating F(t) with the proper dependence on P. Eventually, I also want to autodifferentiate the solution with respect to P, and that this accurately propagates through F(t).

Here is the closest I have come, trying to use remake on a simple example:

using ModelingToolkit
using ModelingToolkit: t_nounits as t, D_nounits as D
using DifferentialEquations
using DataInterpolations
using ForwardDiff

P0 = 2.0 # test value for the parameter P

# F(t) is an external function that depends on the parameter P
function create_F(P)
    ts = range(0.0, 1.0, step=0.1)
    Fs = ts .^ P # toy example: F = t^P
    return QuadraticSpline(ts, Fs)
end
F_spline = create_F(NaN) # uninitialized spline (replace NaN -> P0 for code to run)
F(t) = F_spline(t)

# an ODE that depends on the parameter P both explicitly, and implicitly through F(t)
@parameters P
@variables x(t)
@register_symbolic F(t) # following https://docs.sciml.ai/ModelingToolkit/stable/tutorials/ode_modeling/#Specifying-a-time-variable-forcing-function
sys = structural_simplify(ODESystem([D(x) ~ P * F(t)], t, [x], [P]; name = :sys))
prob = ODEProblem(sys, [x => 0.0], (0.0, 1.0), [P => NaN]) # uninitialized problem

# solution of the ODE for a given value of the parameter P (must be fast!)
function solve_instance_fast(P)
    F_instance = create_F(P)
    prob_instance = remake(prob; p = [P]) # QUESTION: how to update F in prob to F_instance?
    return solve(prob_instance)
end

ForwardDiff.derivative(P -> solve_instance_fast(P)(1.0; idxs=x), P0)

Is there a way to accomplish this with ModelingToolkit?

[ChrisRackauckas]

Why not just make the spline the parameter and update that?

[hersle]

That works! Here is a working example:

using ModelingToolkit
using ModelingToolkit: t_nounits as t, D_nounits as D
using DifferentialEquations
using DataInterpolations
using ForwardDiff

P0 = 2.0 # test value for the parameter P

# F(t) is an external function that depends on the parameter P
function create_F(P)
    ts = range(0.0, 1.0, step=0.1)
    Fs = ts .^ P # toy example: F = t^P
    return CubicSpline(Fs, ts)
end
F(t, spline) = spline(t) # intermediate function for F(t) (possible to avoid this?)

# an ODE that depends on the parameter P both explicitly, and implicitly through F(t)
@parameters P F_spline
@variables x(t)
@register_symbolic F(t, spline) # following https://docs.sciml.ai/ModelingToolkit/stable/tutorials/ode_modeling/#Specifying-a-time-variable-forcing-function
sys = structural_simplify(ODESystem([D(x) ~ P * F(t, F_spline)], t, [x], [P, F_spline]; name = :sys))
prob = ODEProblem(sys, [x => 0.0], (0.0, 1.0), [P => NaN, F_spline => NaN]) # uninitialized problem

# solution of the ODE for a given value of the parameter P (must be fast!)
function solve_instance_fast(P)
    F_instance = create_F(P)
    prob_instance = remake(prob; p = [sys.P => P, sys.F_spline => F_instance]) # update P and spline for F(t)
    return solve(prob_instance)
end

x_at_1(P) = solve_instance_fast(P)(1.0; idxs=x) # analytical result: x(1) = P/(P+1)
ForwardDiff.derivative(x_at_1, P0) # analytical result: d(x(t=1))/dP = 1/(P+1)^2

Thank you! I thought I was perhaps dealing with a restriction of MTK, and this simple solution just did not cross my mind.