Fix first-order transformation

Author: HarrisonGrodin

src/systems/diffeqs/first_order_transform.jl

@@ -1,14 +1,20 @@
 extract_idv(eq::Equation) = eq.lhs.op.x
 
-function lower_varname(var::Variable, naming_scheme; lower=false)
-    D = var.diff
-    D === nothing && return var
+function lower_varname(O::Operation, naming_scheme; lower=false)
+    @assert isa(O.op, Differential)
+
+    D, x = O.op, O.args[1]
     order = lower ? D.order-1 : D.order
-    lower_varname(var.name, D.x, order, var.subtype, naming_scheme)
+
+    sym = x.name
+    name = order == 0 ? sym : Symbol(sym, naming_scheme, string(D.x.name)^order)
+
+    Variable(name, x.subtype, x.dependents)
 end
-function lower_varname(sym::Symbol, idv, order::Int, subtype::Symbol, naming_scheme)
+function lower_varname(var::Variable, idv, order::Int, naming_scheme)
+    sym = var.name
     name = order == 0 ? sym : Symbol(sym, naming_scheme, string(idv.name)^order)
-    return Variable(name, subtype=subtype)
+    return Variable(name, var.subtype, var.dependents)
 end
 
 function ode_order_lowering(sys::DiffEqSystem; kwargs...)
@@ -21,27 +27,30 @@ ode_order_lowering(eqs; naming_scheme = "_") = ode_order_lowering!(deepcopy(eqs)
 function ode_order_lowering!(eqs, naming_scheme)
     idv = extract_idv(eqs[1])
     D   = Differential(idv, 1)
-    sym_order = Dict{Symbol, Int}()
-    dv_name = eqs[1].lhs.subtype
+    var_order = Dict{Variable,Int}()
+    dv_name = eqs[1].lhs.args[1].subtype
+
     for eq in eqs
-        sym, maxorder = extract_symbol_order(eq)
+        var, maxorder = extract_var_order(eq)
         maxorder == 1 && continue # fast pass
-        if maxorder > get(sym_order, sym, 0)
-            sym_order[sym] = maxorder
+        if maxorder > get(var_order, var, 0)
+            var_order[var] = maxorder
         end
-        eq = lhs_renaming!(eq, D, naming_scheme)
-        eq = rhs_renaming!(eq, naming_scheme)
+        lhs_renaming!(eq, D, naming_scheme)
+        rhs_renaming!(eq, naming_scheme)
     end
-    for sym in keys(sym_order)
-        order = sym_order[sym]
+
+    for var ∈ keys(var_order)
+        order = var_order[var]
         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)
+            lhs = D(lower_varname(var, idv, o-1, naming_scheme))
+            rhs = lower_varname(var, idv, o, naming_scheme)
             eq = Equation(lhs, rhs)
             push!(eqs, eq)
         end
     end
-    eqs
+
+    return eqs
 end
 
 function lhs_renaming!(eq, D, naming_scheme)
@@ -51,7 +60,7 @@ 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)
+    isa(rhs, Operation) && isa(rhs.op, Differential) && return lower_varname(rhs, naming_scheme)
     if rhs isa Operation
         args = rhs.args
         for i in eachindex(args)
@@ -61,12 +70,12 @@ function _rec_renaming!(rhs, naming_scheme)
     rhs
 end
 
-function extract_symbol_order(eq)
+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
-    sym = dv.name
-    order = dv.diff.order
-    sym, order
+    var = dv.args[1]
+    order = dv.op.order
+    return (var, order)
 end
 
 export ode_order_lowering

test/system_construction.jl

@@ -38,7 +38,7 @@ test_vars_extraction(de, de2)
 eqs = [D3(u) ~ 2(D2(u)) + D(u) + D(x) + 1
        D2(x) ~ D(x) + 2]
 de = DiffEqSystem(eqs, [t])
-@test_broken de1 = ode_order_lowering(de)
+de1 = ode_order_lowering(de)
 lowered_eqs = [D(u_tt) ~ 2u_tt + u_t + x_t + 1
                D(x_t)  ~ x_t + 2
                D(u_t)  ~ u_tt
@@ -57,7 +57,7 @@ function test_eqs(eqs1, eqs2)
     end
     eq
 end
-@test_broken test_eqs(de1.eqs, lowered_eqs)
+@test test_eqs(de1.eqs, lowered_eqs)
 
 # Internal calculations
 a = y - x