feat: relax type restrictions in MTKParameters construction

Author: AayushSabharwal

src/systems/index_cache.jl

@@ -77,7 +77,7 @@ function IndexCache(sys::AbstractSystem)
 
     function insert_by_type!(buffers::Dict{Any, Set{BasicSymbolic}}, sym)
         sym = unwrap(sym)
-        ctype = concrete_symtype(sym)
+        ctype = symtype(sym)
         buf = get!(buffers, ctype, Set{BasicSymbolic}())
         push!(buf, sym)
     end
@@ -114,7 +114,7 @@ function IndexCache(sys::AbstractSystem)
 
     for p in parameters(sys)
         p = unwrap(p)
-        ctype = concrete_symtype(p)
+        ctype = symtype(p)
         haskey(disc_buffers, ctype) && p in disc_buffers[ctype] && continue
         haskey(dependent_buffers, ctype) && p in dependent_buffers[ctype] && continue
         insert_by_type!(
@@ -310,9 +310,3 @@ function reorder_parameters(ic::IndexCache, ps; drop_missing = false)
     end
     return result
 end
-
-concrete_symtype(x::BasicSymbolic) = concrete_symtype(symtype(x))
-concrete_symtype(::Type{Real}) = Float64
-concrete_symtype(::Type{Integer}) = Int
-concrete_symtype(::Type{A}) where {T, N, A <: Array{T, N}} = Array{concrete_symtype(T), N}
-concrete_symtype(::Type{T}) where {T} = T

src/systems/parameter_buffer.jl

@@ -1,5 +1,8 @@
 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))
+
 struct MTKParameters{T, D, C, E, N, F, G}
     tunable::T
     discrete::D
@@ -107,7 +110,7 @@ function MTKParameters(
     for (sym, val) in p
         sym = unwrap(sym)
         val = unwrap(val)
-        ctype = concrete_symtype(sym)
+        ctype = symtype(sym)
         if symbolic_type(val) !== NotSymbolic()
             continue
         end
@@ -126,19 +129,27 @@ function MTKParameters(
             end
         end
     end
+    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)
 
     if has_parameter_dependencies(sys) &&
        (pdeps = get_parameter_dependencies(sys)) !== nothing
         pdeps = Dict(k => fixpoint_sub(v, pdeps) for (k, v) in pdeps)
-        dep_exprs = ArrayPartition((wrap.(v) for v in dep_buffer)...)
+        dep_exprs = ArrayPartition((Any[missing for _ in 1:length(v)] for v in dep_buffer)...)
         for (sym, val) in pdeps
             i, j = ic.dependent_idx[sym]
             dep_exprs.x[i][j] = wrap(val)
         end
+        dep_exprs = identity.(dep_exprs)
         p = reorder_parameters(ic, full_parameters(sys))
         oop, iip = build_function(dep_exprs, p...)
         update_function_iip, update_function_oop = RuntimeGeneratedFunctions.@RuntimeGeneratedFunction(iip),
         RuntimeGeneratedFunctions.@RuntimeGeneratedFunction(oop)
+        update_function_iip(ArrayPartition(dep_buffer), tunable_buffer..., disc_buffer...,
+            const_buffer..., nonnumeric_buffer..., dep_buffer...)
+        dep_buffer = narrow_buffer_type.(dep_buffer)
     else
         update_function_iip = update_function_oop = nothing
     end
@@ -148,12 +159,26 @@ function MTKParameters(
         typeof(dep_buffer), typeof(nonnumeric_buffer), typeof(update_function_iip),
         typeof(update_function_oop)}(tunable_buffer, disc_buffer, const_buffer, dep_buffer,
         nonnumeric_buffer, update_function_iip, update_function_oop)
-    if mtkps.dependent_update_iip !== nothing
-        mtkps.dependent_update_iip(ArrayPartition(mtkps.dependent), mtkps...)
-    end
     return mtkps
 end
 
+function narrow_buffer_type(buffer::AbstractArray)
+    type = Union{}
+    for x in buffer
+        type = promote_type(type, typeof(x))
+    end
+    return convert.(type, buffer)
+end
+
+function narrow_buffer_type(buffer::AbstractArray{<:AbstractArray})
+    buffer = narrow_buffer_type.(buffer)
+    type = Union{}
+    for x in buffer
+        type = promote_type(type, eltype(x))
+    end
+    return broadcast.(convert, type, buffer)
+end
+
 function buffer_to_arraypartition(buf)
     return ArrayPartition(ntuple(i -> _buffer_to_arrp_helper(buf[i]), Val(length(buf))))
 end

test/odesystem.jl

@@ -6,7 +6,7 @@ using DiffEqBase, SparseArrays
 using StaticArrays
 using Test
 using SymbolicUtils: issym
-
+using ForwardDiff
 using ModelingToolkit: value
 using ModelingToolkit: t_nounits as t, D_nounits as D
 
@@ -1159,3 +1159,17 @@ for sys in [sys1, sys2]
         @test variable_index(sys, x[i]) == variable_index(sys, x)[i]
     end
 end
+
+# Issue#2667
+@testset "ForwardDiff through ODEProblem constructor" begin
+    @parameters P
+    @variables x(t)
+    sys = structural_simplify(ODESystem([D(x) ~ P], t, [x], [P]; name = :sys))
+
+    function x_at_1(P)
+        prob = ODEProblem(sys, [x => 0.0], (0.0, 1.0), [sys.P => P])
+        return solve(prob)(1.0)
+    end
+
+    @test_nowarn ForwardDiff.derivative(P -> x_at_1(P), 1.0)
+end