feat: support callable parameters

Author: AayushSabharwal

src/ModelingToolkit.jl

@@ -63,7 +63,8 @@ using Symbolics: _parse_vars, value, @derivatives, get_variables,
                  VariableSource, getname, variable, Connection, connect,
                  NAMESPACE_SEPARATOR, set_scalar_metadata, setdefaultval,
                  initial_state, transition, activeState, entry, hasnode,
-                 ticksInState, timeInState, fixpoint_sub, fast_substitute
+                 ticksInState, timeInState, fixpoint_sub, fast_substitute,
+                 CallWithMetadata
 const NAMESPACE_SEPARATOR_SYMBOL = Symbol(NAMESPACE_SEPARATOR)
 import Symbolics: rename, get_variables!, _solve, hessian_sparsity,
                   jacobian_sparsity, isaffine, islinear, _iszero, _isone,

src/systems/index_cache.jl

@@ -32,6 +32,8 @@ struct DiscreteIndex
 end
 
 const ParamIndexMap = Dict{BasicSymbolic, Tuple{Int, Int}}
+const NonnumericMap = Dict{
+    Union{BasicSymbolic, Symbolics.CallWithMetadata}, Tuple{Int, Int}}
 const UnknownIndexMap = Dict{
     BasicSymbolic, Union{Int, UnitRange{Int}, AbstractArray{Int}}}
 const TunableIndexMap = Dict{BasicSymbolic,
@@ -45,20 +47,20 @@ struct IndexCache
     callback_to_clocks::Dict{Any, Vector{Int}}
     tunable_idx::TunableIndexMap
     constant_idx::ParamIndexMap
-    nonnumeric_idx::ParamIndexMap
+    nonnumeric_idx::NonnumericMap
     observed_syms::Set{BasicSymbolic}
     dependent_pars::Set{BasicSymbolic}
     discrete_buffer_sizes::Vector{Vector{BufferTemplate}}
     tunable_buffer_size::BufferTemplate
     constant_buffer_sizes::Vector{BufferTemplate}
     nonnumeric_buffer_sizes::Vector{BufferTemplate}
-    symbol_to_variable::Dict{Symbol, BasicSymbolic}
+    symbol_to_variable::Dict{Symbol, Union{BasicSymbolic, CallWithMetadata}}
 end
 
 function IndexCache(sys::AbstractSystem)
     unks = solved_unknowns(sys)
     unk_idxs = UnknownIndexMap()
-    symbol_to_variable = Dict{Symbol, BasicSymbolic}()
+    symbol_to_variable = Dict{Symbol, Union{BasicSymbolic, CallWithMetadata}}()
 
     let idx = 1
         for sym in unks
@@ -105,12 +107,12 @@ function IndexCache(sys::AbstractSystem)
 
     tunable_buffers = Dict{Any, Set{BasicSymbolic}}()
     constant_buffers = Dict{Any, Set{BasicSymbolic}}()
-    nonnumeric_buffers = Dict{Any, Set{BasicSymbolic}}()
+    nonnumeric_buffers = Dict{Any, Set{Union{BasicSymbolic, CallWithMetadata}}}()
 
-    function insert_by_type!(buffers::Dict{Any, Set{BasicSymbolic}}, sym)
+    function insert_by_type!(buffers::Dict{Any, S}, sym) where {S}
         sym = unwrap(sym)
         ctype = symtype(sym)
-        buf = get!(buffers, ctype, Set{BasicSymbolic}())
+        buf = get!(buffers, ctype, S())
         push!(buf, sym)
     end
 
@@ -216,8 +218,8 @@ function IndexCache(sys::AbstractSystem)
         )
     end
 
-    function get_buffer_sizes_and_idxs(buffers::Dict{Any, Set{BasicSymbolic}})
-        idxs = ParamIndexMap()
+    function get_buffer_sizes_and_idxs(T, buffers::Dict)
+        idxs = T()
         buffer_sizes = BufferTemplate[]
         for (i, (T, buf)) in enumerate(buffers)
             for (j, p) in enumerate(buf)
@@ -229,13 +231,18 @@ function IndexCache(sys::AbstractSystem)
                 idxs[rp] = (i, j)
                 idxs[rttp] = (i, j)
             end
+            if T <: Symbolics.FnType
+                T = Any
+            end
             push!(buffer_sizes, BufferTemplate(T, length(buf)))
         end
         return idxs, buffer_sizes
     end
 
-    const_idxs, const_buffer_sizes = get_buffer_sizes_and_idxs(constant_buffers)
-    nonnumeric_idxs, nonnumeric_buffer_sizes = get_buffer_sizes_and_idxs(nonnumeric_buffers)
+    const_idxs, const_buffer_sizes = get_buffer_sizes_and_idxs(
+        ParamIndexMap, constant_buffers)
+    nonnumeric_idxs, nonnumeric_buffer_sizes = get_buffer_sizes_and_idxs(
+        NonnumericMap, nonnumeric_buffers)
 
     tunable_idxs = TunableIndexMap()
     tunable_buffer_size = 0
@@ -397,7 +404,8 @@ function reorder_parameters(ic::IndexCache, ps; drop_missing = false)
     for temp in ic.discrete_buffer_sizes)
     const_buf = Tuple(BasicSymbolic[unwrap(variable(:DEF)) for _ in 1:(temp.length)]
     for temp in ic.constant_buffer_sizes)
-    nonnumeric_buf = Tuple(BasicSymbolic[unwrap(variable(:DEF)) for _ in 1:(temp.length)]
+    nonnumeric_buf = Tuple(Union{BasicSymbolic, CallWithMetadata}[unwrap(variable(:DEF))
+                                                                  for _ in 1:(temp.length)]
     for temp in ic.nonnumeric_buffer_sizes)
     for p in ps
         p = unwrap(p)

src/systems/parameter_buffer.jl

@@ -2,6 +2,9 @@ symconvert(::Type{Symbolics.Struct{T}}, x) where {T} = convert(T, x)
 symconvert(::Type{T}, x) where {T} = convert(T, x)
 symconvert(::Type{Real}, x::Integer) = convert(Float64, x)
 symconvert(::Type{V}, x) where {V <: AbstractArray} = convert(V, symconvert.(eltype(V), x))
+function symconvert(::Type{T}, x) where {T <: FnType}
+    isempty(methods(x)) ? error("Expected value $x to be a callable") : x
+end
 
 struct MTKParameters{T, D, C, N}
     tunable::T

src/systems/systemstructure.jl

@@ -661,7 +661,8 @@ function structural_simplify!(state::TearingState, io = nothing; simplify = fals
             @set! sys.defaults = merge(ModelingToolkit.defaults(sys),
                 Dict(v => 0.0 for v in Iterators.flatten(inputs)))
         end
-        ps = [setmetadata(sym, VariableTimeDomain, get(time_domains, sym, Continuous))
+        ps = [sym isa CallWithMetadata ? sym :
+              setmetadata(sym, VariableTimeDomain, get(time_domains, sym, Continuous))
               for sym in get_ps(sys)]
         @set! sys.ps = ps
     else

test/split_parameters.jl

@@ -219,3 +219,24 @@ S = get_sensitivity(closed_loop, :u)
     @test ps[ParameterIndex(Tunable(), 1:8)] == collect(1.0:8.0) .+ 0.5
     @test ps[ParameterIndex(Discrete(), (2, 1, 2, 2))] == 5
 end
+
+@testset "Callable parameters" begin
+    _f1(x) = 2x
+    struct Foo end
+    (::Foo)(x) = 3x
+    @variables x(t)
+    @parameters fn(..) = _f1
+    @mtkbuild sys = ODESystem(D(x) ~ fn(x), t, [x], [fn])
+    @test is_parameter(sys, fn)
+    @test ModelingToolkit.defaults(sys)[fn] == _f1
+    prob = ODEProblem(sys, [x => 1.0], (0.0, 1.0))
+    @test_broken @inferred prob.ps[fn]
+    @test_broken @inferred prob.f(prob.u0, prob.p, prob.tspan[1])
+    sol = solve(prob; abstol = 1e-10, reltol = 1e-10)
+    @test sol.u[end][] ≈ exp(2.0)
+    prob = ODEProblem(sys, [x => 1.0], (0.0, 1.0), [fn => Foo()])
+    @test_broken @inferred prob.ps[fn]
+    @test_broken @inferred prob.f(prob.u0, prob.p, prob.tspan[1])
+    sol = solve(prob; abstol = 1e-10, reltol = 1e-10)
+    @test sol.u[end][] ≈ exp(3.0)
+end