Fix first-order transform

Author: HarrisonGrodin

src/systems/diffeqs/diffeqsystem.jl

@@ -29,33 +29,53 @@ function to_diffeq(eq::Equation)
         throw(ArgumentError("invalid dependent variable $x"))
     return t.op, DiffEq(x.op, n, eq.rhs)
 end
-Base.:(==)(a::DiffEq, b::DiffEq) = isequal((a.x, a.t, a.n, a.rhs), (b.x, b.t, b.n, b.rhs))
+Base.:(==)(a::DiffEq, b::DiffEq) = isequal((a.x, a.n, a.rhs), (b.x, b.n, b.rhs))
 
 struct ODESystem <: AbstractSystem
     eqs::Vector{DiffEq}
     iv::Variable
     dvs::Vector{Variable}
     ps::Vector{Variable}
     jac::RefValue{Matrix{Expression}}
-    function ODESystem(eqs)
-        reformatted = to_diffeq.(eqs)
+end
+
+function ODESystem(eqs)
+    reformatted = to_diffeq.(eqs)
 
-        ivs = unique(r[1] for r ∈ reformatted)
-        length(ivs) == 1 || throw(ArgumentError("one independent variable currently supported"))
-        iv = first(ivs)
+    ivs = unique(r[1] for r ∈ reformatted)
+    length(ivs) == 1 || throw(ArgumentError("one independent variable currently supported"))
+    iv = first(ivs)
 
-        deqs = [r[2] for r ∈ reformatted]
+    deqs = [r[2] for r ∈ reformatted]
 
-        dvs = [deq.x for deq ∈ deqs]
-        ps = filter(vars(deq.rhs for deq ∈ deqs)) do x
-            x.known & !isequal(x, iv)
-        end |> collect
+    dvs = [deq.x for deq ∈ deqs]
+    ps = filter(vars(deq.rhs for deq ∈ deqs)) do x
+        x.known & !isequal(x, iv)
+    end |> collect
 
-        jac = RefValue(Matrix{Expression}(undef, 0, 0))
+    ODESystem(deqs, iv, dvs, ps)
+end
+function ODESystem(deqs, iv, dvs, ps)
+    jac = RefValue(Matrix{Expression}(undef, 0, 0))
+    ODESystem(deqs, iv, dvs, ps, jac)
+end
 
-        new(deqs, iv, dvs, ps, jac)
+function _eq_unordered(a, b)
+    length(a) === length(b) || return false
+    n = length(a)
+    idxs = Set(1:n)
+    for x ∈ a
+        idx = findfirst(isequal(x), b)
+        idx === nothing && return false
+        idx ∈ idxs      || return false
+        delete!(idxs, idx)
     end
+    return true
 end
+Base.:(==)(sys1::ODESystem, sys2::ODESystem) =
+    _eq_unordered(sys1.eqs, sys2.eqs) && isequal(sys1.iv, sys2.iv) &&
+    _eq_unordered(sys1.dvs, sys2.dvs) && _eq_unordered(sys1.ps, sys2.ps)
+# NOTE: equality does not check cached Jacobian
 
 
 function calculate_jacobian(sys::ODESystem)

src/systems/diffeqs/first_order_transform.jl

@@ -8,43 +8,47 @@ function lower_varname(var::Variable, idv, order)
 end
 
 function ode_order_lowering(sys::ODESystem)
-    eqs_lowered = ode_order_lowering(sys.eqs, sys.iv)
-    ODESystem(eqs_lowered, sys.iv, sys.dvs, sys.ps)
+    (eqs_lowered, new_vars) = ode_order_lowering(sys.eqs, sys.iv)
+    ODESystem(eqs_lowered, sys.iv, [sys.dvs; new_vars], sys.ps)
 end
 function ode_order_lowering(eqs, iv)
     var_order = Dict{Variable,Int}()
     vars = Variable[]
     new_eqs = similar(eqs, DiffEq)
+    new_vars = Variable[]
 
     for (i, eq) ∈ enumerate(eqs)
         var, maxorder = eq.x, eq.n
         if maxorder > get(var_order, var, 0)
             var_order[var] = maxorder
             any(isequal(var), vars) || push!(vars, var)
         end
-        var′ = lower_varname(eq.x, eq.t, eq.n - 1)
+        var′ = lower_varname(eq.x, iv, eq.n - 1)
         rhs′ = rename(eq.rhs)
-        new_eqs[i] = DiffEq(var′, iv, 1, rhs′)
+        new_eqs[i] = DiffEq(var′, 1, rhs′)
     end
 
     for var ∈ vars
         order = var_order[var]
         for o in (order-1):-1:1
             lvar = lower_varname(var, iv, o-1)
-            rhs = lower_varname(var, iv, o)
-            eq = DiffEq(lvar, iv, 1, rhs)
+            rvar = lower_varname(var, iv, o)
+            push!(new_vars, rvar)
+
+            rhs = rvar(iv())
+            eq = DiffEq(lvar, 1, rhs)
             push!(new_eqs, eq)
         end
     end
 
-    return new_eqs
+    return (new_eqs, new_vars)
 end
 
 function rename(O::Expression)
     isa(O, Operation) || return O
     if is_derivative(O)
         (x, t, order) = flatten_differential(O)
-        return lower_varname(x, t, order)
+        return lower_varname(x.op, t.op, order)(x.args...)
     end
     return Operation(O.op, rename.(O.args))
 end

test/system_construction.jl

@@ -73,22 +73,20 @@ ModelingToolkit.generate_ode_iW(de)
     end
 end
 
-@test_broken begin
 # Conversion to first-order ODEs #17
 @derivatives D3'''~t
 @derivatives D2''~t
 @variables u(t) u_tt(t) u_t(t) x_t(t)
 eqs = [D3(u) ~ 2(D2(u)) + D(u) + D(x) + 1
        D2(x) ~ D(x) + 2]
-de = ODESystem(eqs, t)
+de = ODESystem(eqs)
 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
                D(u)    ~ u_t
                D(x)    ~ x_t]
-@test de1.eqs == convert.(ModelingToolkit.DiffEq, lowered_eqs)
-end
+@test de1 == ODESystem(lowered_eqs)
 
 # Internal calculations
 a = y - x