fixup! feat: support callable parameters

Author: AayushSabharwal

Project.toml

@@ -24,6 +24,7 @@ ExprTools = "e2ba6199-217a-4e67-a87a-7c52f15ade04"
 Expronicon = "6b7a57c9-7cc1-4fdf-b7f5-e857abae3636"
 FindFirstFunctions = "64ca27bc-2ba2-4a57-88aa-44e436879224"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+FunctionWrappers = "069b7b12-0de2-55c6-9aab-29f3d0a68a2e"
 FunctionWrappersWrappers = "77dc65aa-8811-40c2-897b-53d922fa7daf"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
@@ -91,6 +92,7 @@ ExprTools = "0.1.10"
 Expronicon = "0.8"
 FindFirstFunctions = "1"
 ForwardDiff = "0.10.3"
+FunctionWrappers = "1.1"
 FunctionWrappersWrappers = "0.1"
 Graphs = "1.5.2"
 InteractiveUtils = "1"

src/ModelingToolkit.jl

@@ -38,6 +38,7 @@ using Base: RefValue
 using Combinatorics
 import Distributions
 import FunctionWrappersWrappers
+import FunctionWrappers: FunctionWrapper
 using URIs: URI
 using SciMLStructures
 using Compat

src/systems/index_cache.jl

@@ -109,9 +109,8 @@ function IndexCache(sys::AbstractSystem)
     constant_buffers = Dict{Any, Set{BasicSymbolic}}()
     nonnumeric_buffers = Dict{Any, Set{Union{BasicSymbolic, CallWithMetadata}}}()
 
-    function insert_by_type!(buffers::Dict{Any, S}, sym) where {S}
+    function insert_by_type!(buffers::Dict{Any, S}, sym, ctype) where {S}
         sym = unwrap(sym)
-        ctype = symtype(sym)
         buf = get!(buffers, ctype, S())
         push!(buf, sym)
     end
@@ -144,7 +143,7 @@ function IndexCache(sys::AbstractSystem)
                 clocks = get!(() -> Set{Int}(), disc_param_callbacks, sym)
                 push!(clocks, i)
             else
-                insert_by_type!(constant_buffers, sym)
+                insert_by_type!(constant_buffers, sym, symtype(sym))
             end
         end
     end
@@ -199,6 +198,9 @@ function IndexCache(sys::AbstractSystem)
     for p in parameters(sys)
         p = unwrap(p)
         ctype = symtype(p)
+        if ctype <: FnType
+            ctype = fntype_to_function_type(ctype)
+        end
         haskey(disc_idxs, p) && continue
         haskey(constant_buffers, ctype) && p in constant_buffers[ctype] && continue
         insert_by_type!(
@@ -214,7 +216,8 @@ function IndexCache(sys::AbstractSystem)
             else
                 nonnumeric_buffers
             end,
-            p
+            p,
+            ctype,
         )
     end
 
@@ -485,3 +488,6 @@ function get_buffer_template(ic::IndexCache, pidx::ParameterIndex)
         error("Unhandled portion $portion")
     end
 end
+
+fntype_to_function_type(::Type{FnType{A, R, T}}) where {A, R, T} = T
+fntype_to_function_type(::Type{FnType{A, R, Nothing}}) where {A, R} = FunctionWrapper{R, A}

src/systems/parameter_buffer.jl

@@ -2,9 +2,6 @@ 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
@@ -155,6 +152,9 @@ function MTKParameters(
         if symbolic_type(val) !== NotSymbolic()
             error("Could not evaluate value of parameter $sym. Missing values for variables in expression $val.")
         end
+        if ctype <: FnType
+            ctype = fntype_to_function_type(ctype)
+        end
         val = symconvert(ctype, val)
         done = set_value(sym, val)
         if !done && Symbolics.isarraysymbolic(sym)