Merge pull request #87 from JuliaDiffEq/hg/refactor/variable

Remove diff field from Variable

Author: ChrisRackauckas

README.md

@@ -152,7 +152,6 @@ context-aware single variable of the IR. Its fields are described as follows:
   the core identifier of the `Variable` in the sense of equality.
 - `subtype`: the main denotation of context. Variables within systems
   are grouped according to their `subtype`.
-- `diff`: the `Differential` object representing the quantity the variable is differentiated with respect to, or `nothing`
 - `dependents`: the vector of variables on which the current variable
   is dependent. For example, `u(t,x)` has dependents `[t,x]`. Derivatives thus
   require this information in order to simplify down.

src/differentials.jl

@@ -11,10 +11,7 @@ Base.convert(::Type{Expr}, D::Differential) = D
 function (D::Differential)(x::Variable)
     D.x === x             && return Constant(1)
     has_dependent(x, D.x) || return Constant(0)
-
-    x′ = copy(x)
-    x′.diff = D
-    return x′
+    return Operation(D, Expression[x])
 end
 Base.:(==)(D1::Differential, D2::Differential) = D1.order == D2.order && D1.x == D2.x
 

src/equations.jl

@@ -12,7 +12,6 @@ Base.:~(lhs::Expression, rhs::Number    ) = Equation(lhs, rhs)
 Base.:~(lhs::Number    , rhs::Expression) = Equation(lhs, rhs)
 
 
-_is_derivative(x::Variable) = x.diff !== nothing
 _is_dependent(x::Variable) = x.subtype === :Unknown && !isempty(x.dependents)
 _is_parameter(ivs) = x -> x.subtype === :Parameter && x ∉ ivs
 _subtype(subtype::Symbol) = x -> x.subtype === subtype

src/systems/diffeqs/diffeqsystem.jl

@@ -4,24 +4,31 @@ mutable struct DiffEqSystem <: AbstractSystem
     dvs::Vector{Variable}
     ps::Vector{Variable}
     jac::Matrix{Expression}
+    function DiffEqSystem(eqs, ivs, dvs, ps, jac)
+        all(!isintermediate, eqs) ||
+            throw(ArgumentError("no intermediate equations permitted in DiffEqSystem"))
+
+        new(eqs, ivs, dvs, ps, jac)
+    end
 end
 
 DiffEqSystem(eqs, ivs, dvs, ps) = DiffEqSystem(eqs, ivs, dvs, ps, Matrix{Expression}(undef,0,0))
 
 function DiffEqSystem(eqs)
-    predicates = [_is_derivative, _is_dependent]
-    _, dvs = extract_elements(eqs, predicates)
+    dvs, = extract_elements(eqs, [_is_dependent])
     ivs = unique(vcat((dv.dependents for dv ∈ dvs)...))
     ps, = extract_elements(eqs, [_is_parameter(ivs)])
     DiffEqSystem(eqs, ivs, dvs, ps, Matrix{Expression}(undef,0,0))
 end
 
 function DiffEqSystem(eqs, ivs)
-    predicates = [_is_derivative, _is_dependent, _is_parameter(ivs)]
-    _, dvs, ps = extract_elements(eqs, predicates)
+    dvs, ps = extract_elements(eqs, [_is_dependent, _is_parameter(ivs)])
     DiffEqSystem(eqs, ivs, dvs, ps, Matrix{Expression}(undef,0,0))
 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)]
     param_exprs = [:($(sys.ps[i].name) = p[$i]) for i in eachindex(sys.ps)]
@@ -44,25 +51,15 @@ function generate_ode_function(sys::DiffEqSystem;version = ArrayFunction)
     end
 end
 
-isintermediate(eq::Equation) = eq.lhs.diff === nothing
-
 function build_equals_expr(eq::Equation)
-    @assert typeof(eq.lhs) <: Variable
-
-    lhs = eq.lhs.name
-    isintermediate(eq) || (lhs = Symbol(lhs, :_, "$(eq.lhs.diff.x.name)"))
+    @assert !isintermediate(eq)
 
+    lhs = Symbol(eq.lhs.args[1].name, :_, eq.lhs.op.x.name)
     return :($lhs = $(convert(Expr, eq.rhs)))
 end
 
 function calculate_jacobian(sys::DiffEqSystem, simplify=true)
-    calcs, diff_exprs = partition(isintermediate, sys.eqs)
-    rhs = [eq.rhs for eq in diff_exprs]
-
-    # Handle intermediate calculations by substitution
-    for calc ∈ calcs
-        find_replace!.(rhs, calc.lhs, calc.rhs)
-    end
+    rhs = [eq.rhs for eq in sys.eqs]
 
     sys_exprs = calculate_jacobian(rhs, sys.dvs)
     sys_exprs = Expression[expand_derivatives(expr) for expr in sys_exprs]
@@ -72,7 +69,6 @@ end
 function generate_ode_jacobian(sys::DiffEqSystem, simplify=true)
     var_exprs = [:($(sys.dvs[i].name) = u[$i]) for i in eachindex(sys.dvs)]
     param_exprs = [:($(sys.ps[i].name) = p[$i]) for i in eachindex(sys.ps)]
-    diff_exprs = filter(!isintermediate, sys.eqs)
     jac = calculate_jacobian(sys, simplify)
     sys.jac = jac
     jac_exprs = [:(J[$i,$j] = $(convert(Expr, jac[i,j]))) for i in 1:size(jac,1), j in 1:size(jac,2)]
@@ -84,7 +80,6 @@ end
 function generate_ode_iW(sys::DiffEqSystem, simplify=true)
     var_exprs = [:($(sys.dvs[i].name) = u[$i]) for i in eachindex(sys.dvs)]
     param_exprs = [:($(sys.ps[i].name) = p[$i]) for i in eachindex(sys.ps)]
-    diff_exprs = filter(!isintermediate, sys.eqs)
     jac = sys.jac
 
     gam = Parameter(:gam)

src/systems/diffeqs/first_order_transform.jl

@@ -1,14 +1,17 @@
-extract_idv(eq::Equation) = eq.lhs.diff.x
+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)
+
+    lower_varname(x, D.x, order, naming_scheme)
 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...)
@@ -19,31 +22,34 @@ function ode_order_lowering(sys::DiffEqSystem; kwargs...)
 end
 ode_order_lowering(eqs; naming_scheme = "_") = ode_order_lowering!(deepcopy(eqs), naming_scheme)
 function ode_order_lowering!(eqs, naming_scheme)
-    ind = findfirst(x->!(isintermediate(x)), eqs)
-    idv = extract_idv(eqs[ind])
+    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}()
+    vars = Variable[]
+    dv_name = eqs[1].lhs.args[1].subtype
+
     for eq in eqs
-        isintermediate(eq) && continue
-        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
+            var ∈ vars || push!(vars, var)
         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 ∈ vars
+        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)
@@ -53,7 +59,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)
@@ -63,12 +69,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

src/variables.jl

@@ -1,22 +1,15 @@
-mutable struct Variable <: Expression
+struct Variable <: Expression
     name::Symbol
     subtype::Symbol
-    diff::Union{Function,Nothing}  # FIXME
     dependents::Vector{Variable}
 end
 
-Variable(name; subtype::Symbol, dependents::Vector{Variable} = Variable[]) =
-    Variable(name, subtype, nothing, dependents)
-
-Parameter(name; kwargs...) = Variable(name; subtype=:Parameter, kwargs...)
-Unknown(name, ;kwargs...) = Variable(name; subtype=:Unknown, kwargs...)
+Parameter(name; dependents = Variable[]) = Variable(name, :Parameter, dependents)
+Unknown(name; dependents = Variable[]) = Variable(name, :Unknown, dependents)
 
 export Variable, Unknown, Parameter, Constant, @Unknown, @Param, @Const
 
 
-Base.copy(x::Variable) = Variable(x.name, x.subtype, x.diff, x.dependents)
-
-
 struct Constant <: Expression
     value::Number
 end
@@ -28,7 +21,7 @@ Base.isone(ex::Expression)  = isa(ex, Constant) && isone(ex.value)
 
 
 # Variables use isequal for equality since == is an Operation
-Base.:(==)(x::Variable, y::Variable) = (x.name, x.subtype, x.diff) == (y.name, y.subtype, y.diff)
+Base.:(==)(x::Variable, y::Variable) = (x.name, x.subtype) == (y.name, y.subtype)
 Base.:(==)(::Variable, ::Number) = false
 Base.:(==)(::Number, ::Variable) = false
 Base.:(==)(::Variable, ::Constant) = false
@@ -37,16 +30,10 @@ Base.:(==)(c::Constant, n::Number) = c.value == n
 Base.:(==)(n::Number, c::Constant) = c.value == n
 Base.:(==)(a::Constant, b::Constant) = a.value == b.value
 
-function Base.convert(::Type{Expr}, x::Variable)
-    x.diff === nothing && return x.name
-    return Symbol("$(x.name)_$(x.diff.x.name)")
-end
+Base.convert(::Type{Expr}, x::Variable) = x.name
 Base.convert(::Type{Expr}, c::Constant) = c.value
 
-function Base.show(io::IO, x::Variable)
-    print(io, x.subtype, '(', x.name, ')')
-    x.diff === nothing || print(io, ", diff = ", x.diff)
-end
+Base.show(io::IO, x::Variable) = print(io, x.subtype, '(', x.name, ')')
 
 # Build variables more easily
 function _parse_vars(macroname, fun, x)

test/system_construction.jl

@@ -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