Merge pull request #4459 from SciML/as/fix-dual-cache

fix: correctly promote cache buffers in `remake`

Author: AayushSabharwal

lib/ModelingToolkitBase/src/systems/nonlinear/initializesystem.jl

@@ -670,6 +670,11 @@ function promote_with_nothing(::Type{T}, p::MTKParameters) where {T}
     p = SciMLStructures.replace(SciMLStructures.Tunable(), p, tunables)
     initials = promote_with_nothing(T, p.initials)
     p = SciMLStructures.replace(SciMLStructures.Initials(), p, initials)
+    for i in eachindex(p.caches)
+        if eltype(p.caches[i]) <: AbstractFloat
+            @set! p.caches[i] = promote_with_nothing(T, p.caches[i])
+        end
+    end
     return p
 end
 

lib/ModelingToolkitBase/src/systems/problem_utils.jl

@@ -845,15 +845,19 @@ function get_mtkparameters_reconstructor(
     getter = let getters = getters, diffcache_buffer_idx = diffcache_buffer_idx
         function _getter(valp, initprob)
             oldcache = parameter_values(initprob).caches
+            tunablevals = getters[1](valp)
             initialvals = getters[2](valp)
             nonnumerics = getters[5](valp)
             if !iszero(diffcache_buffer_idx)
                 @set! nonnumerics[diffcache_buffer_idx] = DiffCacheAllocatorAPIWrapper{eltype(initialvals)}.(nonnumerics[diffcache_buffer_idx])
             end
-            return MTKParameters(
-                getters[1](valp), initialvals, getters[3](valp),
-                getters[4](valp), nonnumerics, oldcache isa Tuple{} ? () :
-                    copy.(oldcache)
+            return promote_with_nothing(
+                promote_type_with_nothing(eltype(tunablevals), initialvals),
+                MTKParameters(
+                    tunablevals, initialvals, getters[3](valp),
+                    getters[4](valp), nonnumerics, oldcache isa Tuple{} ? () :
+                        copy.(oldcache)
+                )
             )
         end
     end

lib/ModelingToolkitBase/test/initializationsystem.jl

@@ -2013,3 +2013,30 @@ end
         @test !(prob.f.initialization_data.initializeprob isa SCCNonlinearProblem)
     end
 end
+
+if @isdefined(ModelingToolkit)
+    # `SCCNonlinearProblem` uses the cache buffers
+    @testset "Cache buffers are correctly promoted during initialization" begin
+        @parameters g
+        @variables x(t) y(t) [state_priority = 10] λ(t) yˍt(t) xˍt(t) xˍtt(t)
+        @mtkcomplete pend = index_reduced_pend()
+        g_true = 9.81
+        prob = ODEProblem(
+            pend, [x => 1, D(y) => 0, g => g_true], (0.0, 0.5);
+            guesses = [λ => 0, y => 1, x => 1], missing_guess_value
+        )
+        @test !isempty(prob.f.initialization_data.initializeprob.p.caches)
+        saveat = range(prob.tspan..., length = 30)
+        sol = solve(prob, Rodas5P(); saveat)
+        @test SciMLBase.successful_retcode(sol)
+        setter = setp_oop(prob, [g])
+
+        function costfn(theta)
+            p = setter(prob, theta)
+            newprob = SciMLBase.remake(prob; p)
+            sol = solve(newprob, Rodas5P(); saveat = 0.1)
+            sum(abs2, sol[x])
+        end
+        @test_nowarn ForwardDiff.gradient(costfn, [1.2])
+    end
+end

lib/ModelingToolkitBase/test/split_parameters.jl

@@ -162,7 +162,13 @@ if @isdefined(ModelingToolkit)
     @named d = Step(start_time = 1.0, duration = 10.0, offset = 0.0, height = 1.0) # Disturbance
     model_outputs = [model.inertia1.w, model.inertia2.w, model.inertia1.phi, model.inertia2.phi] # This is the state realization we want to control
     inputs = [model.torque.tau.u]
-    op = [model.torque.tau.u => 0.0]
+    op = [
+        model.inertia1.w => 1.0
+        model.inertia2.w => 1.0
+        model.inertia1.phi => 1.0
+        model.inertia2.phi => 1.0
+        model.torque.tau.u => 0.0
+    ]
     matrices, ssys = ModelingToolkit.linearize(
         wr(model), inputs, model_outputs; op,
         guesses = [model.inertia2.flange_a.phi => 0.0, model.inertia1.flange_b.phi => 0.0]
@@ -197,7 +203,7 @@ if @isdefined(ModelingToolkit)
         connect(add.output, :u, model.torque.tau)
     ]
     @named closed_loop = System(connections, t, systems = [model, state_feedback, add, d])
-    S = get_sensitivity(closed_loop, :u)
+    S = get_sensitivity(closed_loop, :u; op = x_costs)
 end
 
 @testset "Indexing MTKParameters with ParameterIndex" begin