Merge pull request #3216 from AayushSabharwal/as/fix-daeprob

fix: fix DAEProblem with array parameters

src/systems/diffeqs/abstractodesystem.jl

@@ -221,9 +221,11 @@ function generate_function(sys::AbstractODESystem, dvs = unknowns(sys),
         pre, sol_states = get_substitutions_and_solved_unknowns(sys)
 
         if implicit_dae
+            # inputs = [] makes `wrap_array_vars` offset by 1 since there is an extra
+            # argument
             build_function(rhss, ddvs, u, p..., t; postprocess_fbody = pre,
                 states = sol_states,
-                wrap_code = wrap_code .∘ wrap_array_vars(sys, rhss; dvs, ps) .∘
+                wrap_code = wrap_code .∘ wrap_array_vars(sys, rhss; dvs, ps, inputs = []) .∘
                             wrap_parameter_dependencies(sys, false),
                 kwargs...)
         else
@@ -790,12 +792,6 @@ function DiffEqBase.ODEProblem{false}(sys::AbstractODESystem, args...; kwargs...
     ODEProblem{false, SciMLBase.FullSpecialize}(sys, args...; kwargs...)
 end
 
-struct DiscreteSaveAffect{F, S} <: Function
-    f::F
-    s::S
-end
-(d::DiscreteSaveAffect)(args...) = d.f(args..., d.s)
-
 function DiffEqBase.ODEProblem{iip, specialize}(sys::AbstractODESystem, u0map = [],
         tspan = get_tspan(sys),
         parammap = DiffEqBase.NullParameters();

src/systems/problem_utils.jl

@@ -641,8 +641,7 @@ function process_SciMLProblem(
         ddvs = map(Differential(iv), dvs)
         du0map = to_varmap(du0map, ddvs)
         merge!(op, du0map)
-
-        du0 = varmap_to_vars(du0map, ddvs; toterm = identity,
+        du0 = varmap_to_vars(op, ddvs; toterm = identity,
             tofloat = true)
         kwargs = merge(kwargs, (; ddvs))
     else

test/odesystem.jl

@@ -1515,3 +1515,12 @@ end
     sol = solve(prob, Tsit5())
     @test sol[obs] ≈ 1:7
 end
+
+@testset "DAEProblem with array parameters" begin
+    @variables x(t)=1.0 y(t) [guess = 1.0]
+    @parameters p[1:2] = [1.0, 2.0]
+    @mtkbuild sys = ODESystem([D(x) ~ x, y^2 ~ x + sum(p)], t)
+    prob = DAEProblem(sys, [D(x) => x, D(y) => D(x) / 2y], [], (0.0, 1.0))
+    sol = solve(prob, DFBDF(), abstol=1e-8, reltol=1e-8)
+    @test sol[x]≈sol[y^2 - sum(p)] atol=1e-5
+end