Merge pull request SciML#2935 from contradict/flatten_for_complex

Flatten for complex

Author: ChrisRackauckas

src/systems/model_parsing.jl

@@ -35,6 +35,12 @@ for f in (:connector, :mtkmodel)
     end
 end
 
+flatten_equations(eqs::Vector{Equation}, eq::Equation) = vcat(eqs, [eq])
+flatten_equations(eq::Vector{Equation}, eqs::Vector{Equation}) = vcat(eq, eqs)
+function flatten_equations(eqs::Vector{Union{Equation, Vector{Equation}}})
+    foldl(flatten_equations, eqs; init = Equation[])
+end
+
 function _model_macro(mod, name, expr, isconnector)
     exprs = Expr(:block)
     dict = Dict{Symbol, Any}(
@@ -56,7 +62,7 @@ function _model_macro(mod, name, expr, isconnector)
     push!(exprs.args, :(variables = []))
     push!(exprs.args, :(parameters = []))
     push!(exprs.args, :(systems = ODESystem[]))
-    push!(exprs.args, :(equations = Equation[]))
+    push!(exprs.args, :(equations = Union{Equation, Vector{Equation}}[]))
     push!(exprs.args, :(defaults = Dict{Num, Union{Number, Symbol, Function}}()))
 
     Base.remove_linenums!(expr)
@@ -106,7 +112,7 @@ function _model_macro(mod, name, expr, isconnector)
     @inline pop_structure_dict!.(
         Ref(dict), [:constants, :defaults, :kwargs, :structural_parameters])
 
-    sys = :($ODESystem($Equation[equations...], $iv, variables, parameters;
+    sys = :($ODESystem($(flatten_equations)(equations), $iv, variables, parameters;
         name, systems, gui_metadata = $gui_metadata, defaults))
 
     if ext[] === nothing

test/complex.jl

@@ -0,0 +1,16 @@
+using ModelingToolkit
+using ModelingToolkit: t_nounits as t
+using Test
+
+@mtkmodel ComplexModel begin
+    @variables begin
+        x(t)
+        y(t)
+        z(t)::Complex
+    end
+    @equations begin
+        z ~ x + im * y
+    end
+end
+@named mixed = ComplexModel()
+@test length(equations(mixed)) == 2

test/runtests.jl

@@ -73,6 +73,7 @@ end
             @safetestset "Initial Values Test" include("initial_values.jl")
             @safetestset "Discrete System" include("discrete_system.jl")
             @safetestset "Equation Type Accessors Test" include("equation_type_accessors.jl")
+            @safetestset "Equations with complex values" include("complex.jl")
         end
     end