fix: improve type promotion in remake_initializeprob

Author: AayushSabharwal

src/systems/nonlinear/initializesystem.jl

@@ -173,15 +173,54 @@ function is_parameter_solvable(p, pmap, defs, guesses)
     # either (missing is a default or was passed to the ODEProblem) or (nothing was passed to
     # the ODEProblem and it has a default and a guess)
     return ((_val1 === missing || _val2 === missing) ||
-           (_val1 === nothing && _val2 !== nothing)) && _val3 !== nothing
+            (_val1 === nothing && _val2 !== nothing)) && _val3 !== nothing
 end
 
 function SciMLBase.remake_initializeprob(sys::ODESystem, odefn, u0, t0, p)
-    if (u0 === missing || !(eltype(u0) <: Pair) || isempty(u0)) &&
-       (p === missing || !(eltype(p) <: Pair) || isempty(p))
+    if u0 === missing && p === missing
         return odefn.initializeprob, odefn.update_initializeprob!, odefn.initializeprobmap,
         odefn.initializeprobpmap
     end
+    if !(eltype(u0) <: Pair) && !(eltype(p) <: Pair)
+        oldinitprob = odefn.initializeprob
+        if oldinitprob === nothing || !SciMLBase.has_sys(oldinitprob.f) ||
+           !(oldinitprob.f.sys isa NonlinearSystem)
+            return oldinitprob, odefn.update_initializeprob!, odefn.initializeprobmap,
+            odefn.initializeprobpmap
+        end
+        pidxs = ParameterIndex[]
+        pvals = []
+        u0idxs = Int[]
+        u0vals = []
+        for sym in variable_symbols(oldinitprob)
+            if is_variable(sys, sym)
+                u0 !== missing || continue
+                idx = variable_index(oldinitprob, sym)
+                push!(u0idxs, idx)
+                push!(u0vals, eltype(u0)(state_values(oldinitprob, idx)))
+            else
+                p !== missing || continue
+                idx = variable_index(oldinitprob, sym)
+                push!(u0idxs, idx)
+                push!(u0vals, typeof(getp(sys, sym)(p))(state_values(oldinitprob, idx)))
+            end
+        end
+        if p !== missing
+            for sym in parameter_symbols(oldinitprob)
+                push!(pidxs, parameter_index(oldinitprob, sym))
+                if isequal(sym, get_iv(sys))
+                    push!(pvals, t0)
+                else
+                    push!(pvals, getp(sys, sym)(p))
+                end
+            end
+        end
+        newu0 = remake_buffer(oldinitprob.f.sys, state_values(oldinitprob), u0idxs, u0vals)
+        newp = remake_buffer(oldinitprob.f.sys, parameter_values(oldinitprob), pidxs, pvals)
+        initprob = remake(oldinitprob; u0 = newu0, p = newp)
+        return initprob, odefn.update_initializeprob!, odefn.initializeprobmap,
+        odefn.initializeprobpmap
+    end
     if u0 === missing || isempty(u0)
         u0 = Dict()
     elseif !(eltype(u0) <: Pair)

test/initializationsystem.jl

@@ -1,5 +1,7 @@
 using ModelingToolkit, OrdinaryDiffEq, NonlinearSolve, Test
-using SymbolicIndexingInterface
+using ForwardDiff
+using SymbolicIndexingInterface, SciMLStructures
+using SciMLStructures: Tunable
 using ModelingToolkit: t_nounits as t, D_nounits as D
 
 @parameters g
@@ -746,3 +748,35 @@ end
     prob5 = remake(prob)
     @test init(prob, Tsit5()).ps[p] ≈ 2.0
 end
+
+@testset "`remake` changes initialization problem types" begin
+    @variables x(t) y(t) z(t)
+    @parameters p q
+    @mtkbuild sys = ODESystem(
+        [D(x) ~ x * p + y * q, y^2 * q + q^2 * x ~ 0, z * p - p^2 * x * z ~ 0],
+        t; guesses = [x => 0.0, y => 0.0, z => 0.0, p => 0.0, q => 0.0])
+    prob = ODEProblem(sys, [x => 1.0], (0.0, 1.0), [p => 1.0, q => missing])
+    @test is_variable(prob.f.initializeprob, q)
+    ps = prob.p
+    newps = SciMLStructures.replace(Tunable(), ps, ForwardDiff.Dual.(ps.tunable))
+    prob2 = remake(prob; p = newps)
+    @test eltype(prob2.f.initializeprob.u0) <: ForwardDiff.Dual
+    @test eltype(prob2.f.initializeprob.p.tunable) <: ForwardDiff.Dual
+    @test prob2.f.initializeprob.u0 ≈ prob.f.initializeprob.u0
+
+    prob2 = remake(prob; u0 = ForwardDiff.Dual.(prob.u0))
+    @test eltype(prob2.f.initializeprob.u0) <: ForwardDiff.Dual
+    @test eltype(prob2.f.initializeprob.p.tunable) <: Float64
+    @test prob2.f.initializeprob.u0 ≈ prob.f.initializeprob.u0
+
+    prob2 = remake(prob; u0 = ForwardDiff.Dual.(prob.u0), p = newps)
+    @test eltype(prob2.f.initializeprob.u0) <: ForwardDiff.Dual
+    @test eltype(prob2.f.initializeprob.p.tunable) <: ForwardDiff.Dual
+    @test prob2.f.initializeprob.u0 ≈ prob.f.initializeprob.u0
+
+    prob2 = remake(prob; u0 = [x => ForwardDiff.Dual(1.0)],
+        p = [p => ForwardDiff.Dual(1.0), q => missing])
+    @test eltype(prob2.f.initializeprob.u0) <: ForwardDiff.Dual
+    @test eltype(prob2.f.initializeprob.p.tunable) <: ForwardDiff.Dual
+    @test prob2.f.initializeprob.u0 ≈ prob.f.initializeprob.u0
+end