Merge pull request #80 from JuliaDiffEq/hg/fix/equation

Update IR for equations

Author: ChrisRackauckas

README.md

@@ -167,7 +167,13 @@ context-aware single variable of the IR. Its fields are described as follows:
 ### Operations
 
 Operations are the basic composition of variables and puts together the pieces
-with a function. The `~` function denotes equality between the arguments.
+with a function.
+
+### Equations
+
+Equations are stored using the `Equation` datatype. Given expressions for the
+left-hand and right-hand sides, an equation is constructed as `Equation(lhs, rhs)`,
+or equivalently `lhs ~ rhs`.
 
 ### Differentials
 

src/ModelingToolkit.jl

@@ -24,6 +24,7 @@ function caclulate_jacobian end
 
 include("operations.jl")
 include("differentials.jl")
+include("equations.jl")
 include("systems/diffeqs/diffeqsystem.jl")
 include("systems/diffeqs/first_order_transform.jl")
 include("systems/nonlinear/nonlinear_system.jl")

src/equations.jl

@@ -0,0 +1,51 @@
+export Equation
+
+
+mutable struct Equation
+    lhs::Expression
+    rhs::Expression
+end
+Base.broadcastable(eq::Equation) = Ref(eq)
+
+Base.:~(lhs::Expression, rhs::Expression) = Equation(lhs, rhs)
+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)
+    end
+    Tuple(elems[target] for target in targets)
+end
+# Walk the tree recursively and push variables into the right set
+function extract_elements!(op, elems, names, targetmap, default)
+    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
+    end
+end

src/function_registration.jl

@@ -32,7 +32,7 @@ for (M, f, arity) in DiffRules.diffrules()
     @eval @register $sig
 end
 
-for fun = (:<, :>, :(==), :~, :!, :&, :|, :div)
+for fun = (:<, :>, :(==), :!, :&, :|, :div)
     basefun = Expr(:., Base, QuoteNode(fun))
     sig = :($basefun(x,y))
     @eval @register $sig

src/operations.jl

@@ -21,13 +21,13 @@ Base.show(io::IO, O::Operation) = print(io, convert(Expr, O))
 
 
 """
-find_replace(O::Operation,x::Variable,y::Expression)
+find_replace(O::Operation, x::Expression, y::Expression)
 
-Finds the variable `x` in Operation `O` and replaces it with the Expression `y`
+Finds the expression `x` in Operation `O` and replaces it with the Expression `y`
 """
-function find_replace!(O::Operation,x::Variable,y::Expression)
+function find_replace!(O::Operation, x::Expression, y::Expression)
     for i in eachindex(O.args)
-        if isequal(O.args[i],x)
+        if isequal(O.args[i], x)
             O.args[i] = y
         elseif typeof(O.args[i]) <: Operation
             find_replace!(O.args[i],x,y)

src/systems/diffeqs/diffeqsystem.jl

@@ -1,5 +1,5 @@
 mutable struct DiffEqSystem <: AbstractSystem
-    eqs::Vector{Operation}
+    eqs::Vector{Equation}
     ivs::Vector{Variable}
     dvs::Vector{Variable}
     vs::Vector{Variable}
@@ -41,71 +41,64 @@ function generate_ode_function(sys::DiffEqSystem;version = ArrayFunction)
     param_exprs = [:($(sys.ps[i].name) = p[$i]) for i in 1:length(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)]
-      exprs = vcat(var_exprs,param_exprs,sys_exprs,dvar_exprs)
-      block = expr_arr_to_block(exprs)
-      :((du,u,p,t)->$(toexpr(block)))
+        dvar_exprs = [:(du[$i] = $(Symbol("$(sys.dvs[i].name)_$(sys.ivs[1].name)"))) for i in 1:length(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)]
-      svector_expr = quote
-        E = eltype(tuple($(dvar_exprs...)))
-        T = StaticArrays.similar_type(typeof(u), E)
-        T($(dvar_exprs...))
-      end
-      exprs = vcat(var_exprs,param_exprs,sys_exprs,svector_expr)
-      block = expr_arr_to_block(exprs)
-      :((u,p,t)->$(toexpr(block)))
+        dvar_exprs = [:($(Symbol("$(sys.dvs[i].name)_$(sys.ivs[1].name)"))) for i in 1:length(sys.dvs)]
+        svector_expr = quote
+            E = eltype(tuple($(dvar_exprs...)))
+            T = StaticArrays.similar_type(typeof(u), E)
+            T($(dvar_exprs...))
+        end
+        exprs = vcat(var_exprs,param_exprs,sys_exprs,svector_expr)
+        block = expr_arr_to_block(exprs)
+        :((u,p,t)->$(toexpr(block)))
     end
 end
 
-isintermediate(eq) = eq.args[1].diff == nothing
+isintermediate(eq::Equation) = eq.lhs.diff === nothing
 
-function build_equals_expr(eq)
-    @assert typeof(eq.args[1]) <: Variable
-    if !(isintermediate(eq))
-        # Differential statement
-        :($(Symbol("$(eq.args[1].name)_$(eq.args[1].diff.x.name)")) = $(eq.args[2]))
-    else
-        # Intermediate calculation
-        :($(Symbol("$(eq.args[1].name)")) = $(eq.args[2]))
-    end
+function build_equals_expr(eq::Equation)
+    @assert typeof(eq.lhs) <: Variable
+
+    lhs = eq.lhs.name
+    isintermediate(eq) || (lhs = Symbol(lhs, :_, "$(eq.lhs.diff.x.name)"))
+
+    return :($lhs = $(convert(Expr, eq.rhs)))
 end
 
-function calculate_jacobian(sys::DiffEqSystem,simplify=true)
-    diff_idxs = map(eq->eq.args[1].diff !=nothing,sys.eqs)
-    diff_exprs = sys.eqs[diff_idxs]
-    rhs = [eq.args[2] for eq in diff_exprs]
+function calculate_jacobian(sys::DiffEqSystem, simplify=true)
+    calcs, diff_exprs = partition(isintermediate, sys.eqs)
+    rhs = [eq.rhs for eq in diff_exprs]
+
     # Handle intermediate calculations by substitution
-    calcs = sys.eqs[.!(diff_idxs)]
-    for i in 1:length(calcs)
-        find_replace!.(rhs,calcs[i].args[1],calcs[i].args[2])
+    for calc ∈ calcs
+        find_replace!.(rhs, calc.lhs, calc.rhs)
     end
-    sys_exprs = calculate_jacobian(rhs,sys.dvs)
+
+    sys_exprs = calculate_jacobian(rhs, sys.dvs)
     sys_exprs = Expression[expand_derivatives(expr) for expr in sys_exprs]
-    if simplify
-        sys_exprs = Expression[simplify_constants(expr) for expr in sys_exprs]
-    end
     sys_exprs
 end
 
-function generate_ode_jacobian(sys::DiffEqSystem,simplify=true)
+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)]
-    diff_idxs = map(eq->eq.args[1].diff !=nothing,sys.eqs)
-    diff_exprs = sys.eqs[diff_idxs]
-    jac = calculate_jacobian(sys,simplify)
+    diff_exprs = filter(!isintermediate, sys.eqs)
+    jac = calculate_jacobian(sys, simplify)
     sys.jac = jac
     jac_exprs = [:(J[$i,$j] = $(convert(Expr, jac[i,j]))) for i in 1:size(jac,1), j in 1:size(jac,2)]
     exprs = vcat(var_exprs,param_exprs,vec(jac_exprs))
     block = expr_arr_to_block(exprs)
     :((J,u,p,t)->$(block))
 end
 
-function generate_ode_iW(sys::DiffEqSystem,simplify=true)
+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)]
-    diff_idxs = map(eq->eq.args[1].diff !=nothing,sys.eqs)
-    diff_exprs = sys.eqs[diff_idxs]
+    diff_exprs = filter(!isintermediate, sys.eqs)
     jac = sys.jac
 
     gam = Variable(:gam)
@@ -144,5 +137,6 @@ function DiffEqBase.ODEFunction(sys::DiffEqSystem;version = ArrayFunction,kwargs
     end
 end
 
+
 export DiffEqSystem, ODEFunction
 export generate_ode_function

src/systems/diffeqs/first_order_transform.jl

@@ -1,3 +1,5 @@
+extract_idv(eq::Equation) = eq.lhs.diff.x
+
 function lower_varname(var::Variable, naming_scheme; lower=false)
     D = var.diff
     D === nothing && return var
@@ -21,7 +23,7 @@ function ode_order_lowering!(eqs, naming_scheme)
     idv = extract_idv(eqs[ind])
     D   = Differential(idv, 1)
     sym_order = Dict{Symbol, Int}()
-    dv_name = eqs[1].args[1].subtype
+    dv_name = eqs[1].lhs.subtype
     for eq in eqs
         isintermediate(eq) && continue
         sym, maxorder = extract_symbol_order(eq)
@@ -37,21 +39,18 @@ function ode_order_lowering!(eqs, naming_scheme)
         for o in (order-1):-1:1
             lhs = D(lower_varname(sym, idv, o-1, dv_name, naming_scheme))
             rhs = lower_varname(sym, idv, o, dv_name, naming_scheme)
-            eq = Operation(==, [lhs, rhs])
+            eq = Equation(lhs, rhs)
             push!(eqs, eq)
         end
     end
     eqs
 end
 
 function lhs_renaming!(eq, D, naming_scheme)
-    eq.args[1] = D(lower_varname(eq.args[1], naming_scheme, lower=true))
+    eq.lhs = D(lower_varname(eq.lhs, naming_scheme, lower=true))
     return eq
 end
-function rhs_renaming!(eq, naming_scheme)
-    rhs = eq.args[2]
-    _rec_renaming!(rhs, naming_scheme)
-end
+rhs_renaming!(eq, naming_scheme) = _rec_renaming!(eq.rhs, naming_scheme)
 
 function _rec_renaming!(rhs, naming_scheme)
     rhs isa Variable && rhs.diff != nothing && return lower_varname(rhs, naming_scheme)
@@ -66,7 +65,7 @@ end
 
 function extract_symbol_order(eq)
     # We assume that the differential with the highest order is always going to be in the LHS
-    dv = eq.args[1]
+    dv = eq.lhs
     sym = dv.name
     order = dv.diff.order
     sym, order

src/systems/nonlinear/nonlinear_system.jl

@@ -1,5 +1,5 @@
 struct NonlinearSystem <: AbstractSystem
-    eqs::Vector{Operation}
+    eqs::Vector{Equation}
     vs::Vector{Variable}
     ps::Vector{Variable}
     v_name::Vector{Symbol}
@@ -26,32 +26,25 @@ end
 function generate_nlsys_function(sys::NonlinearSystem)
     var_exprs = [:($(sys.vs[i].name) = u[$i]) for i in 1:length(sys.vs)]
     param_exprs = [:($(sys.ps[i].name) = p[$i]) for i in 1:length(sys.ps)]
-    sys_idxs = map(eq->isequal(eq.args[1],Constant(0)),sys.eqs)
-    sys_eqs = sys.eqs[sys_idxs]
-    calc_eqs = sys.eqs[.!(sys_idxs)]
-    calc_exprs = [:($(Symbol("$(eq.args[1].name)")) = $(eq.args[2])) for eq in calc_eqs]
-    sys_exprs = [:($(Symbol("resid[$i]")) = $(sys_eqs[i].args[2])) for i in eachindex(sys_eqs)]
+    sys_eqs, calc_eqs = partition(eq -> isequal(eq.lhs, Constant(0)), sys.eqs)
+    calc_exprs = [:($(eq.lhs.name) = $(eq.rhs)) for eq in calc_eqs if isa(eq.lhs, Variable)]
+    sys_exprs = [:($(Symbol("resid[$i]")) = $(sys_eqs[i].rhs)) for i in eachindex(sys_eqs)]
 
     exprs = vcat(var_exprs,param_exprs,calc_exprs,sys_exprs)
     block = expr_arr_to_block(exprs)
     :((du,u,p)->$(block))
 end
 
 function calculate_jacobian(sys::NonlinearSystem,simplify=true)
-    sys_idxs = map(eq->isequal(eq.args[1],Constant(0)),sys.eqs)
-    sys_eqs = sys.eqs[sys_idxs]
-    calc_eqs = sys.eqs[.!(sys_idxs)]
-    rhs = [eq.args[2] for eq in sys_eqs]
+    sys_eqs, calc_eqs = partition(eq -> isequal(eq.lhs, Constant(0)), sys.eqs)
+    rhs = [eq.rhs for eq in sys_eqs]
 
-    for i in 1:length(calc_eqs)
-        find_replace!.(rhs,calc_eqs[i].args[1],calc_eqs[i].args[2])
+    for calc_eq ∈ calc_eqs
+        find_replace!.(rhs, calc_eq.lhs, calc_eq.rhs)
     end
 
     sys_exprs = calculate_jacobian(rhs,sys.vs)
     sys_exprs = Expression[expand_derivatives(expr) for expr in sys_exprs]
-    if simplify
-        sys_exprs = Expression[simplify_constants(expr) for expr in sys_exprs]
-    end
     sys_exprs
 end
 

src/utils.jl

@@ -33,6 +33,12 @@ end
 
 toexpr(ex) = MacroTools.postwalk(x -> isa(x, Expression) ? convert(Expr, x) : x, ex)
 
+function partition(f, xs)
+    idxs = map(f, xs)
+    not_idxs = eachindex(xs) .∉ (idxs,)
+    return (xs[idxs], xs[not_idxs])
+end
+
 is_constant(::Constant) = true
 is_constant(::Any) = false
 

src/variables.jl

@@ -55,44 +55,6 @@ function Base.show(io::IO, x::Variable)
     x.diff === nothing || print(io, ", diff = ", x.diff)
 end
 
-extract_idv(eq) = eq.args[1].diff.x
-
-function extract_elements(ops, 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 op in ops
-        extract_elements!(op, elems, names, targetmap, default)
-    end
-    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, targetmap, default)
-    for arg in op.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
-    end
-end
-
 # Build variables more easily
 function _parse_vars(macroname, fun, x)
     ex = Expr(:block)