feat: implement SII.remake_buffer for MTKParameters

Author: AayushSabharwal

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,43 @@ 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)
+                           (DEPENDENT_PORTION, oldbuf.dependent)
+                           (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
+    newbuf = MTKParameters(
+        newbufs..., oldbuf.dependent_update_iip, oldbuf.dependent_update_oop)
+    for (p, val) in vals
+        _set_parameter_unchecked!(
+            newbuf, val, parameter_index(sys, p); update_dependent = false)
+    end
+    if newbuf.dependent_update_iip !== nothing
+        newbuf.dependent_update_iip(ArrayPartition(newbuf.dependent), newbuf...)
+    end
+    return newbuf
+end
+
 _subarrays(v::AbstractVector) = isempty(v) ? () : (v,)
 _subarrays(v::ArrayPartition) = v.x
 _subarrays(v::Tuple) = v