Reshape data to mirror implicit invariants

Author: HarrisonGrodin

src/equations.jl

@@ -29,10 +29,10 @@ function extract_elements(eqs, predicates)
     return result
 end
 
+get_args(O::Operation) = O.args
+get_args(eq::Equation) = Expression[eq.lhs, eq.rhs]
 function vars!(vars, op)
-    args = isa(op, Equation) ? Expression[op.lhs, op.rhs] : op.args
-
-    for arg ∈ args
+    for arg ∈ get_args(op)
         if isa(arg, Operation)
             vars!(vars, arg)
         elseif isa(arg, Variable)

src/systems/diffeqs/diffeqsystem.jl

@@ -1,16 +1,26 @@
 using Base: RefValue
 
 
+isintermediate(eq::Equation) = !(isa(eq.lhs, Operation) && isa(eq.lhs.op, Differential))
+
+mutable struct DiffEq  # D(x) = t
+    D::Differential  # D
+    var::Variable    # x
+    rhs::Expression  # t
+end
+function Base.convert(::Type{DiffEq}, eq::Equation)
+    isintermediate(eq) && throw(ArgumentError("intermediate equation received"))
+    return DiffEq(eq.lhs.op, eq.lhs.args[1], eq.rhs)
+end
+get_args(eq::DiffEq) = Expression[eq.var, eq.rhs]
+
 struct DiffEqSystem <: AbstractSystem
-    eqs::Vector{Equation}
+    eqs::Vector{DiffEq}
     ivs::Vector{Variable}
     dvs::Vector{Variable}
     ps::Vector{Variable}
     jac::RefValue{Matrix{Expression}}
     function DiffEqSystem(eqs, ivs, dvs, ps)
-        all(!isintermediate, eqs) ||
-            throw(ArgumentError("no intermediate equations permitted in DiffEqSystem"))
-
         jac = RefValue(Matrix{Expression}(undef, 0, 0))
         new(eqs, ivs, dvs, ps, jac)
     end
@@ -28,8 +38,6 @@ function DiffEqSystem(eqs, ivs)
     DiffEqSystem(eqs, ivs, dvs, ps)
 end
 
-isintermediate(eq::Equation) = !(isa(eq.lhs, Operation) && isa(eq.lhs.op, Differential))
-
 
 function generate_ode_function(sys::DiffEqSystem;version = ArrayFunction)
     var_exprs = [:($(sys.dvs[i].name) = u[$i]) for i in eachindex(sys.dvs)]
@@ -53,10 +61,8 @@ function generate_ode_function(sys::DiffEqSystem;version = ArrayFunction)
     end
 end
 
-function build_equals_expr(eq::Equation)
-    @assert !isintermediate(eq)
-
-    lhs = Symbol(eq.lhs.args[1].name, :_, eq.lhs.op.x.name)
+function build_equals_expr(eq::DiffEq)
+    lhs = Symbol(eq.var.name, :_, eq.D.x.name)
     return :($lhs = $(convert(Expr, eq.rhs)))
 end
 

src/systems/diffeqs/first_order_transform.jl

@@ -1,12 +1,8 @@
-extract_idv(eq::Equation) = eq.lhs.op.x
+extract_idv(eq::DiffEq) = eq.D.x
 
-function lower_varname(O::Operation, naming_scheme; lower=false)
-    @assert isa(O.op, Differential)
-
-    D, x = O.op, O.args[1]
+function lower_varname(D::Differential, x, naming_scheme; lower=false)
     order = lower ? D.order-1 : D.order
-
-    lower_varname(x, D.x, order, naming_scheme)
+    return lower_varname(x, D.x, order, naming_scheme)
 end
 function lower_varname(var::Variable, idv, order::Int, naming_scheme)
     sym = var.name
@@ -26,7 +22,7 @@ function ode_order_lowering!(eqs, naming_scheme)
     D   = Differential(idv, 1)
     var_order = Dict{Variable,Int}()
     vars = Variable[]
-    dv_name = eqs[1].lhs.args[1].subtype
+    dv_name = eqs[1].var.subtype
 
     for eq in eqs
         var, maxorder = extract_var_order(eq)
@@ -35,7 +31,7 @@ function ode_order_lowering!(eqs, naming_scheme)
             var_order[var] = maxorder
             var ∈ vars || push!(vars, var)
         end
-        lhs_renaming!(eq, D, naming_scheme)
+        lhs_renaming!(eq, naming_scheme)
         rhs_renaming!(eq, naming_scheme)
     end
 
@@ -52,14 +48,15 @@ function ode_order_lowering!(eqs, naming_scheme)
     return eqs
 end
 
-function lhs_renaming!(eq, D, naming_scheme)
-    eq.lhs = D(lower_varname(eq.lhs, naming_scheme, lower=true))
+function lhs_renaming!(eq::DiffEq, naming_scheme)
+    eq.var = lower_varname(eq.D, eq.var, naming_scheme, lower=true)
     return eq
 end
-rhs_renaming!(eq, naming_scheme) = _rec_renaming!(eq.rhs, naming_scheme)
+rhs_renaming!(eq::DiffEq, naming_scheme) = _rec_renaming!(eq.rhs, naming_scheme)
 
 function _rec_renaming!(rhs, naming_scheme)
-    isa(rhs, Operation) && isa(rhs.op, Differential) && return lower_varname(rhs, naming_scheme)
+    isa(rhs, Operation) && isa(rhs.op, Differential) &&
+        return lower_varname(rhs.op, rhs.args[1], naming_scheme)
     if rhs isa Operation
         args = rhs.args
         for i in eachindex(args)
@@ -69,12 +66,6 @@ function _rec_renaming!(rhs, naming_scheme)
     rhs
 end
 
-function extract_var_order(eq)
-    # We assume that the differential with the highest order is always going to be in the LHS
-    dv = eq.lhs
-    var = dv.args[1]
-    order = dv.op.order
-    return (var, order)
-end
+extract_var_order(eq::DiffEq) = (eq.var, eq.D.order)
 
 export ode_order_lowering

test/system_construction.jl

@@ -44,20 +44,7 @@ lowered_eqs = [D(u_tt) ~ 2u_tt + u_t + x_t + 1
                D(u_t)  ~ u_tt
                D(u)    ~ u_t
                D(x)    ~ x_t]
-function test_eqs(eqs1, eqs2)
-    length(eqs1) == length(eqs2) || return false
-    eq = true
-    for (eq1, eq2) in zip(eqs1, eqs2)
-        lhs1, lhs2 = eq1.lhs, eq2.lhs
-        typeof(lhs1) === typeof(lhs2) || return false
-        for f in fieldnames(typeof(lhs1))
-            eq = eq & isequal(getfield(lhs1, f), getfield(lhs2, f))
-        end
-        eq = eq & isequal(eq1.rhs, eq2.rhs)
-    end
-    eq
-end
-@test test_eqs(de1.eqs, lowered_eqs)
+@test_broken de1.eqs == lowered_eqs
 
 # Internal calculations
 a = y - x