fix: fix remake with u0 dependent on Symbol parameter

Author: AayushSabharwal

src/remake.jl

@@ -567,6 +567,7 @@ end
 
 function fill_u0(prob, u0; defs = nothing, use_defaults = false)
     vsyms = variable_symbols(prob)
+    idx_to_vsym = anydict(variable_index(prob, sym) => sym for sym in vsyms)
     sym_to_idx = anydict()
     idx_to_sym = anydict()
     idx_to_val = anydict()
@@ -580,6 +581,7 @@ function fill_u0(prob, u0; defs = nothing, use_defaults = false)
             v = (v,)
         end
         for (kk, vv, ii) in zip(k, v, idx)
+            kk = idx_to_vsym[ii]
             sym_to_idx[kk] = ii
             idx_to_sym[ii] = kk
             idx_to_val[ii] = vv
@@ -612,6 +614,7 @@ end
 
 function fill_p(prob, p; defs = nothing, use_defaults = false)
     psyms = parameter_symbols(prob)
+    idx_to_psym = anydict(parameter_index(prob, sym) => sym for sym in psyms)
     sym_to_idx = anydict()
     idx_to_sym = anydict()
     idx_to_val = anydict()
@@ -625,6 +628,7 @@ function fill_p(prob, p; defs = nothing, use_defaults = false)
             v = (v,)
         end
         for (kk, vv, ii) in zip(k, v, idx)
+            kk = idx_to_psym[ii]
             sym_to_idx[kk] = ii
             idx_to_sym[ii] = kk
             idx_to_val[ii] = vv

test/downstream/modelingtoolkit_remake.jl

@@ -229,3 +229,13 @@ end
     prob2 = remake(prob; u0 = [y => 3.0], p = Dict())
     @test prob2.ps[p] ≈ 4.0
 end
+
+@testset "u0 dependent on parameter given as Symbol" begin
+    @variables x(t)
+    @parameters p
+    @mtkbuild sys = ODESystem([D(x) ~ x * p], t)
+    prob = ODEProblem(sys, [x => 1.0], (0.0, 1.0), [p => 1.0])
+    @test prob.ps[p] ≈ 1.0
+    prob2 = remake(prob; u0 = [x => p], p = [:p => 2.0])
+    @test prob2[x] ≈ 2.0
+end