feat: add parameter type and size validation in remake_buffer and setp

Author: AayushSabharwal

src/systems/index_cache.jl

@@ -1,5 +1,5 @@
 struct BufferTemplate
-    type::DataType
+    type::Union{DataType, UnionAll}
     length::Int
 end
 
@@ -16,8 +16,11 @@ const NONNUMERIC_PORTION = Nonnumeric()
 struct ParameterIndex{P, I}
     portion::P
     idx::I
+    validate_size::Bool
 end
 
+ParameterIndex(portion, idx) = ParameterIndex(portion, idx, false)
+
 const ParamIndexMap = Dict{Union{Symbol, BasicSymbolic}, Tuple{Int, Int}}
 const UnknownIndexMap = Dict{
     Union{Symbol, BasicSymbolic}, Union{Int, UnitRange{Int}, AbstractArray{Int}}}
@@ -34,11 +37,14 @@ struct IndexCache
     constant_buffer_sizes::Vector{BufferTemplate}
     dependent_buffer_sizes::Vector{BufferTemplate}
     nonnumeric_buffer_sizes::Vector{BufferTemplate}
+    symbol_to_variable::Dict{Symbol, BasicSymbolic}
 end
 
 function IndexCache(sys::AbstractSystem)
     unks = solved_unknowns(sys)
     unk_idxs = UnknownIndexMap()
+    symbol_to_variable = Dict{Symbol, BasicSymbolic}()
+
     let idx = 1
         for sym in unks
             usym = unwrap(sym)
@@ -50,7 +56,9 @@ function IndexCache(sys::AbstractSystem)
             unk_idxs[usym] = sym_idx
 
             if hasname(sym) && (!istree(sym) || operation(sym) !== getindex)
-                unk_idxs[getname(usym)] = sym_idx
+                name = getname(usym)
+                unk_idxs[name] = sym_idx
+                symbol_to_variable[name] = sym
             end
             idx += length(sym)
         end
@@ -66,7 +74,9 @@ function IndexCache(sys::AbstractSystem)
             end
             unk_idxs[arrsym] = idxs
             if hasname(arrsym)
-                unk_idxs[getname(arrsym)] = idxs
+                name = getname(arrsym)
+                unk_idxs[name] = idxs
+                symbol_to_variable[name] = arrsym
             end
         end
     end
@@ -144,14 +154,15 @@ function IndexCache(sys::AbstractSystem)
                 idxs[default_toterm(p)] = (i, j)
                 if hasname(p) && (!istree(p) || operation(p) !== getindex)
                     idxs[getname(p)] = (i, j)
+                    symbol_to_variable[getname(p)] = p
                     idxs[getname(default_toterm(p))] = (i, j)
+                    symbol_to_variable[getname(default_toterm(p))] = p
                 end
             end
             push!(buffer_sizes, BufferTemplate(T, length(buf)))
         end
         return idxs, buffer_sizes
     end
-
     disc_idxs, discrete_buffer_sizes = get_buffer_sizes_and_idxs(disc_buffers)
     tunable_idxs, tunable_buffer_sizes = get_buffer_sizes_and_idxs(tunable_buffers)
     const_idxs, const_buffer_sizes = get_buffer_sizes_and_idxs(constant_buffers)
@@ -169,7 +180,8 @@ function IndexCache(sys::AbstractSystem)
         tunable_buffer_sizes,
         const_buffer_sizes,
         dependent_buffer_sizes,
-        nonnumeric_buffer_sizes
+        nonnumeric_buffer_sizes,
+        symbol_to_variable
     )
 end
 
@@ -190,16 +202,21 @@ function SymbolicIndexingInterface.is_parameter(ic::IndexCache, sym)
 end
 
 function SymbolicIndexingInterface.parameter_index(ic::IndexCache, sym)
+    if sym isa Symbol
+        sym = ic.symbol_to_variable[sym]
+    end
+    validate_size = Symbolics.isarraysymbolic(sym) &&
+                    Symbolics.shape(sym) !== Symbolics.Unknown()
     return if (idx = check_index_map(ic.tunable_idx, sym)) !== nothing
-        ParameterIndex(SciMLStructures.Tunable(), idx)
+        ParameterIndex(SciMLStructures.Tunable(), idx, validate_size)
     elseif (idx = check_index_map(ic.discrete_idx, sym)) !== nothing
-        ParameterIndex(SciMLStructures.Discrete(), idx)
+        ParameterIndex(SciMLStructures.Discrete(), idx, validate_size)
     elseif (idx = check_index_map(ic.constant_idx, sym)) !== nothing
-        ParameterIndex(SciMLStructures.Constants(), idx)
+        ParameterIndex(SciMLStructures.Constants(), idx, validate_size)
     elseif (idx = check_index_map(ic.nonnumeric_idx, sym)) !== nothing
-        ParameterIndex(NONNUMERIC_PORTION, idx)
+        ParameterIndex(NONNUMERIC_PORTION, idx, validate_size)
     elseif (idx = check_index_map(ic.dependent_idx, sym)) !== nothing
-        ParameterIndex(DEPENDENT_PORTION, idx)
+        ParameterIndex(DEPENDENT_PORTION, idx, validate_size)
     else
         nothing
     end
@@ -224,26 +241,6 @@ function check_index_map(idxmap, sym)
     end
 end
 
-function ParameterIndex(ic::IndexCache, p, sub_idx = ())
-    p = unwrap(p)
-    return if haskey(ic.tunable_idx, p)
-        ParameterIndex(SciMLStructures.Tunable(), (ic.tunable_idx[p]..., sub_idx...))
-    elseif haskey(ic.discrete_idx, p)
-        ParameterIndex(SciMLStructures.Discrete(), (ic.discrete_idx[p]..., sub_idx...))
-    elseif haskey(ic.constant_idx, p)
-        ParameterIndex(SciMLStructures.Constants(), (ic.constant_idx[p]..., sub_idx...))
-    elseif haskey(ic.dependent_idx, p)
-        ParameterIndex(DEPENDENT_PORTION, (ic.dependent_idx[p]..., sub_idx...))
-    elseif haskey(ic.nonnumeric_idx, p)
-        ParameterIndex(NONNUMERIC_PORTION, (ic.nonnumeric_idx[p]..., sub_idx...))
-    elseif istree(p) && operation(p) === getindex
-        _p, sub_idx... = arguments(p)
-        ParameterIndex(ic, _p, sub_idx)
-    else
-        nothing
-    end
-end
-
 function discrete_linear_index(ic::IndexCache, idx::ParameterIndex)
     idx.portion isa SciMLStructures.Discrete || error("Discrete variable index expected")
     ind = sum(temp.length for temp in ic.tunable_buffer_sizes; init = 0)

src/systems/parameter_buffer.jl

@@ -132,7 +132,8 @@ function MTKParameters(
     tunable_buffer = narrow_buffer_type.(tunable_buffer)
     disc_buffer = narrow_buffer_type.(disc_buffer)
     const_buffer = narrow_buffer_type.(const_buffer)
-    nonnumeric_buffer = narrow_buffer_type.(nonnumeric_buffer)
+    # Don't narrow nonnumeric types
+    nonnumeric_buffer = nonnumeric_buffer
 
     if has_parameter_dependencies(sys) &&
        (pdeps = get_parameter_dependencies(sys)) !== nothing
@@ -308,22 +309,31 @@ end
 
 function SymbolicIndexingInterface.set_parameter!(
         p::MTKParameters, val, idx::ParameterIndex)
-    @unpack portion, idx = idx
+    @unpack portion, idx, validate_size = idx
     i, j, k... = idx
     if portion isa SciMLStructures.Tunable
         if isempty(k)
+            if validate_size && size(val) !== size(p.tunable[i][j])
+                throw(InvalidParameterSizeException(size(p.tunable[i][j]), size(val)))
+            end
             p.tunable[i][j] = val
         else
             p.tunable[i][j][k...] = val
         end
     elseif portion isa SciMLStructures.Discrete
         if isempty(k)
+            if validate_size && size(val) !== size(p.discrete[i][j])
+                throw(InvalidParameterSizeException(size(p.discrete[i][j]), size(val)))
+            end
             p.discrete[i][j] = val
         else
             p.discrete[i][j][k...] = val
         end
     elseif portion isa SciMLStructures.Constants
         if isempty(k)
+            if validate_size && size(val) !== size(p.constant[i][j])
+                throw(InvalidParameterSizeException(size(p.constant[i][j]), size(val)))
+            end
             p.constant[i][j] = val
         else
             p.constant[i][j][k...] = val
@@ -392,14 +402,73 @@ function narrow_buffer_type_and_fallback_undefs(oldbuf::Vector, newbuf::Vector)
         isassigned(newbuf, i) || continue
         type = promote_type(type, typeof(newbuf[i]))
     end
+    if type == Union{}
+        type = eltype(oldbuf)
+    end
     for i in eachindex(newbuf)
         isassigned(newbuf, i) && continue
         newbuf[i] = convert(type, oldbuf[i])
     end
     return convert(Vector{type}, newbuf)
 end
 
-function SymbolicIndexingInterface.remake_buffer(sys, oldbuf::MTKParameters, vals::Dict)
+function validate_parameter_type(ic::IndexCache, p, index, 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
+    end
+    (; 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))
+    end
+    stype = symtype(p)
+    # Array parameters need array values...
+    if stype <: AbstractArray && !isa(val, AbstractArray)
+        throw(ParameterTypeException(:validate_parameter_type, p, stype, val))
+    end
+    # ... and must match sizes
+    if stype <: AbstractArray && Symbolics.shape(p) !== Symbolics.Unknown() &&
+       size(val) != size(p)
+        throw(InvalidParameterSizeException(p, val))
+    end
+    # Early exit
+    val isa stype && return nothing
+    if stype <: AbstractArray
+        # Arrays need handling when eltype is `Real` (accept any real array)
+        etype = eltype(stype)
+        if etype <: Real
+            etype = Real
+        end
+        # 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))
+    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))
+    end
+end
+
+function indp_to_system(indp)
+    while hasmethod(symbolic_container, Tuple{typeof(indp)})
+        indp = symbolic_container(indp)
+    end
+    return indp
+end
+
+function SymbolicIndexingInterface.remake_buffer(indp, oldbuf::MTKParameters, vals::Dict)
     newbuf = @set oldbuf.tunable = Tuple(Vector{Any}(undef, length(buf))
     for buf in oldbuf.tunable)
     @set! newbuf.discrete = Tuple(Vector{Any}(undef, length(buf))
@@ -409,9 +478,15 @@ function SymbolicIndexingInterface.remake_buffer(sys, oldbuf::MTKParameters, val
     @set! newbuf.nonnumeric = Tuple(Vector{Any}(undef, length(buf))
     for buf in newbuf.nonnumeric)
 
+    # 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)
+        validate_parameter_type(ic, p, idx, val)
         _set_parameter_unchecked!(
-            newbuf, val, parameter_index(sys, p); update_dependent = false)
+            newbuf, val, idx; update_dependent = false)
     end
 
     @set! newbuf.tunable = narrow_buffer_type_and_fallback_undefs.(
@@ -588,3 +663,15 @@ function Base.showerror(io::IO, e::MissingParametersError)
     println(io, MISSING_PARAMETERS_MESSAGE)
     println(io, e.vars)
 end
+
+function InvalidParameterSizeException(param, val)
+    DimensionMismatch("InvalidParameterSizeException: For parameter $(param) expected value of size $(size(param)). Received value $(val) of size $(size(val)).")
+end
+
+function InvalidParameterSizeException(param::Tuple, val::Tuple)
+    DimensionMismatch("InvalidParameterSizeException: Expected value of size $(param). Received value of size $(val).")
+end
+
+function ParameterTypeException(func, param, expected, val)
+    TypeError(func, "Parameter $param", expected, val)
+end

test/index_cache.jl

@@ -0,0 +1,45 @@
+using ModelingToolkit, SymbolicIndexingInterface
+using ModelingToolkit: t_nounits as t
+
+# Ensure indexes of array symbolics are cached appropriately
+@variables x(t)[1:2]
+@named sys = ODESystem(Equation[], t, [x], [])
+sys1 = complete(sys)
+@named sys = ODESystem(Equation[], t, [x...], [])
+sys2 = complete(sys)
+for sys in [sys1, sys2]
+    for (sym, idx) in [(x, 1:2), (x[1], 1), (x[2], 2)]
+        @test is_variable(sys, sym)
+        @test variable_index(sys, sym) == idx
+    end
+end
+
+@variables x(t)[1:2, 1:2]
+@named sys = ODESystem(Equation[], t, [x], [])
+sys1 = complete(sys)
+@named sys = ODESystem(Equation[], t, [x...], [])
+sys2 = complete(sys)
+for sys in [sys1, sys2]
+    @test is_variable(sys, x)
+    @test variable_index(sys, x) == [1 3; 2 4]
+    for i in eachindex(x)
+        @test is_variable(sys, x[i])
+        @test variable_index(sys, x[i]) == variable_index(sys, x)[i]
+    end
+end
+
+# Ensure Symbol to symbolic map is correct
+@parameters p1 p2[1:2] p3::String
+@variables x(t) y(t)[1:2] z(t)
+
+@named sys = ODESystem(Equation[], t, [x, y, z], [p1, p2, p3])
+sys = complete(sys)
+
+ic = ModelingToolkit.get_index_cache(sys)
+
+@test isequal(ic.symbol_to_variable[:p1], p1)
+@test isequal(ic.symbol_to_variable[:p2], p2)
+@test isequal(ic.symbol_to_variable[:p3], p3)
+@test isequal(ic.symbol_to_variable[:x], x)
+@test isequal(ic.symbol_to_variable[:y], y)
+@test isequal(ic.symbol_to_variable[:z], z)

test/mtkparameters.jl

@@ -241,3 +241,45 @@ newps = remake_buffer(
 VDual = Vector{<:ForwardDiff.Dual}
 VVDual = Vector{<:Vector{<:ForwardDiff.Dual}}
 @test newps.dependent isa Union{Tuple{VDual, VVDual}, Tuple{VVDual, VDual}}
+
+@testset "Parameter type validation" begin
+    struct Foo{T}
+        x::T
+    end
+
+    @parameters a b::Int c::Vector{Float64} d[1:2, 1:2]::Int e::Foo{Int} f::Foo
+    @named sys = ODESystem(Equation[], t, [], [a, b, c, d, e, f])
+    sys = complete(sys)
+    ps = MTKParameters(sys,
+        Dict(a => 1.0, b => 2, c => 3ones(2),
+            d => 3ones(Int, 2, 2), e => Foo(1), f => Foo("a")))
+    @test_nowarn setp(sys, c)(ps, ones(4)) # so this is fixed when SII is fixed
+    @test_throws DimensionMismatch set_parameter!(
+        ps, 4ones(Int, 3, 2), parameter_index(sys, d))
+    @test_throws DimensionMismatch set_parameter!(
+        ps, 4ones(Int, 4), parameter_index(sys, d)) # size has to match, not just length
+    @test_nowarn setp(sys, f)(ps, Foo(:a)) # can change non-concrete type
+
+    # 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 getp(sys, sym)(newps) isa Float32
+        newps = @test_nowarn remake_buffer(sys, ps, Dict(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
+    end
+
+    newps = @test_nowarn remake_buffer(sys, ps, Dict(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
+
+    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
+    @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)))
+end

test/runtests.jl

@@ -24,6 +24,7 @@ end
             @safetestset "Parsing Test" include("variable_parsing.jl")
             @safetestset "Simplify Test" include("simplify.jl")
             @safetestset "Direct Usage Test" include("direct.jl")
+            @safetestset "IndexCache Test" include("index_cache.jl")
             @safetestset "System Linearity Test" include("linearity.jl")
             @safetestset "Input Output Test" include("input_output_handling.jl")
             @safetestset "Clock Test" include("clock.jl")