Merge pull request #2686 from AayushSabharwal/as/mtkparams-types

feat: relax type restrictions in `MTKParameters` construction

Author: ChrisRackauckas

src/systems/index_cache.jl

@@ -79,7 +79,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
@@ -116,7 +116,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!(
@@ -312,9 +312,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
@@ -67,7 +70,7 @@ function MTKParameters(
         end
     end
 
-    isempty(missing_params) || throw(MissingVariablesError(collect(missing_params)))
+    isempty(missing_params) || throw(MissingParametersError(collect(missing_params)))
 
     tunable_buffer = Tuple(Vector{temp.type}(undef, temp.length)
     for temp in ic.tunable_buffer_sizes)
@@ -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
@@ -550,3 +575,17 @@ function as_duals(p::MTKParameters, dualtype)
     discrete = dualtype.(p.discrete)
     return MTKParameters{typeof(tunable), typeof(discrete)}(tunable, discrete)
 end
+
+const MISSING_PARAMETERS_MESSAGE = """
+                                Some parameters are missing from the variable map.
+                                Please provide a value or default for the following variables:
+                                """
+
+struct MissingParametersError <: Exception
+    vars::Any
+end
+
+function Base.showerror(io::IO, e::MissingParametersError)
+    println(io, MISSING_PARAMETERS_MESSAGE)
+    println(io, e.vars)
+end

test/mtkparameters.jl

@@ -130,7 +130,7 @@ u0 = [X => 1.0]
 tspan = (0.0, 100.0)
 ps = [p => 1.0] # Value for `d` is missing
 
-@test_throws ModelingToolkit.MissingVariablesError ODEProblem(sys, u0, tspan, ps)
+@test_throws ModelingToolkit.MissingParametersError ODEProblem(sys, u0, tspan, ps)
 @test_nowarn ODEProblem(sys, u0, tspan, [ps..., d => 1.0])
 
 # JET tests

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
 
@@ -1168,3 +1168,17 @@ end
 @named sys = ODESystem(Equation[], t)
 @test getname(unknowns(sys, x)) == :sys₊x
 @test size(unknowns(sys, x)) == size(x)
+
+# 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, Tsit5())(1.0)
+    end
+
+    @test_nowarn ForwardDiff.derivative(P -> x_at_1(P), 1.0)
+end