Merge pull request #2570 from AayushSabharwal/as/remake-hook

feat: implement SII.remake_buffer for MTKParameters, fix bugs, add tests

Author: ChrisRackauckas

Project.toml

@@ -102,7 +102,7 @@ SimpleNonlinearSolve = "0.1.0, 1"
 SparseArrays = "1"
 SpecialFunctions = "0.7, 0.8, 0.9, 0.10, 1.0, 2"
 StaticArrays = "0.10, 0.11, 0.12, 1.0"
-SymbolicIndexingInterface = "0.3.11"
+SymbolicIndexingInterface = "0.3.12"
 SymbolicUtils = "1.0"
 Symbolics = "5.26"
 URIs = "1"

src/systems/parameter_buffer.jl

@@ -318,6 +318,46 @@ function _set_parameter_unchecked!(
         p.dependent_update_iip(ArrayPartition(p.dependent), p...)
 end
 
+function SymbolicIndexingInterface.remake_buffer(sys, oldbuf::MTKParameters, vals::Dict)
+    buftypes = Dict{Tuple{Any, Int}, Any}()
+    for (p, val) in vals
+        (idx = parameter_index(sys, p)) isa ParameterIndex || continue
+        k = (idx.portion, idx.idx[1])
+        buftypes[k] = Union{get(buftypes, k, Union{}), typeof(val)}
+    end
+
+    newbufs = []
+    for (portion, old) in [(SciMLStructures.Tunable(), oldbuf.tunable)
+                           (SciMLStructures.Discrete(), oldbuf.discrete)
+                           (SciMLStructures.Constants(), oldbuf.constant)
+                           (NONNUMERIC_PORTION, oldbuf.nonnumeric)]
+        if isempty(old)
+            push!(newbufs, old)
+            continue
+        end
+        new = Any[copy(i) for i in old]
+        for i in eachindex(new)
+            buftype = get(buftypes, (portion, i), eltype(new[i]))
+            new[i] = similar(new[i], buftype)
+        end
+        push!(newbufs, Tuple(new))
+    end
+    tmpbuf = MTKParameters(
+        newbufs[1], newbufs[2], newbufs[3], oldbuf.dependent, newbufs[4], nothing, nothing)
+    for (p, val) in vals
+        _set_parameter_unchecked!(
+            tmpbuf, val, parameter_index(sys, p); update_dependent = false)
+    end
+    if oldbuf.dependent_update_oop !== nothing
+        dependent = oldbuf.dependent_update_oop(tmpbuf...)
+    else
+        dependent = ()
+    end
+    newbuf = MTKParameters(newbufs[1], newbufs[2], newbufs[3], dependent, newbufs[4],
+        oldbuf.dependent_update_iip, oldbuf.dependent_update_oop)
+    return newbuf
+end
+
 _subarrays(v::AbstractVector) = isempty(v) ? () : (v,)
 _subarrays(v::ArrayPartition) = v.x
 _subarrays(v::Tuple) = v

test/mtkparameters.jl

@@ -2,6 +2,7 @@ using ModelingToolkit
 using ModelingToolkit: t_nounits as t, D_nounits as D, MTKParameters
 using SymbolicIndexingInterface
 using SciMLStructures: SciMLStructures, canonicalize, Tunable, Discrete, Constants
+using ForwardDiff
 
 @parameters a b c d::Integer e[1:3] f[1:3, 1:3]::Int g::Vector{AbstractFloat} h::String
 @named sys = ODESystem(
@@ -70,3 +71,33 @@ setp(sys, f[2, 2])(ps, 42) # with a sub-index
 
 setp(sys, h)(ps, "bar") # with a non-numeric
 @test getp(sys, h)(ps) == "bar"
+
+newps = remake_buffer(sys,
+    ps,
+    Dict(a => 1.0f0, b => 5.0f0, c => 2.0, d => 0x5, e => [0.4, 0.5, 0.6],
+        f => 3ones(UInt, 3, 3), g => ones(Float32, 4), h => "bar"))
+
+for fname in (:tunable, :discrete, :constant, :dependent)
+    # ensure same number of sub-buffers
+    @test length(getfield(ps, fname)) == length(getfield(newps, fname))
+end
+@test ps.dependent_update_iip === newps.dependent_update_iip
+@test ps.dependent_update_oop === newps.dependent_update_oop
+
+@test getp(sys, a)(newps) isa Float32
+@test getp(sys, b)(newps) == 2.0f0 # ensure dependent update still happened, despite explicit value
+@test getp(sys, c)(newps) isa Float64
+@test getp(sys, d)(newps) isa UInt8
+@test getp(sys, f)(newps) isa Matrix{UInt}
+# SII bug
+@test_broken getp(sys, g)(newps) isa Vector{Float32}
+
+ps = MTKParameters(sys, ivs)
+function loss(value, sys, ps)
+    @test value isa ForwardDiff.Dual
+    vals = merge(Dict(parameters(sys) .=> getp(sys, parameters(sys))(ps)), Dict(a => value))
+    ps = remake_buffer(sys, ps, vals)
+    getp(sys, a)(ps) + getp(sys, b)(ps)
+end
+
+@test ForwardDiff.derivative(x -> loss(x, sys, ps), 1.5) == 3.0