working codegen

Author: vyudu

src/systems/diffeqs/abstractodesystem.jl

@@ -899,14 +899,14 @@ get_callback(prob::BVProblem) = error("BVP solvers do not support callbacks.")
 """
 function generate_function_bc(sys::ODESystem, u0, u0_idxs, tspan; kwargs...)
     iv = get_iv(sys)
-    sts = get_unknowns(sys)
-    ps = get_ps(sys)
+    sts = unknowns(sys)
+    ps = parameters(sys)
     np = length(ps)
     ns = length(sts)
     stidxmap = Dict([v => i for (i, v) in enumerate(sts)])
     pidxmap = Dict([v => i for (i, v) in enumerate(ps)])
 
-    @variables sol(..)[1:ns] p[1:np]
+    @variables sol(..)[1:ns]
 
     conssys = get_constraintsystem(sys)
     cons = Any[]
@@ -931,7 +931,7 @@ function generate_function_bc(sys::ODESystem, u0, u0_idxs, tspan; kwargs...)
     exprs = vcat(init_conds, cons)
     _p = reorder_parameters(sys, ps)
 
-    build_function_wrapper(sys, exprs, sol, _p..., t; kwargs...)
+    build_function_wrapper(sys, exprs, sol, _p..., t; output_type = Array, kwargs...)
 end
 
 """

src/systems/diffeqs/odesystem.jl

@@ -687,7 +687,6 @@ function process_constraint_system(constraints::Vector{Equation}, sts, ps, iv; c
     constraintsts = OrderedSet()
     constraintps = OrderedSet()
 
-    # Hack? to extract parameters from callable variables in constraints.
     for cons in constraints
         collect_vars!(constraintsts, constraintps, cons, iv)
     end
@@ -712,5 +711,6 @@ function process_constraint_system(constraints::Vector{Equation}, sts, ps, iv; c
         end
     end
 
+    @show constraints
     ConstraintsSystem(constraints, collect(constraintsts), collect(constraintps); name = consname)
 end

src/systems/optimization/constraints_system.jl

@@ -123,7 +123,7 @@ function ConstraintsSystem(constraints, unknowns, ps;
     name === nothing &&
         throw(ArgumentError("The `name` keyword must be provided. Please consider using the `@named` macro"))
 
-    cstr = value.(Symbolics.canonical_form.(scalarize(constraints)))
+    cstr = value.(Symbolics.canonical_form.(vcat(scalarize(constraints)...)))
     unknowns′ = value.(scalarize(unknowns))
     ps′ = value.(ps)
 

test/odesystem.jl

@@ -1670,4 +1670,16 @@ end
     @test_throws ArgumentError @mtkbuild sys = ODESystem(eqs, t; constraints = cons)
     cons = [x(t) * v ~ 3]
     @test_throws ArgumentError @mtkbuild sys = ODESystem(eqs, t; constraints = cons) # Need time argument.
+
+    # Test array variables
+    @variables x(..)[1:5]
+    mat = [1 2 0 3 2
+           0 0 3 2 0
+           0 1 3 0 4
+           2 0 0 2 1
+           0 0 2 0 5]
+    eqs = D(x(t)) ~ mat * x(t)
+    cons = [x(3) ~ [2,3,3,5,4]]
+    @mtkbuild ode = ODESystem(D(x(t)) ~ mat * x(t), t; constraints = cons)
+    @test length(constraints(ModelingToolkit.get_constraintsystem(ode))) == 5
 end