Handle inhomogeneous parameters using a Tuple of Vectors

src/parameters.jl

@@ -61,3 +61,29 @@
         xs,
         toparam) |> esc
 end
+
+function split_parameters_by_type(ps)
+    by = let set = Dict{Any, Int}(), counter = Ref(1)
+        x -> begin
+            t = typeof(x)
+            get!(set, typeof(x)) do
+                if t == Float64
+                    1
+                else
+                    counter[] += 1
+                end
+            end
+        end
+    end
+    idxs = by.(ps)
+    split_idxs = [Int[]]
+    for (i, idx) in enumerate(idxs)
+        if idx > length(split_idxs)
+            push!(split_idxs, Int[])
+        end
+        push!(split_idxs[idx], i)
+    end
+    tighten_types = x -> identity.(x)
+    split_ps = tighten_types.(Base.Fix1(getindex, ps).(split_idxs))
+    (split_ps...,), split_idxs
+end

src/systems/diffeqs/abstractodesystem.jl

@@ -122,6 +122,7 @@
            nothing,
     isdde = false,
     has_difference = false,
+    split_parameters = false,
     kwargs...)
     if isdde
         eqs = delay_to_function(sys)
@@ -151,6 +152,9 @@
             build_function(rhss, ddvs, u, p, t; postprocess_fbody = pre,
                 states = sol_states,
                 kwargs...)
+        elseif split_parameters
+            build_function(rhss, u, p..., t; postprocess_fbody = pre, states = sol_states,
+                kwargs...)
         else
             build_function(rhss, u, p, t; postprocess_fbody = pre, states = sol_states,
                 kwargs...)
@@ -325,15 +329,23 @@
     checkbounds = false,
     sparsity = false,
     analytic = nothing,
+    split_parameters = false,
     kwargs...) where {iip, specialize}
     f_gen = generate_function(sys, dvs, ps; expression = Val{eval_expression},
-        expression_module = eval_module, checkbounds = checkbounds,
+        expression_module = eval_module, checkbounds = checkbounds, split_parameters,
         kwargs...)
     f_oop, f_iip = eval_expression ?
                    (drop_expr(@RuntimeGeneratedFunction(eval_module, ex)) for ex in f_gen) :
                    f_gen
-    f(u, p, t) = f_oop(u, p, t)
-    f(du, u, p, t) = f_iip(du, u, p, t)
+    if split_parameters
+        g(u, p, t) = f_oop(u, p..., t)
+        g(du, u, p, t) = f_iip(du, u, p..., t)
+        f = g
+    else
+        k(u, p, t) = f_oop(u, p, t)
+        k(du, u, p, t) = f_iip(du, u, p, t)
+        f = k
+    end
 
     if specialize === SciMLBase.FunctionWrapperSpecialize && iip
         if u0 === nothing || p === nothing || t === nothing
@@ -686,6 +698,7 @@
     parammap;
     use_union = false,
     tofloat = !use_union,
+    split_parameters = false,
     symbolic_u0 = false)
     eqs = equations(sys)
     dvs = states(sys)
@@ -698,7 +711,7 @@
     if symbolic_u0
         u0 = varmap_to_vars(u0map, dvs; defaults = defs, tofloat = false, use_union = false)
     else
-        u0 = varmap_to_vars(u0map, dvs; defaults = defs, tofloat = true)
+        u0 = varmap_to_vars(u0map, dvs; defaults = defs, tofloat = !split_parameters)
     end
     p = varmap_to_vars(parammap, ps; defaults = defs, tofloat, use_union)
     p = p === nothing ? SciMLBase.NullParameters() : p
@@ -716,13 +729,24 @@
     use_union = false,
     tofloat = !use_union,
     symbolic_u0 = false,
+    split_parameters = false,
     kwargs...)
     eqs = equations(sys)
     dvs = states(sys)
     ps = parameters(sys)
     iv = get_iv(sys)
 
-    u0, p, defs = get_u0_p(sys, u0map, parammap; tofloat, use_union, symbolic_u0)
+    u0, p, defs = get_u0_p(sys,
+        u0map,
+        parammap;
+        tofloat,
+        use_union,
+        symbolic_u0,
+        split_parameters)
+    if split_parameters
+        p, split_idxs = split_parameters_by_type(p)
+        ps = Base.Fix1(getindex, parameters(sys)).(split_idxs)
+    end
 
     if implicit_dae && du0map !== nothing
         ddvs = map(Differential(iv), dvs)
@@ -739,7 +763,8 @@
     f = constructor(sys, dvs, ps, u0; ddvs = ddvs, tgrad = tgrad, jac = jac,
         checkbounds = checkbounds, p = p,
         linenumbers = linenumbers, parallel = parallel, simplify = simplify,
-        sparse = sparse, eval_expression = eval_expression, kwargs...)
+        sparse = sparse, eval_expression = eval_expression, split_parameters,
+        kwargs...)
     implicit_dae ? (f, du0, u0, p) : (f, u0, p)
 end
 

src/utils.jl

@@ -243,6 +243,10 @@
 
 hasdefault(v) = hasmetadata(v, Symbolics.VariableDefaultValue)
 getdefault(v) = value(getmetadata(v, Symbolics.VariableDefaultValue))
+function getdefaulttype(v)
+    def = value(getmetadata(unwrap(v), Symbolics.VariableDefaultValue, nothing))
+    def === nothing ? Float64 : typeof(def)
+end
 function setdefault(v, val)
     val === nothing ? v : setmetadata(v, Symbolics.VariableDefaultValue, value(val))
 end

test/runtests.jl

@@ -23,6 +23,7 @@ using SafeTestsets, Test
 @safetestset "JumpSystem Test" include("jumpsystem.jl")
 @safetestset "Constraints Test" include("constraints.jl")
 @safetestset "Reduction Test" include("reduction.jl")
+@safetestset "Split Parameters Test" include("split_parameters.jl")
 @safetestset "ODAEProblem Test" include("odaeproblem.jl")
 @safetestset "Components Test" include("components.jl")
 @safetestset "Model Parsing Test" include("model_parsing.jl")

test/split_parameters.jl

@@ -0,0 +1,31 @@
+using ModelingToolkit, Test
+using ModelingToolkitStandardLibrary.Blocks
+using OrdinaryDiffEq
+
+dt = 4e-4
+t_end = 10.0
+time = 0:dt:t_end
+x = @. time^2 + 1.0
+
+@parameters t
+D = Differential(t)
+
+vars = @variables y(t)=1 dy(t)=0 ddy(t)=0
+@named src = SampledData(; data = Float64[], dt)
+@named int = Integrator()
+
+eqs = [y ~ src.output.u
+    D(y) ~ dy
+    D(dy) ~ ddy
+    connect(src.output, int.input)]
+
+@named sys = ODESystem(eqs, t, vars, []; systems = [int, src])
+s = complete(sys)
+sys = structural_simplify(sys)
+
+prob = ODEProblem(sys, [], (0.0, t_end), [s.src.data => x]; split_parameters = true)
+@test prob.p isa Tuple{Vector{Float64}, Vector{Int}, Vector{Vector{Float64}}}
+@time sol = solve(prob, ImplicitEuler());
+prob2 = ODEProblem(sys, [], (0.0, t_end), [s.src.data => x]; to_float = false)
+@test prob2.p isa Vector{Union{Float64, Int64, Vector{Float64}}}
+@time sol2 = solve(prob2, ImplicitEuler());