Fix DiffEqSystem

Author: HarrisonGrodin

src/equations.jl

@@ -12,40 +12,31 @@ Base.:~(lhs::Expression, rhs::Number    ) = Equation(lhs, rhs)
 Base.:~(lhs::Number    , rhs::Expression) = Equation(lhs, rhs)
 
 
-function extract_elements(eqs, targetmap, default = nothing)
-    elems = Dict{Symbol, Vector{Variable}}()
-    names = Dict{Symbol, Set{Symbol}}()
-    if default == nothing
-        targets =  unique(collect(values(targetmap)))
-    else
-        targets = [unique(collect(values(targetmap))), default]
-    end
-    for target in targets
-        elems[target] = Vector{Variable}()
-        names[target] = Set{Symbol}()
-    end
-    for eq in eqs
-        extract_elements!(eq, elems, names, targetmap, default)
+_is_derivative(x::Variable) = x.diff !== nothing
+_is_dependent(x::Variable) = x.subtype === :DependentVariable && !isempty(x.dependents)
+_subtype(subtype::Symbol) = x -> x.subtype === subtype
+
+function extract_elements(eqs, predicates)
+    result = [Variable[] for p ∈ predicates]
+    vars = foldl(vars!, eqs; init=Set{Variable}())
+
+    for var ∈ vars
+        for (i, p) ∈ enumerate(predicates)
+            p(var) && (push!(result[i], var); break)
+        end
     end
-    Tuple(elems[target] for target in targets)
+
+    return result
 end
-# Walk the tree recursively and push variables into the right set
-function extract_elements!(op, elems, names, targetmap, default)
+
+function vars!(vars, op)
     args = isa(op, Equation) ? Expression[op.lhs, op.rhs] : op.args
 
-    for arg in args
-        if arg isa Operation
-            extract_elements!(arg, elems, names, targetmap, default)
-        elseif arg isa Variable
-            if default == nothing
-                target = haskey(targetmap, arg.subtype) ? targetmap[arg.subtype] : continue
-            else
-                target = haskey(targetmap, arg.subtype) ? targetmap[arg.subtype] : default
-            end
-            if !in(arg.name, names[target])
-                push!(names[target], arg.name)
-                push!(elems[target], arg)
-            end
-        end
+    for arg ∈ args
+        isa(arg, Operation) ? vars!(vars, arg) :
+        isa(arg, Variable)  ? push!(vars, arg) :
+        nothing
     end
+
+    return vars
 end

src/systems/diffeqs/diffeqsystem.jl

@@ -19,34 +19,31 @@ end
 
 function DiffEqSystem(eqs; iv_name = :IndependentVariable,
                            dv_name = :DependentVariable,
-                           v_name = :Variable,
                            p_name = :Parameter)
-    targetmap =  Dict(iv_name => iv_name, dv_name => dv_name, v_name => v_name,
-                       p_name => p_name)
-    ivs, dvs, vs, ps = extract_elements(eqs, targetmap)
-    DiffEqSystem(eqs, ivs, dvs, vs, ps, iv_name, dv_name, p_name, Matrix{Expression}(0,0))
+    predicates = [_is_derivative, _subtype(iv_name), _is_dependent, _subtype(dv_name), _subtype(p_name)]
+    _, ivs, dvs, vs, ps = extract_elements(eqs, predicates)
+    DiffEqSystem(eqs, ivs, dvs, vs, ps, iv_name, dv_name, p_name, Matrix{Expression}(undef,0,0))
 end
 
 function DiffEqSystem(eqs, ivs;
                       dv_name = :DependentVariable,
-                      v_name = :Variable,
                       p_name = :Parameter)
-    targetmap =  Dict(dv_name => dv_name, v_name => v_name, p_name => p_name)
-    dvs, vs, ps = extract_elements(eqs, targetmap)
+    predicates = [_is_derivative, _is_dependent, _subtype(dv_name), _subtype(p_name)]
+    _, dvs, vs, ps = extract_elements(eqs, predicates)
     DiffEqSystem(eqs, ivs, dvs, vs, ps, ivs[1].subtype, dv_name, p_name, Matrix{Expression}(undef,0,0))
 end
 
 function generate_ode_function(sys::DiffEqSystem;version = ArrayFunction)
-    var_exprs = [:($(sys.dvs[i].name) = u[$i]) for i in 1:length(sys.dvs)]
-    param_exprs = [:($(sys.ps[i].name) = p[$i]) for i in 1:length(sys.ps)]
+    var_exprs = [:($(sys.dvs[i].name) = u[$i]) for i in eachindex(sys.dvs)]
+    param_exprs = [:($(sys.ps[i].name) = p[$i]) for i in eachindex(sys.ps)]
     sys_exprs = build_equals_expr.(sys.eqs)
     if version == ArrayFunction
-        dvar_exprs = [:(du[$i] = $(Symbol("$(sys.dvs[i].name)_$(sys.ivs[1].name)"))) for i in 1:length(sys.dvs)]
+        dvar_exprs = [:(du[$i] = $(Symbol("$(sys.dvs[i].name)_$(sys.ivs[1].name)"))) for i in eachindex(sys.dvs)]
         exprs = vcat(var_exprs,param_exprs,sys_exprs,dvar_exprs)
         block = expr_arr_to_block(exprs)
         :((du,u,p,t)->$(toexpr(block)))
     elseif version == SArrayFunction
-        dvar_exprs = [:($(Symbol("$(sys.dvs[i].name)_$(sys.ivs[1].name)"))) for i in 1:length(sys.dvs)]
+        dvar_exprs = [:($(Symbol("$(sys.dvs[i].name)_$(sys.ivs[1].name)"))) for i in eachindex(sys.dvs)]
         svector_expr = quote
             E = eltype(tuple($(dvar_exprs...)))
             T = StaticArrays.similar_type(typeof(u), E)
@@ -84,8 +81,8 @@ function calculate_jacobian(sys::DiffEqSystem, simplify=true)
 end
 
 function generate_ode_jacobian(sys::DiffEqSystem, simplify=true)
-    var_exprs = [:($(sys.dvs[i].name) = u[$i]) for i in 1:length(sys.dvs)]
-    param_exprs = [:($(sys.ps[i].name) = p[$i]) for i in 1:length(sys.ps)]
+    var_exprs = [:($(sys.dvs[i].name) = u[$i]) for i in eachindex(sys.dvs)]
+    param_exprs = [:($(sys.ps[i].name) = p[$i]) for i in eachindex(sys.ps)]
     diff_exprs = filter(!isintermediate, sys.eqs)
     jac = calculate_jacobian(sys, simplify)
     sys.jac = jac
@@ -96,8 +93,8 @@ function generate_ode_jacobian(sys::DiffEqSystem, simplify=true)
 end
 
 function generate_ode_iW(sys::DiffEqSystem, simplify=true)
-    var_exprs = [:($(sys.dvs[i].name) = u[$i]) for i in 1:length(sys.dvs)]
-    param_exprs = [:($(sys.ps[i].name) = p[$i]) for i in 1:length(sys.ps)]
+    var_exprs = [:($(sys.dvs[i].name) = u[$i]) for i in eachindex(sys.dvs)]
+    param_exprs = [:($(sys.ps[i].name) = p[$i]) for i in eachindex(sys.ps)]
     diff_exprs = filter(!isintermediate, sys.eqs)
     jac = sys.jac
 

src/systems/nonlinear/nonlinear_system.jl

@@ -15,8 +15,7 @@ end
 function NonlinearSystem(eqs;
                          v_name = :DependentVariable,
                          p_name = :Parameter)
-    targetmap = Dict(v_name => v_name, p_name => p_name)
-    vs, ps = extract_elements(eqs, targetmap)
+    vs, ps = extract_elements(eqs, [_subtype(v_name), _subtype(p_name)])
     NonlinearSystem(eqs, vs, ps, [v_name], p_name)
 end
 

test/system_construction.jl

@@ -20,7 +20,7 @@ jac = ModelingToolkit.calculate_jacobian(de)
 f = ODEFunction(de)
 ModelingToolkit.generate_ode_iW(de)
 
-# Differential equation with automatic extraction of variables on rhs
+# Differential equation with automatic extraction of variables
 de2 = DiffEqSystem(eqs, [t])
 
 function test_vars_extraction(de, de2)