Map subtypes in extract_elements

Specify mapping of original subtypes to final subtypes.
Allows a default mapping for subtypes not included in `targetmap`. If `default == nothing`, subtypes not included in `targetmap` are ignored.

Author: AleMorales

src/systems/diffeqs/diffeqsystem.jl

@@ -20,15 +20,18 @@ function DiffEqSystem(eqs; iv_name = :IndependentVariable,
                            dv_name = :DependentVariable,
                            v_name = :Variable,
                            p_name = :Parameter)
-    ivs, dvs, vs, ps = extract_elements(eqs, (iv_name, dv_name, v_name, p_name))
+    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)
 end
 
 function DiffEqSystem(eqs, ivs;
                       dv_name = :DependentVariable,
                       v_name = :Variable,
                       p_name = :Parameter)
-    dvs, vs, ps = extract_elements(eqs, (dv_name, v_name, p_name))
+    targetmap =  Dict(dv_name => dv_name, v_name => v_name, p_name => p_name)
+    dvs, vs, ps = extract_elements(eqs, targetmap)
     DiffEqSystem(eqs, ivs, dvs, vs, ps, ivs[1].subtype, dv_name, p_name)
 end
 

src/systems/nonlinear/nonlinear_system.jl

@@ -18,8 +18,8 @@ function NonlinearSystem(eqs;
                          dv_name = :DependentVariable,
                          p_name = :Parameter)
     # Allow the use of :DependentVariable to make it seamless with DE use
-    dvs, vs, ps = extract_elements(eqs, (dv_name, v_name, p_name))
-    vs = [dvs;vs]
+    targetmap = Dict(v_name => v_name, dv_name => v_name, p_name => p_name)
+    vs, ps = extract_elements(eqs, targetmap)
     NonlinearSystem(eqs, vs, ps, [v_name,dv_name], p_name)
 end
 

src/variables.jl

@@ -100,26 +100,38 @@ end
 
 extract_idv(eq) = eq.args[1].diff.x
 
-function extract_elements(ops, eltypes)
+function extract_elements(ops, targetmap, default = nothing)
     elems = Dict{Symbol, Vector{Variable}}()
     names = Dict{Symbol, Set{Symbol}}()
-    for el in eltypes
-        elems[el] = Vector{Variable}()
-        names[el] = 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 op in ops
-        extract_elements!(op, elems, names)
+        extract_elements!(op, elems, names, targetmap, default)
     end
-    Tuple(elems[el] for el in eltypes)
+    Tuple(elems[target] for target in targets)
 end
 # Walk the tree recursively and push variables into the right set
-function extract_elements!(op::AbstractOperation, elems, names)
+function extract_elements!(op::AbstractOperation, elems, names, targetmap, default)
     for arg in op.args
         if arg isa Operation
-            extract_elements!(arg, elems, names)
-        elseif arg isa Variable && haskey(elems, arg.subtype) && !in(arg.name, names[arg.subtype])
-            push!(names[arg.subtype], arg.name)
-            push!(elems[arg.subtype], arg)
+                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
     end
 end