Minor cleanup

Author: HarrisonGrodin

src/differentials.jl

@@ -2,6 +2,7 @@ struct Differential <: Function
     x::Expression
 end
 
+Base.show(io::IO, D::Differential) = print(io, "(D'~", D.x, ")")
 Base.convert(::Type{Expr}, D::Differential) = D
 
 (D::Differential)(x::Operation) = Operation(D, Expression[x])

src/systems/diffeqs/diffeqsystem.jl

@@ -6,12 +6,11 @@ using Base: RefValue
 
 isintermediate(eq::Equation) = !(isa(eq.lhs, Operation) && isa(eq.lhs.op, Differential))
 
-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"))
+function flatten_differential(O::Operation)
+    @assert is_derivative(O) "invalid differential: $O"
+    is_derivative(O.args[1]) || return (O.args[1], O.op.x, 1)
+    (x, t, order) = flatten_differential(O.args[1])
+    t == O.op.x || throw(ArgumentError("non-matching differentials on lhs: $t, $(O.op.x)"))
     return (x, t, order + 1)
 end
 
@@ -24,7 +23,7 @@ struct DiffEq  # dⁿx/dtⁿ = rhs
 end
 function Base.convert(::Type{DiffEq}, eq::Equation)
     isintermediate(eq) && throw(ArgumentError("intermediate equation received"))
-    (x, t, n) = _unwrap_differenital(eq.lhs)
+    (x, t, n) = flatten_differential(eq.lhs)
     return DiffEq(x, t, n, eq.rhs)
 end
 Base.:(==)(a::DiffEq, b::DiffEq) = (a.x, a.t, a.n, a.rhs) == (b.x, b.t, b.n, b.rhs)

src/systems/diffeqs/first_order_transform.jl

@@ -1,4 +1,4 @@
-function lower_varname(var::Variable, idv, order = 0)
+function lower_varname(var::Variable, idv, order)
     order == 0 && return var
     name = Symbol(var.name, :_, string(idv.name)^order)
     return Variable(name, var.known, var.dependents)
@@ -39,8 +39,8 @@ end
 
 function rename(O::Expression)
     isa(O, Operation) || return O
-    if isa(O.op, Differential)
-        (x, t, order) = _unwrap_differenital(O)
+    if is_derivative(O)
+        (x, t, order) = flatten_differential(O)
         return lower_varname(x, t, order)
     end
     return Operation(O.op, rename.(O.args))

src/utils.jl

@@ -60,6 +60,9 @@ is_constant(::Any) = false
 is_operation(::Operation) = true
 is_operation(::Any) = false
 
+is_derivative(O::Operation) = isa(O.op, Differential)
+is_derivative(::Any) = false
+
 has_dependent(t::Variable) = Base.Fix2(has_dependent, t)
 has_dependent(x::Variable, t::Variable) =
     t ∈ x.dependents || any(has_dependent(t), x.dependents)

src/variables.jl

@@ -38,12 +38,7 @@ function Base.convert(::Type{Expr}, x::Variable)
 end
 Base.convert(::Type{Expr}, c::Constant) = c.value
 
-function Base.show(io::IO, x::Variable)
-    subtype = x.known ? :Parameter : :Unknown
-    print(io, subtype, '(', repr(x.name))
-    isempty(x.dependents) || print(io, ", ", x.dependents)
-    print(io, ')')
-end
+Base.show(io::IO, x::Variable) = print(io, x.name)
 
 # Build variables more easily
 function _parse_vars(macroname, fun, x)