Unsafe parameter mutation breaks latest release

[hersle]

On the latest 9.61.0 release:

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

@variables x(t)
@parameters p
@named sys = ODESystem([D(x) ~ 0], t, [x], [p]; initialization_eqs = [x ~ p]) # solution is constant x(t) ~ p
sys = structural_simplify(sys)

prob = ODEProblem(sys, [], (0.0, 1.0), [p => 1.0])
sol1 = solve(prob, Tsit5())
sol1[x+p] # gives 2.0, as it should

ModelingToolkit.setp(prob, p)(prob, 2.0)
sol2 = solve(prob, Tsit5())

sol1[x+p] # gives 3.0, but should give 2.0
sol2[x+p] # gives 4.0, as it should

This is bad. It looks like the setp effectively changes the value of p when reading sol1, which it should not.

I have narrowed down that this code ran as expected before commit 39f8271.

[ChrisRackauckas]

This is likely a result of the eager application of the initialization. This is not really a bug but a feature collision: if we want remake to apply u0 and parameter changes eagerly (a long time ask of @isaacsas for Catalyst, removing the need for rebuild with the latest versions), then you cannot validate changes afterwards (because otherwise they would be trivialized and undone). Basically the current setup is now during remake:

1. If the initialization has a trivial solution, update u0 and p immediately. Don't run the initialization check process during solve (as it's already resolved).
2. If it has a non-trivial solution, do not update the values (other than the guesses) and resolve at solve time.

(2) is of course required to resolve late because you require a solver to do this. The previous behavior of (1) was to also resolve at solve time, but we got some requests to not do that. I think then if this special case is too weird, then we should figure out some other way to resolve this that is both correct and easy to understand, since a lot of the previous proposals on rebuild were simply not correct so I wanted to avoid those.

[hersle]

I see, thank you for explaining. So the eager initialization effectively copied the parameters before 39f8271?

How would one go about now, then, to update the value of a single parameter such that it always remains "safe" to access in the solution, in a way that is as efficient as absolutely possible (like setp offered before)? Following the FAQ, duplicating the problem with something like this?

using ModelingToolkit, OrdinaryDiffEq, Setfield
using ModelingToolkit: t_nounits as t, D_nounits as D
using SciMLStructures, SymbolicIndexingInterface

@variables x(t)
@parameters P
@named sys = ODESystem([D(x) ~ 0], t, [x, y], [P]; initialization_eqs = [x ~ P]) # solution is constant x(t) ~ p
sys = structural_simplify(sys)

prob1 = ODEProblem(sys, [], (0.0, 1.0), [P => 1.0])
sol1 = solve(prob1, Tsit5())
sol1[x+P] # gives 2.0, as it should

p = copy(prob1.p) # copy, so mutation does not affect sol1 parameters
Pidx = parameter_index(prob1, P)
setindex!(p, 2.0, Pidx) # set P => 2.0
prob2 = @set prob1.p = p # create another problem with the copied parameter object
sol2 = solve(prob2, Tsit5())
sol2[x+P] # gives 4.0, as it should

sol1[x+P] # gives 2.0, as it should
sol2[x+P] # gives 4.0, as it should

This works, but is needlessly complicated. If the current setp is mutating, should MTK also offer an equivalent non-mutating (i.e. duplicating) single-parameter-updating method that does something like this?

[SebastianM-C]

I think that using setsym_oop and remake would be a better way to do an out of place mutation for the parameters.

[hersle]

Thank you. That does also work, but it is slower than my (hacky?) solution:

using ModelingToolkit, OrdinaryDiffEq, Setfield, BenchmarkTools
using ModelingToolkit: t_nounits as t, D_nounits as D
using SciMLStructures, SymbolicIndexingInterface

@variables x(t)
@parameters P
@named sys = ODESystem([D(x) ~ 0], t, [x], [P]; initialization_eqs = [x ~ P]) # solution is constant x(t) ~ p
sys = structural_simplify(sys)
prob = ODEProblem(sys, [], (0.0, 1.0), [P => 1.0])

Pidx = parameter_index(prob, P)
function remake1(prob, newP)
    p = copy(prob.p) # copy, so mutation does not affect sol1 parameters
    setindex!(p, newP, Pidx) # set P => 2.0
    return @set prob.p = p # create another problem with the copied parameter object
end

setter = setsym_oop(prob, P)
function remake2(prob, newP)
    u0, p = setter(prob, newP)
    return remake(prob; u0, p)
end

@btime prob1 = remake1(prob, 2.0) # 61.641 ns (5 allocations: 224 bytes)
@btime prob2 = remake2(prob, 2.0) # 25.423 μs (173 allocations: 6.17 KiB)

@test solve(prob1, Tsit5())[x+P][1] == 4.0 # passes
@test solve(prob2, Tsit5())[x+P][1] == 4.0 # passes

[ChrisRackauckas]

Your solution isn't AD compatible though for all ADs. If you don't need AD compatibility, then copy into setsym should be fine?

[hersle]

I guess you are right, and I would like to keep AD compatibility.

In my scenario, one CosmologySolution is made up of ~1000 ODESolutions, all of which are solutions of exactly the same ODEProblem, only on a "grid" with ~1000 different values of the parameter P. All P[i] are always of the same type. I want to avoid an "expensive" remake() for every P[i] in the hot loop. Instead, I want to do a single remake() before the loop (say for P[1]) to promote types and preserve AD compatibility, and then simply update the value to P[i] as cheaply as absolutely possible inside the hot loop. Is there a better way to do that?