Refactor differentials

Remove `order`, fixes higher-order derivatives.

Author: HarrisonGrodin

src/differentials.jl

@@ -1,10 +1,7 @@
 struct Differential <: Function
     x::Expression
-    order::Int
 end
-Differential(x) = Differential(x,1)
 
-Base.show(io::IO, D::Differential) = print(io,"($(D.x),$(D.order))")
 Base.convert(::Type{Expr}, D::Differential) = D
 
 (D::Differential)(x::Operation) = Operation(D, Expression[x])
@@ -13,8 +10,8 @@ function (D::Differential)(x::Variable)
     has_dependent(x, D.x) || return Constant(0)
     return Operation(D, Expression[x])
 end
-(::Differential)(::Constant) = Constant(0)
-Base.:(==)(D1::Differential, D2::Differential) = D1.order == D2.order && D1.x == D2.x
+(::Differential)(::Any) = Constant(0)
+Base.:(==)(D1::Differential, D2::Differential) = D1.x == D2.x
 
 function expand_derivatives(O::Operation)
     @. O.args = expand_derivatives(O.args)
@@ -56,6 +53,7 @@ function count_order(x)
     n, x.args[1]
 end
 
+_repeat_apply(f, n) = n == 1 ? f : f ∘ _repeat_apply(f, n-1)
 function _differential_macro(x)
     ex = Expr(:block)
     lhss = Symbol[]
@@ -66,7 +64,7 @@ function _differential_macro(x)
         rhs = di.args[3]
         order, lhs = count_order(lhs)
         push!(lhss, lhs)
-        expr = :($lhs = Differential($rhs, $order))
+        expr = :($lhs = $_repeat_apply(Differential($rhs), $order))
         push!(ex.args,  expr)
     end
     push!(ex.args, Expr(:tuple, lhss...))

src/systems/diffeqs/diffeqsystem.jl

@@ -6,17 +6,29 @@ using Base: RefValue
 
 isintermediate(eq::Equation) = !(isa(eq.lhs, Operation) && isa(eq.lhs.op, Differential))
 
-struct DiffEq  # D(x) = t
-    D::Differential  # D
-    var::Variable    # x
-    rhs::Expression  # t
+function _unwrap_differenital(O)
+    isa(O, Operation)       || return (O, nothing, 0)
+    isa(O.op, Differential) || return (O, nothing, 0)
+    (x, t, order) = _unwrap_differenital(O.args[1])
+    t === nothing && (t = O.op.x)
+    t == O.op.x || throw(ArgumentError("non-matching differentials on lhs"))
+    return (x, t, order + 1)
+end
+
+
+struct DiffEq  # dⁿx/dtⁿ = rhs
+    x::Expression
+    t::Variable
+    n::Int
+    rhs::Expression
 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)
+    (x, t, n) = _unwrap_differenital(eq.lhs)
+    return DiffEq(x, t, n, eq.rhs)
 end
-Base.:(==)(a::DiffEq, b::DiffEq) = (a.D, a.var, a.rhs) == (b.D, b.var, b.rhs)
-get_args(eq::DiffEq) = Expression[eq.var, eq.rhs]
+Base.:(==)(a::DiffEq, b::DiffEq) = (a.x, a.t, a.n, a.rhs) == (b.x, b.t, b.n, b.rhs)
+get_args(eq::DiffEq) = Expression[eq.x, eq.t, eq.rhs]
 
 struct DiffEqSystem <: AbstractSystem
     eqs::Vector{DiffEq}

src/systems/diffeqs/first_order_transform.jl

@@ -1,41 +1,35 @@
-function lower_varname(D::Differential, x; lower=false)
-    order = lower ? D.order-1 : D.order
-    return lower_varname(x, D.x, order)
-end
-function lower_varname(var::Variable, idv, order::Int)
-    sym = var.name
-    name = order == 0 ? sym : Symbol(sym, :_, string(idv.name)^order)
+function lower_varname(var::Variable, idv, order = 0)
+    order == 0 && return var
+    name = Symbol(var.name, :_, string(idv.name)^order)
     return Variable(name, var.known, var.dependents)
 end
 
 function ode_order_lowering(sys::DiffEqSystem)
     eqs_lowered = ode_order_lowering(sys.eqs, sys.iv)
-    DiffEqSystem(eqs_lowered, sys.iv)
+    DiffEqSystem(eqs_lowered, sys.iv, sys.dvs, sys.ps)
 end
 function ode_order_lowering(eqs, iv)
-    D = Differential(iv, 1)
     var_order = Dict{Variable,Int}()
     vars = Variable[]
     new_eqs = similar(eqs, DiffEq)
 
     for (i, eq) ∈ enumerate(eqs)
-        var, maxorder = eq.var, eq.D.order
-        maxorder == 1 && continue # fast pass
+        var, maxorder = eq.x, eq.n
         if maxorder > get(var_order, var, 0)
             var_order[var] = maxorder
             var ∈ vars || push!(vars, var)
         end
-        var′ = lower_varname(eq.D, eq.var, lower = true)
+        var′ = lower_varname(eq.x, eq.t, eq.n - 1)
         rhs′ = rename(eq.rhs)
-        new_eqs[i] = DiffEq(D, var′, rhs′)
+        new_eqs[i] = DiffEq(var′, iv, 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(D, lvar, rhs)
+            eq = DiffEq(lvar, iv, 1, rhs)
             push!(new_eqs, eq)
         end
     end
@@ -45,7 +39,10 @@ end
 
 function rename(O::Expression)
     isa(O, Operation) || return O
-    isa(O.op, Differential) && return lower_varname(O.op, O.args[1])
+    if isa(O.op, Differential)
+        (x, t, order) = _unwrap_differenital(O)
+        return lower_varname(x, t, order)
+    end
     return Operation(O.op, rename.(O.args))
 end
 

test/variable_parsing.jl

@@ -27,7 +27,7 @@ s1 = Parameter(:s)
 @test convert(Expr, s) == :s
 @test convert(Expr, cos(t + sin(s))) == :(cos(t + sin(s)))
 
-@Deriv D''~t
-D1 = Differential(t, 2)
+@Deriv D'~t
+D1 = Differential(t)
 @test D1 == D
 @test convert(Expr, D) == D