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,8 @@ function fill_u0(prob, u0; defs = nothing, use_defaults = false)
             v = (v,)
         end
         for (kk, vv, ii) in zip(k, v, idx)
+            sym_to_idx[kk] = ii
+            kk = idx_to_vsym[ii]
             sym_to_idx[kk] = ii
             idx_to_sym[ii] = kk
             idx_to_val[ii] = vv
@@ -612,6 +615,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 +629,8 @@ function fill_p(prob, p; defs = nothing, use_defaults = false)
             v = (v,)
         end
         for (kk, vv, ii) in zip(k, v, idx)
+            sym_to_idx[kk] = ii
+            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

test/remake_tests.jl

@@ -297,26 +297,54 @@ a = Remake_Test1(p = 1)
 @test @inferred remake(a, kwargs = (; a = 1)) == Remake_Test1(p = 1, a = 1)
 
 @testset "fill_u0 and fill_p ignore identical variables with different names" begin
-    sys = SymbolCache(Dict(:x => 1, :x2 => 1, :y => 2), Dict(:a => 1, :a2 => 1, :b => 2),
-        :t; defaults = Dict(:x => 1, :y => 2, :a => 3, :b => 4))
+    struct SCWrapper{S}
+        sc::S
+    end
+    SymbolicIndexingInterface.symbolic_container(s::SCWrapper) = s.sc
+    function SymbolicIndexingInterface.is_variable(s::SCWrapper, i::Symbol)
+        if i == :x2
+            return is_variable(s.sc, :x)
+        end
+        is_variable(s.sc, i)
+    end
+    function SymbolicIndexingInterface.variable_index(s::SCWrapper, i::Symbol)
+        if i == :x2
+            return variable_index(s.sc, :x)
+        end
+        variable_index(s.sc, i)
+    end
+    function SymbolicIndexingInterface.is_parameter(s::SCWrapper, i::Symbol)
+        if i == :a2
+            return is_parameter(s.sc, :a)
+        end
+        is_parameter(s.sc, i)
+    end
+    function SymbolicIndexingInterface.parameter_index(s::SCWrapper, i::Symbol)
+        if i == :a2
+            return parameter_index(s.sc, :a)
+        end
+        parameter_index(s.sc, i)
+    end
+    sys = SCWrapper(SymbolCache(Dict(:x => 1, :y => 2), Dict(:a => 1, :b => 2),
+        :t; defaults = Dict(:x => 1, :y => 2, :a => 3, :b => 4)))
     function foo(du, u, p, t)
         du .= u .* p
     end
     prob = ODEProblem(ODEFunction(foo; sys), [1.5, 2.5], (0.0, 1.0), [3.5, 4.5])
     u0 = Dict(:x2 => 2)
     newu0 = SciMLBase.fill_u0(prob, u0; defs = default_values(sys))
     @test length(newu0) == 2
-    @test get(newu0, :x2, 0) == 2
+    @test get(newu0, :x, 0) == 2
     @test get(newu0, :y, 0) == 2.5
     p = Dict(:a2 => 3)
     newp = SciMLBase.fill_p(prob, p; defs = default_values(sys))
     @test length(newp) == 2
-    @test get(newp, :a2, 0) == 3
+    @test get(newp, :a, 0) == 3
     @test get(newp, :b, 0) == 4.5
 end
 
 @testset "value of `nothing` is ignored" begin
-    sys = SymbolCache(Dict(:x => 1, :x2 => 1, :y => 2), Dict(:a => 1, :a2 => 1, :b => 2),
+    sys = SymbolCache(Dict(:x => 1, :y => 2), Dict(:a => 1, :b => 2),
         :t; defaults = Dict(:x => 1, :y => 2, :a => 3, :b => 4))
     function foo(du, u, p, t)
         du .= u .* p