Merge pull request #3020 from AayushSabharwal/as/fix-remake-buffer

feat: update to new `remake_buffer` signature

Author: ChrisRackauckas

Project.toml

@@ -109,15 +109,15 @@ PrecompileTools = "1"
 RecursiveArrayTools = "3.26"
 Reexport = "0.2, 1"
 RuntimeGeneratedFunctions = "0.5.9"
-SciMLBase = "2.46"
+SciMLBase = "2.52.1"
 SciMLStructures = "1.0"
 Serialization = "1"
 Setfield = "0.7, 0.8, 1"
 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.28"
+SymbolicIndexingInterface = "0.3.29"
 SymbolicUtils = "3.2"
 Symbolics = "6.3"
 URIs = "1"

src/systems/index_cache.jl

@@ -18,6 +18,7 @@ struct ParameterIndex{P, I}
 end
 
 ParameterIndex(portion, idx) = ParameterIndex(portion, idx, false)
+ParameterIndex(p::ParameterIndex) = ParameterIndex(p.portion, p.idx, false)
 
 struct DiscreteIndex
     # of all buffers corresponding to types, which one
@@ -318,7 +319,8 @@ function SymbolicIndexingInterface.parameter_index(ic::IndexCache, sym)
         sym = get(ic.symbol_to_variable, sym, nothing)
         sym === nothing && return nothing
     end
-    validate_size = Symbolics.isarraysymbolic(sym) &&
+    sym = unwrap(sym)
+    validate_size = Symbolics.isarraysymbolic(sym) && symtype(sym) <: AbstractArray &&
                     Symbolics.shape(sym) !== Symbolics.Unknown()
     return if (idx = check_index_map(ic.tunable_idx, sym)) !== nothing
         ParameterIndex(SciMLStructures.Tunable(), idx, validate_size)
@@ -459,3 +461,19 @@ function iterated_buffer_index(ic::IndexCache, ind::ParameterIndex)
     end
     error("Unhandled portion $(ind.portion)")
 end
+
+function get_buffer_template(ic::IndexCache, pidx::ParameterIndex)
+    (; portion, idx) = pidx
+
+    if portion isa SciMLStructures.Tunable
+        return ic.tunable_buffer_size
+    elseif portion isa SciMLStructures.Discrete
+        return ic.discrete_buffer_sizes[idx[1]][1]
+    elseif portion isa SciMLStructures.Constants
+        return ic.constant_buffer_sizes[idx[1]]
+    elseif portion isa Nonnumeric
+        return ic.nonnumeric_buffer_sizes[idx[1]]
+    else
+        error("Unhandled portion $portion")
+    end
+end

src/systems/parameter_buffer.jl

@@ -396,37 +396,6 @@ function SymbolicIndexingInterface.set_parameter!(
     return nothing
 end
 
-function _set_parameter_unchecked!(
-        p::MTKParameters, val, idx::ParameterIndex; update_dependent = true)
-    @unpack portion, idx = idx
-    if portion isa SciMLStructures.Tunable
-        p.tunable[idx] = val
-    else
-        i, j, k... = idx
-        if portion isa SciMLStructures.Discrete
-            if isempty(k)
-                p.discrete[i][j] = val
-            else
-                p.discrete[i][j][k...] = val
-            end
-        elseif portion isa SciMLStructures.Constants
-            if isempty(k)
-                p.constant[i][j] = val
-            else
-                p.constant[i][j][k...] = val
-            end
-        elseif portion === NONNUMERIC_PORTION
-            if isempty(k)
-                p.nonnumeric[i][j] = val
-            else
-                p.nonnumeric[i][j][k...] = val
-            end
-        else
-            error("Unhandled portion $portion")
-        end
-    end
-end
-
 function narrow_buffer_type_and_fallback_undefs(
         oldbuf::AbstractVector, newbuf::AbstractVector)
     type = Union{}
@@ -448,31 +417,42 @@ function narrow_buffer_type_and_fallback_undefs(
     return newerbuf
 end
 
-function validate_parameter_type(ic::IndexCache, p, index, val)
+function validate_parameter_type(ic::IndexCache, p, idx::ParameterIndex, val)
     p = unwrap(p)
     if p isa Symbol
         p = get(ic.symbol_to_variable, p, nothing)
-        if p === nothing
-            @warn "No matching variable found for `Symbol` $p, skipping type validation."
-            return nothing
-        end
+        p === nothing && return validate_parameter_type(ic, idx, val)
+    end
+    stype = symtype(p)
+    sz = if stype <: AbstractArray
+        Symbolics.shape(p) == Symbolics.Unknown() ? Symbolics.Unknown() : size(p)
+    elseif stype <: Number
+        size(p)
+    else
+        Symbolics.Unknown()
     end
+    validate_parameter_type(ic, stype, sz, p, idx, val)
+end
+
+function validate_parameter_type(ic::IndexCache, idx::ParameterIndex, val)
+    validate_parameter_type(
+        ic, get_buffer_template(ic, idx).type, Symbolics.Unknown(), nothing, idx, val)
+end
+
+function validate_parameter_type(ic::IndexCache, stype, sz, sym, index, val)
     (; portion) = index
     # Nonnumeric parameters have to match the type
     if portion === NONNUMERIC_PORTION
-        stype = symtype(p)
         val isa stype && return nothing
-        throw(ParameterTypeException(:validate_parameter_type, p, stype, val))
+        throw(ParameterTypeException(:validate_parameter_type, sym, stype, val))
     end
-    stype = symtype(p)
     # Array parameters need array values...
     if stype <: AbstractArray && !isa(val, AbstractArray)
-        throw(ParameterTypeException(:validate_parameter_type, p, stype, val))
+        throw(ParameterTypeException(:validate_parameter_type, sym, stype, val))
     end
     # ... and must match sizes
-    if stype <: AbstractArray && Symbolics.shape(p) !== Symbolics.Unknown() &&
-       size(val) != size(p)
-        throw(InvalidParameterSizeException(p, val))
+    if stype <: AbstractArray && sz != Symbolics.Unknown() && size(val) != sz
+        throw(InvalidParameterSizeException(sym, val))
     end
     # Early exit
     val isa stype && return nothing
@@ -485,15 +465,15 @@ function validate_parameter_type(ic::IndexCache, p, index, val)
         # This is for duals and other complicated number types
         etype = SciMLBase.parameterless_type(etype)
         eltype(val) <: etype || throw(ParameterTypeException(
-            :validate_parameter_type, p, AbstractArray{etype}, val))
+            :validate_parameter_type, sym, AbstractArray{etype}, val))
     else
         # Real check
         if stype <: Real
             stype = Real
         end
         stype = SciMLBase.parameterless_type(stype)
         val isa stype ||
-            throw(ParameterTypeException(:validate_parameter_type, p, stype, val))
+            throw(ParameterTypeException(:validate_parameter_type, sym, stype, val))
     end
 end
 
@@ -504,45 +484,69 @@ function indp_to_system(indp)
     return indp
 end
 
-function SymbolicIndexingInterface.remake_buffer(indp, oldbuf::MTKParameters, vals::Dict)
-    newbuf = @set oldbuf.tunable = Vector{Any}(undef, length(oldbuf.tunable))
+function SymbolicIndexingInterface.remake_buffer(indp, oldbuf::MTKParameters, idxs, vals)
+    newbuf = @set oldbuf.tunable = similar(oldbuf.tunable, Any)
     @set! newbuf.discrete = Tuple(similar(buf, Any) for buf in newbuf.discrete)
-    @set! newbuf.constant = Tuple(Vector{Any}(undef, length(buf))
-    for buf in newbuf.constant)
-    @set! newbuf.nonnumeric = Tuple(Vector{Any}(undef, length(buf))
-    for buf in newbuf.nonnumeric)
-
-    syms = collect(keys(vals))
-    vals = Dict{Any, Any}(vals)
-    for sym in syms
-        symbolic_type(sym) == ArraySymbolic() || continue
-        is_parameter(indp, sym) && continue
-        stype = symtype(unwrap(sym))
-        stype <: AbstractArray || continue
-        Symbolics.shape(sym) == Symbolics.Unknown() && continue
-        for i in eachindex(sym)
-            vals[sym[i]] = vals[sym][i]
+    @set! newbuf.constant = Tuple(similar(buf, Any) for buf in newbuf.constant)
+    @set! newbuf.nonnumeric = Tuple(similar(buf, Any) for buf in newbuf.nonnumeric)
+
+    function handle_parameter(ic, sym, idx, val)
+        if sym === nothing
+            validate_parameter_type(ic, idx, val)
+        else
+            validate_parameter_type(ic, sym, idx, val)
         end
+        # `ParameterIndex(idx)` turns off size validation since it relies on there
+        # being an existing value
+        set_parameter!(newbuf, val, ParameterIndex(idx))
     end
 
+    handled_idxs = Set{ParameterIndex}()
     # If the parameter buffer is an `MTKParameters` object, `indp` must eventually drill
     # down to an `AbstractSystem` using `symbolic_container`. We leverage this to get
     # the index cache.
     ic = get_index_cache(indp_to_system(indp))
-    for (p, val) in vals
-        idx = parameter_index(indp, p)
-        if idx !== nothing
-            validate_parameter_type(ic, p, idx, val)
-            _set_parameter_unchecked!(
-                newbuf, val, idx; update_dependent = false)
-        elseif symbolic_type(p) == ArraySymbolic()
-            for (i, j) in zip(eachindex(p), eachindex(val))
-                pi = p[i]
-                idx = parameter_index(indp, pi)
-                validate_parameter_type(ic, pi, idx, val[j])
-                _set_parameter_unchecked!(
-                    newbuf, val[j], idx; update_dependent = false)
+    for (idx, val) in zip(idxs, vals)
+        sym = nothing
+        if symbolic_type(idx) == ScalarSymbolic()
+            sym = idx
+            idx = parameter_index(ic, sym)
+            if idx === nothing
+                @warn "Symbolic variable $sym is not a (non-dependent) parameter in the system"
+                continue
+            end
+            idx in handled_idxs && continue
+            handle_parameter(ic, sym, idx, val)
+            push!(handled_idxs, idx)
+        elseif symbolic_type(idx) == ArraySymbolic()
+            sym = idx
+            idx = parameter_index(ic, sym)
+            if idx === nothing
+                Symbolics.shape(sym) == Symbolics.Unknown() &&
+                    throw(ParameterNotInSystem(sym))
+                size(sym) == size(val) || throw(InvalidParameterSizeException(sym, val))
+
+                for (i, vali) in zip(eachindex(sym), eachindex(val))
+                    idx = parameter_index(ic, sym[i])
+                    if idx === nothing
+                        @warn "Symbolic variable $sym is not a (non-dependent) parameter in the system"
+                        continue
+                    end
+                    # Intentionally don't check handled_idxs here because array variables always take priority
+                    # See Issue#2804
+                    handle_parameter(ic, sym[i], idx, val[vali])
+                    push!(handled_idxs, idx)
+                end
+            else
+                idx in handled_idxs && continue
+                handle_parameter(ic, sym, idx, val)
+                push!(handled_idxs, idx)
             end
+        else # NotSymbolic
+            if !(idx isa ParameterIndex)
+                throw(ArgumentError("Expected index for parameter to be a symbolic variable or `ParameterIndex`, got $idx"))
+            end
+            handle_parameter(ic, nothing, idx, val)
         end
     end
 
@@ -688,7 +692,7 @@ function jacobian_wrt_vars(pf::F, p::MTKParameters, input_idxs, chunk::C) where
 
         function (p_small_inner)
             for (i, val) in zip(input_idxs, p_small_inner)
-                _set_parameter_unchecked!(p_big, val, i)
+                set_parameter!(p_big, val, i)
             end
             return if pf isa SciMLBase.ParamJacobianWrapper
                 buffer = Array{dualtype}(undef, size(pf.u))
@@ -735,3 +739,11 @@ end
 function ParameterTypeException(func, param, expected, val)
     TypeError(func, "Parameter $param", expected, val)
 end
+
+struct ParameterNotInSystem <: Exception
+    p::Any
+end
+
+function Base.showerror(io::IO, e::ParameterNotInSystem)
+    println(io, "Symbolic variable $(e.p) is not a parameter in the system")
+end

test/mtkparameters.jl

@@ -9,7 +9,7 @@ using JET
 
 @parameters a b c(t) d::Integer e[1:3] f[1:3, 1:3]::Int g::Vector{AbstractFloat} h::String
 @named sys = ODESystem(
-    Equation[], t, [], [a, c, d, e, f, g, h], parameter_dependencies = [b => 2a],
+    Equation[], t, [], [a, c, d, e, f, g, h], parameter_dependencies = [b ~ 2a],
     continuous_events = [[a ~ 0] => [c ~ 0]], defaults = Dict(a => 0.0))
 sys = complete(sys)
 
@@ -72,10 +72,12 @@ setp(sys, g)(ps, ones(100)) # with non-fixed-length array
 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 => Float32[0.4, 0.5, 0.6],
-        f => 3ones(UInt, 3, 3), g => ones(Float32, 4), h => "bar"))
+varmap = Dict(a => 1.0f0, b => 5.0f0, c => 2.0, d => 0x5, e => Float32[0.4, 0.5, 0.6],
+    f => 3ones(UInt, 3, 3), g => ones(Float32, 4), h => "bar")
+@test_deprecated remake_buffer(sys, ps, varmap)
+@test_warn ["Symbolic variable b", "non-dependent", "parameter"] remake_buffer(
+    sys, ps, keys(varmap), values(varmap))
+newps = remake_buffer(sys, ps, keys(varmap), values(varmap))
 
 for fname in (:tunable, :discrete, :constant)
     # ensure same number of sub-buffers
@@ -92,8 +94,7 @@ end
 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)
+    ps = remake_buffer(sys, ps, (a,), (value,))
     getp(sys, a)(ps) + getp(sys, b)(ps)
 end
 
@@ -115,7 +116,7 @@ p = MTKParameters(osys, ps, u0)
 @named sys = ODESystem(Equation[], t, [], [p, q, r])
 sys = complete(sys)
 ps = MTKParameters(sys, [p => 1.0, q => 2.0, r => 3.0])
-newps = remake_buffer(sys, ps, Dict(p => 1.0f0))
+newps = remake_buffer(sys, ps, (p,), (1.0f0,))
 @test newps.tunable isa Vector{Float32}
 @test newps.tunable == [1.0f0, 2.0f0, 3.0f0]
 
@@ -227,19 +228,6 @@ end
 
 @test_nowarn ForwardDiff.gradient(loss, collect(tunables))
 
-# Ensure dependent parameters are `Tuple{...}` and not `ArrayPartition` when using
-# `remake_buffer`.
-@parameters p1 p2 p3[1:2] p4[1:2]
-@named sys = ODESystem(
-    Equation[], t, [], [p1, p2, p3, p4]; parameter_dependencies = [p2 => 2p1, p4 => 3p3])
-sys = complete(sys)
-ps = MTKParameters(sys, [p1 => 1.0, p3 => [2.0, 3.0]])
-@test getp(sys, p2)(ps) == 2.0
-@test getp(sys, p4)(ps) == [6.0, 9.0]
-
-newps = remake_buffer(
-    sys, ps, Dict(p1 => ForwardDiff.Dual(2.0), p3 => ForwardDiff.Dual.([3.0, 4.0])))
-
 VDual = Vector{<:ForwardDiff.Dual}
 VVDual = Vector{<:Vector{<:ForwardDiff.Dual}}
 
@@ -263,26 +251,26 @@ VVDual = Vector{<:Vector{<:ForwardDiff.Dual}}
 
     # Same flexibility is afforded to `b::Int` to allow for ForwardDiff
     for sym in [a, b]
-        @test_nowarn remake_buffer(sys, ps, Dict(sym => 1))
-        newps = @test_nowarn remake_buffer(sys, ps, Dict(sym => 1.0f0)) # Can change type if it's numeric
+        @test_nowarn remake_buffer(sys, ps, (sym,), (1,))
+        newps = @test_nowarn remake_buffer(sys, ps, (sym,), (1.0f0,)) # Can change type if it's numeric
         @test getp(sys, sym)(newps) isa Float32
-        newps = @test_nowarn remake_buffer(sys, ps, Dict(sym => ForwardDiff.Dual(1.0)))
+        newps = @test_nowarn remake_buffer(sys, ps, sym, ForwardDiff.Dual(1.0))
         @test getp(sys, sym)(newps) isa ForwardDiff.Dual
-        @test_throws TypeError remake_buffer(sys, ps, Dict(sym => :a)) # still has to be numeric
+        @test_throws TypeError remake_buffer(sys, ps, (sym,), (:a,)) # still has to be numeric
     end
 
-    newps = @test_nowarn remake_buffer(sys, ps, Dict(c => view(1.0:4.0, 2:4))) # can change type of array
+    newps = @test_nowarn remake_buffer(sys, ps, (c,), (view(1.0:4.0, 2:4),)) # can change type of array
     @test getp(sys, c)(newps) == 2.0:4.0
     @test parameter_values(newps, parameter_index(sys, c)) ≈ [2.0, 3.0, 4.0]
-    @test_throws TypeError remake_buffer(sys, ps, Dict(c => [:a, :b, :c])) # can't arbitrarily change eltype
-    @test_throws TypeError remake_buffer(sys, ps, Dict(c => :a)) # can't arbitrarily change type
+    @test_throws TypeError remake_buffer(sys, ps, (c,), ([:a, :b, :c],)) # can't arbitrarily change eltype
+    @test_throws TypeError remake_buffer(sys, ps, (c,), (:a,)) # can't arbitrarily change type
 
-    newps = @test_nowarn remake_buffer(sys, ps, Dict(d => ForwardDiff.Dual.(ones(2, 2)))) # can change eltype
-    @test_throws TypeError remake_buffer(sys, ps, Dict(d => [:a :b; :c :d])) # eltype still has to be numeric
+    newps = @test_nowarn remake_buffer(sys, ps, (d,), (ForwardDiff.Dual.(ones(2, 2)),)) # can change eltype
+    @test_throws TypeError remake_buffer(sys, ps, (d,), ([:a :b; :c :d],)) # eltype still has to be numeric
     @test getp(sys, d)(newps) isa Matrix{<:ForwardDiff.Dual}
 
-    @test_throws TypeError remake_buffer(sys, ps, Dict(e => Foo(2.0))) # need exact same type for nonnumeric
-    @test_nowarn remake_buffer(sys, ps, Dict(f => Foo(:a)))
+    @test_throws TypeError remake_buffer(sys, ps, (e,), (Foo(2.0),)) # need exact same type for nonnumeric
+    @test_nowarn remake_buffer(sys, ps, (f,), (Foo(:a),))
 end
 
 @testset "Error on missing parameter defaults" begin
@@ -292,7 +280,7 @@ end
     @test_throws ["Could not evaluate", "b", "Missing", "2c"] MTKParameters(sys, [a => 1.0])
 end
 
-@testset "Issue#3804" begin
+@testset "Issue#2804" begin
     @parameters k[1:4]
     @variables (V(t))[1:2]
     eqs = [