Merge pull request #767 from SciML/myb/print

Modularize BipartiteGraphs and SystemStructures

Author: YingboMa

Project.toml

@@ -21,6 +21,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 NaNMath = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
 RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
+Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 RuntimeGeneratedFunctions = "7e49a35a-f44a-4d26-94aa-eba1b4ca6b47"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
@@ -49,6 +50,7 @@ LightGraphs = "1.3"
 MacroTools = "0.5"
 NaNMath = "0.3"
 RecursiveArrayTools = "2.3"
+Reexport = "1"
 Requires = "1.0"
 RuntimeGeneratedFunctions = "0.4, 0.5"
 SafeTestsets = "0.0.1"

src/ModelingToolkit.jl

@@ -5,6 +5,7 @@ using StaticArrays, LinearAlgebra, SparseArrays, LabelledArrays
 using Latexify, Unitful, ArrayInterface
 using MacroTools
 using UnPack: @unpack
+using Setfield
 using DiffEqJump
 using DataStructures
 using SpecialFunctions, NaNMath
@@ -202,6 +203,7 @@ Get the set of parameters variables for the given system.
 function parameters end
 
 include("bipartite_graph.jl")
+using .BipartiteGraphs
 
 include("variables.jl")
 include("context_dsl.jl")
@@ -216,7 +218,6 @@ include("domains.jl")
 include("register_function.jl")
 
 include("systems/abstractsystem.jl")
-include("systems/systemstructure.jl")
 
 include("systems/diffeqs/odesystem.jl")
 include("systems/diffeqs/sdesystem.jl")
@@ -239,6 +240,9 @@ include("systems/pde/pdesystem.jl")
 include("systems/reaction/reactionsystem.jl")
 include("systems/dependency_graphs.jl")
 
+include("systems/systemstructure.jl")
+using .SystemStructures
+
 include("systems/reduction.jl")
 
 include("latexify_recipes.jl")

src/bipartite_graph.jl

@@ -1,6 +1,15 @@
+module BipartiteGraphs
+
+export BipartiteEdge, BipartiteGraph
+
+export 𝑠vertices, 𝑑vertices, has_𝑠vertex, has_𝑑vertex, 𝑠neighbors, 𝑑neighbors,
+       𝑠edges, 𝑑edges, nsrcs, ndsts, SRC, DST
+
+using DocStringExtensions
+using Reexport
 using UnPack
 using SparseArrays
-using LightGraphs
+@reexport using LightGraphs
 using Setfield
 
 ###
@@ -229,3 +238,5 @@ function LightGraphs.incidence_matrix(g::BipartiteGraph, val=true)
     end
     S = sparse(I, J, val, nsrcs(g), ndsts(g))
 end
+
+end # module

src/solve.jl

@@ -56,34 +56,37 @@ end
 # return the coefficient matrix `A` and a
 # vector of constants (possibly symbolic) `b` such that
 # A \ b will solve the equations for the vars
-function A_b(eqs::AbstractArray, vars::AbstractArray)
+function A_b(eqs::AbstractArray, vars::AbstractArray, check)
     exprs = rhss(eqs) .- lhss(eqs)
-    for ex in exprs
-        @assert islinear(ex, vars)
+    if check
+        for ex in exprs
+            @assert islinear(ex, vars)
+        end
     end
     A = jacobian(exprs, vars)
     b = A * vars - exprs
     A, b
 end
-function A_b(eq, var)
+function A_b(eq, var, check)
     ex = eq.rhs - eq.lhs
-    @assert islinear(ex, [var])
+    check && @assert islinear(ex, [var])
     a = expand_derivatives(Differential(var)(ex))
     b = a * var - ex
     a, b
 end
 
 """
-    solve_for(eqs::Vector, vars::Vector)
+    solve_for(eqs::Vector, vars::Vector; simplify=true, check=true)
 
 Solve the vector of equations `eqs` for a set of variables `vars`.
 
 Assumes `length(eqs) == length(vars)`
 
-Currently only works if all equations are linear.
+Currently only works if all equations are linear. `check` if the expr is linear
+w.r.t `vars`.
 """
-function solve_for(eqs, vars; simplify=true)
-    A, b = A_b(eqs, vars)
+function solve_for(eqs, vars; simplify=true, check=true)
+    A, b = A_b(eqs, vars, check)
     #TODO: we need to make sure that `solve_for(eqs, vars)` contains no `vars`
     _solve(A, b, simplify)
 end

src/systems/systemstructure.jl

@@ -1,3 +1,11 @@
+module SystemStructures
+
+using ..ModelingToolkit
+import ..ModelingToolkit: isdiffeq, var_from_nested_derivative, vars!, flatten
+using SymbolicUtils: arguments
+using ..BipartiteGraphs
+using UnPack
+using Setfield
 using SparseArrays
 
 #=
@@ -27,10 +35,18 @@ for v in 𝑣vertices(graph); active_𝑣vertices[v] || continue
 end
 =#
 
+export SystemStructure, initialize_system_structure
+export diffvars_range, dervars_range, algvars_range
+export isdiffvar, isdervar, isalgvar, isdiffeq, isalgeq
+export DIFFERENTIAL_VARIABLE, ALGEBRAIC_VARIABLE, DERIVATIVE_VARIABLE
+export DIFFERENTIAL_EQUATION, ALGEBRAIC_EQUATION
+export vartype, eqtype
+
 struct SystemStructure
     dxvar_offset::Int
     fullvars::Vector # [xvar; dxvars; algvars]
     varassoc::Vector{Int}
+    algeqs::BitVector
     graph::BipartiteGraph{Int}
     solvable_graph::BipartiteGraph{Int}
     assign::Vector{Int}
@@ -39,12 +55,35 @@ struct SystemStructure
     partitions::Vector{NTuple{4, Vector{Int}}}
 end
 
+diffvars_range(s::SystemStructure) = 1:s.dxvar_offset
+dervars_range(s::SystemStructure) = s.dxvar_offset+1:2s.dxvar_offset
+algvars_range(s::SystemStructure) = 2s.dxvar_offset+1:length(s.fullvars)
+
+isdiffvar(s::SystemStructure, var::Integer) = var in diffvars_range(s)
+isdervar(s::SystemStructure, var::Integer) = var in dervars_range(s)
+isalgvar(s::SystemStructure, var::Integer) = var in algvars_range(s)
+
+@enum VariableType DIFFERENTIAL_VARIABLE ALGEBRAIC_VARIABLE DERIVATIVE_VARIABLE
+
+function vartype(s::SystemStructure, var::Integer)::VariableType
+    isdiffvar(s, var) ? DIFFERENTIAL_VARIABLE :
+    isdervar(s, var)  ? DERIVATIVE_VARIABLE :
+    isalgvar(s, var)  ? ALGEBRAIC_VARIABLE : error("Variable $var out of bounds")
+end
+
+@enum EquationType DIFFERENTIAL_EQUATION ALGEBRAIC_EQUATION
+
+isalgeq(s::SystemStructure, eq::Integer) = s.algeqs[eq]
+isdiffeq(s::SystemStructure, eq::Integer) = !isalgeq(s, eq)
+eqtype(s::SystemStructure, eq::Integer)::EquationType = isalgeq(s, eq) ? ALGEBRAIC_EQUATION : DIFFERENTIAL_EQUATION
+
 function initialize_system_structure(sys)
-    sys, dxvar_offset, fullvars, varassoc, graph, solvable_graph = init_graph(flatten(sys))
+    sys, dxvar_offset, fullvars, varassoc, algeqs, graph, solvable_graph = init_graph(flatten(sys))
     @set sys.structure = SystemStructure(
                                          dxvar_offset,
                                          fullvars,
                                          varassoc,
+                                         algeqs,
                                          graph,
                                          solvable_graph,
                                          Int[],
@@ -74,8 +113,10 @@ end
 function collect_variables(sys)
     dxvars = []
     eqs = equations(sys)
+    algeqs = trues(length(eqs))
     for (i, eq) in enumerate(eqs)
         if isdiffeq(eq)
+            algeqs[i] = false
             lhs = eq.lhs
             # Make sure that the LHS is a first order derivative of a var.
             @assert !(arguments(lhs)[1] isa Differential) "The equation $eq is not first order"
@@ -86,11 +127,11 @@ function collect_variables(sys)
 
     xvars = (first ∘ var_from_nested_derivative).(dxvars)
     algvars  = setdiff(states(sys), xvars)
-    return xvars, dxvars, algvars
+    return xvars, dxvars, algvars, algeqs
 end
 
 function init_graph(sys)
-    xvars, dxvars, algvars = collect_variables(sys)
+    xvars, dxvars, algvars, algeqs = collect_variables(sys)
     dxvar_offset = length(xvars)
     algvar_offset = 2dxvar_offset
 
@@ -119,5 +160,7 @@ function init_graph(sys)
     end
 
     varassoc = Int[(1:dxvar_offset) .+ dxvar_offset; zeros(Int, length(fullvars) - dxvar_offset)] # variable association list
-    sys, dxvar_offset, fullvars, varassoc, graph, solvable_graph
+    sys, dxvar_offset, fullvars, varassoc, algeqs, graph, solvable_graph
 end
+
+end # module