Refactor variables as functions

Author: HarrisonGrodin

src/differentials.jl

@@ -4,27 +4,27 @@ export Differential, expand_derivatives, @derivatives
 struct Differential <: Function
     x::Expression
 end
+(D::Differential)(x) = Operation(D, Expression[x])
 
 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])
-function (D::Differential)(x::Variable)
-    D.x === x             && return Constant(1)
-    has_dependent(x, D.x) || return Constant(0)
-    return Operation(D, Expression[x])
-end
-(::Differential)(::Any) = Constant(0)
 Base.:(==)(D1::Differential, D2::Differential) = isequal(D1.x, D2.x)
 
 function expand_derivatives(O::Operation)
     @. O.args = expand_derivatives(O.args)
 
-    if O.op isa Differential
-        D = O.op
-        o = O.args[1]
-        isa(o, Operation) || return O
-        return simplify_constants(sum(i->derivative(o,i)*expand_derivatives(D(o.args[i])),1:length(o.args)))
+    if isa(O.op, Differential)
+        (D, o) = (O.op, O.args[1])
+
+        isequal(o, D.x)     && return Constant(1)
+        occursin(D.x, o)    || return Constant(0)
+        isa(o, Operation)   || return O
+        isa(o.op, Variable) && return O
+
+        return sum(1:length(o.args)) do i
+            derivative(o, i) * expand_derivatives(D(o.args[i]))
+        end |> simplify_constants
     end
 
     return O
@@ -79,6 +79,6 @@ macro derivatives(x...)
     esc(_differential_macro(x))
 end
 
-function calculate_jacobian(eqs,vars)
-    Expression[Differential(vars[j])(eqs[i]) for i in 1:length(eqs), j in 1:length(vars)]
+function calculate_jacobian(eqs, dvs, iv)
+    Expression[Differential(dv(iv()))(eq) for eq ∈ eqs, dv ∈ dvs]
 end

src/equations.jl

@@ -11,11 +11,9 @@ Base.:~(lhs::Expression, rhs::Expression) = Equation(lhs, rhs)
 Base.:~(lhs::Expression, rhs::Number    ) = Equation(lhs, rhs)
 Base.:~(lhs::Number    , rhs::Expression) = Equation(lhs, rhs)
 
-
-_is_dependent(x::Variable) = !x.known && !isempty(x.dependents)
-_is_parameter(iv) = x -> x.known && !isequal(x, iv)
-_is_known(x::Variable) = x.known
-_is_unknown(x::Variable) = !x.known
+_is_parameter(iv) = (O::Operation) -> O.op.known && !isequal(O, iv)
+_is_known(O::Operation) = O.op.known
+_is_unknown(O::Operation) = !O.op.known
 
 function extract_elements(eqs, predicates)
     result = [Variable[] for p ∈ predicates]
@@ -32,15 +30,13 @@ end
 
 get_args(O::Operation) = O.args
 get_args(eq::Equation) = Expression[eq.lhs, eq.rhs]
-function vars!(vars, op)
-    for arg ∈ get_args(op)
+vars(exprs) = foldl(vars!, exprs; init = Set{Variable}())
+function vars!(vars, O)
+    isa(O, Operation) || return vars
+    for arg ∈ O.args
         if isa(arg, Operation)
+            isa(arg.op, Variable) && push!(vars, arg.op)
             vars!(vars, arg)
-        elseif isa(arg, Variable)
-            push!(vars, arg)
-            for dep ∈ arg.dependents
-                push!(vars, dep)
-            end
         end
     end
 

src/operations.jl

@@ -3,10 +3,8 @@ struct Operation <: Expression
     args::Vector{Expression}
 end
 
-# Recursive ==
-function Base.isequal(x::Operation,y::Operation)
+Base.isequal(x::Operation,y::Operation) =
     x.op == y.op && length(x.args) == length(y.args) && all(isequal.(x.args,y.args))
-end
 Base.isequal(::Operation, ::Number   ) = false
 Base.isequal(::Number   , ::Operation) = false
 Base.isequal(::Operation, ::Variable ) = false

src/systems/diffeqs/diffeqsystem.jl

@@ -16,51 +16,53 @@ end
 
 
 struct DiffEq  # dⁿx/dtⁿ = rhs
-    x::Expression
-    t::Variable
+    x::Variable
     n::Int
     rhs::Expression
 end
-function Base.convert(::Type{DiffEq}, eq::Equation)
+function to_diffeq(eq::Equation)
     isintermediate(eq) && throw(ArgumentError("intermediate equation received"))
     (x, t, n) = flatten_differential(eq.lhs)
-    return DiffEq(x, t, n, eq.rhs)
+    (isa(t, Operation) && isa(t.op, Variable) && isempty(t.args)) ||
+        throw(ArgumentError("invalid independent variable $t"))
+    (isa(x, Operation) && isa(x.op, Variable) && length(x.args) == 1 && isequal(first(x.args), t)) ||
+        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))
-get_args(eq::DiffEq) = Expression[eq.x, eq.t, eq.rhs]
 
 struct DiffEqSystem <: AbstractSystem
     eqs::Vector{DiffEq}
     iv::Variable
     dvs::Vector{Variable}
     ps::Vector{Variable}
     jac::RefValue{Matrix{Expression}}
-    function DiffEqSystem(eqs, iv, dvs, ps)
-        jac = RefValue(Matrix{Expression}(undef, 0, 0))
-        new(eqs, iv, dvs, ps, jac)
-    end
-end
+    function DiffEqSystem(eqs)
+        reformatted = to_diffeq.(eqs)
 
-function DiffEqSystem(eqs)
-    dvs, = extract_elements(eqs, [_is_dependent])
-    ivs = unique(vcat((dv.dependents for dv ∈ dvs)...))
-    length(ivs) == 1 || throw(ArgumentError("one independent variable currently supported"))
-    iv = first(ivs)
-    ps, = extract_elements(eqs, [_is_parameter(iv)])
-    DiffEqSystem(eqs, iv, dvs, ps)
-end
+        ivs = unique(r[1] for r ∈ reformatted)
+        length(ivs) == 1 || throw(ArgumentError("one independent variable currently supported"))
+        iv = first(ivs)
 
-function DiffEqSystem(eqs, iv)
-    dvs, ps = extract_elements(eqs, [_is_dependent, _is_parameter(iv)])
-    DiffEqSystem(eqs, iv, dvs, ps)
+        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
+
+        jac = RefValue(Matrix{Expression}(undef, 0, 0))
+
+        new(deqs, iv, dvs, ps, jac)
+    end
 end
 
 
 function calculate_jacobian(sys::DiffEqSystem)
     isempty(sys.jac[]) || return sys.jac[]  # use cached Jacobian, if possible
-    rhs = [eq.rhs for eq in sys.eqs]
+    rhs = [eq.rhs for eq ∈ sys.eqs]
 
-    jac = expand_derivatives.(calculate_jacobian(rhs, sys.dvs))
+    jac = expand_derivatives.(calculate_jacobian(rhs, sys.dvs, sys.iv))
     sys.jac[] = jac  # cache Jacobian
     return jac
 end
@@ -70,16 +72,30 @@ function generate_jacobian(sys::DiffEqSystem; version::FunctionVersion = ArrayFu
     return build_function(jac, sys.dvs, sys.ps, (sys.iv.name,); version = version)
 end
 
+struct DiffEqToExpr
+    sys::DiffEqSystem
+end
+function (f::DiffEqToExpr)(O::Operation)
+    if isa(O.op, Variable)
+        isequal(O.op, f.sys.iv) && return O.op.name  # independent variable
+        O.op ∈ f.sys.dvs        && return O.op.name  # dependent variables
+        isempty(O.args)         && return O.op.name  # 0-ary parameters
+        return build_expr(:call, Any[O.op.name; f.(O.args)])
+    end
+    return build_expr(:call, Any[O.op; f.(O.args)])
+end
+(f::DiffEqToExpr)(x) = convert(Expr, x)
+
 function generate_function(sys::DiffEqSystem; version::FunctionVersion = ArrayFunction)
-    rhss = [eq.rhs for eq ∈ sys.eqs]
-    return build_function(rhss, sys.dvs, sys.ps, (sys.iv.name,); version = version)
+    rhss = [deq.rhs for deq ∈ sys.eqs]
+    return build_function(rhss, sys.dvs, sys.ps, (sys.iv.name,), DiffEqToExpr(sys); version = version)
 end
 
 
 function generate_ode_iW(sys::DiffEqSystem, simplify=true; version::FunctionVersion = ArrayFunction)
     jac = calculate_jacobian(sys)
 
-    gam = Variable(:gam; known = true)
+    gam = Variable(:gam; known = true)()
 
     W = LinearAlgebra.I - gam*jac
     W = SMatrix{size(W,1),size(W,2)}(W)

src/systems/diffeqs/first_order_transform.jl

@@ -4,7 +4,7 @@ export ode_order_lowering
 function lower_varname(var::Variable, idv, order)
     order == 0 && return var
     name = Symbol(var.name, :_, string(idv.name)^order)
-    return Variable(name, var.dependents; known = var.known)
+    return Variable(name; known = var.known)
 end
 
 function ode_order_lowering(sys::DiffEqSystem)

src/utils.jl

@@ -30,7 +30,7 @@ function flatten_expr!(x)
     x
 end
 
-function build_function(rhss, vs, ps, args = (); version::FunctionVersion)
+function build_function(rhss, vs, ps, args = (), conv = rhs -> convert(Expr, rhs); version::FunctionVersion)
     var_pairs   = [(u.name, :(u[$i])) for (i, u) ∈ enumerate(vs)]
     param_pairs = [(p.name, :(p[$i])) for (i, p) ∈ enumerate(ps)]
     (ls, rs) = zip(var_pairs..., param_pairs...)
@@ -39,11 +39,11 @@ function build_function(rhss, vs, ps, args = (); version::FunctionVersion)
 
     if version === ArrayFunction
         X = gensym()
-        sys_exprs = [:($X[$i] = $(convert(Expr, rhs))) for (i, rhs) ∈ enumerate(rhss)]
+        sys_exprs = [:($X[$i] = $(conv(rhs))) for (i, rhs) ∈ enumerate(rhss)]
         let_expr = Expr(:let, var_eqs, build_expr(:block, sys_exprs))
         :(($X,u,p,$(args...)) -> $let_expr)
     elseif version === SArrayFunction
-        sys_expr = build_expr(:tuple, [convert(Expr, rhs) for rhs ∈ rhss])
+        sys_expr = build_expr(:tuple, [conv(rhs) for rhs ∈ rhss])
         let_expr = Expr(:let, var_eqs, sys_expr)
         :((u,p,$(args...)) -> begin
             X = $let_expr
@@ -63,6 +63,9 @@ 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) =
-    any(isequal(t), x.dependents) || any(has_dependent(t), x.dependents)
+Base.occursin(t::Expression) = Base.Fix1(occursin, t)
+Base.occursin(t::Expression, x::Operation ) = isequal(x, t) || any(occursin(t), x.args)
+Base.occursin(t::Expression, x::Expression) = isequal(x, t)
+
+clean(x::Variable) = x
+clean(O::Operation) = isa(O.op, Variable) ? O.op : throw(ArgumentError("invalid variable: $(O.op)"))

src/variables.jl

@@ -1,13 +1,12 @@
 export Variable, @variables, @parameters
 
 
-struct Variable <: Expression
+struct Variable <: Function
     name::Symbol
-    dependents::Vector{Variable}
     known::Bool
-    Variable(name, dependents = Variable[]; known = false) =
-        new(name, dependents, known)
+    Variable(name; known = false) = new(name, known)
 end
+(x::Variable)(args...) = Operation(x, collect(Expression, args))
 
 
 struct Constant <: Expression
@@ -30,11 +29,7 @@ Base.isequal(c::Constant, n::Number) = c.value == n
 Base.isequal(n::Number, c::Constant) = c.value == n
 Base.isequal(a::Constant, b::Constant) = a.value == b.value
 
-function Base.convert(::Type{Expr}, x::Variable)
-    x.known               || return x.name
-    isempty(x.dependents) && return x.name
-    return :($(x.name)($(convert.(Expr, x.dependents)...)))
-end
+Base.convert(::Type{Expr}, x::Variable) = x
 Base.convert(::Type{Expr}, c::Constant) = c.value
 
 Base.show(io::IO, x::Variable) = print(io, x.name)
@@ -56,15 +51,14 @@ function _parse_vars(macroname, known, x)
         @assert iscall || issym "@$macroname expects a tuple of expressions (`@$macroname x y z(t)`)"
 
         if iscall
-            dependents = :(Variable[$(_var.args[2:end]...)])
             var_name = _var.args[1]
+            expr = :($var_name = $Variable($(Meta.quot(var_name)); known = $known)($(_var.args[2:end]...)))
         else
-            dependents = Variable[]
             var_name = _var
+            expr = :($var_name = $Variable($(Meta.quot(var_name)); known = $known))
         end
 
         push!(var_names, var_name)
-        expr = :($var_name = $Variable($(Meta.quot(var_name)), $dependents; known = $known))
         push!(ex.args, expr)
     end
     push!(ex.args, build_expr(:tuple, var_names))

test/derivatives.jl

@@ -2,9 +2,9 @@ using ModelingToolkit
 using Test
 
 # Derivatives
-@parameters t σ ρ β
+@parameters t() σ() ρ() β()
 @variables x(t) y(t) z(t)
-@derivatives D'~t D2''~t
+@derivatives D'~t D2''~t Dx'~x
 
 @test isequal(expand_derivatives(D(t)), 1)
 @test isequal(expand_derivatives(D(D(t))), 0)
@@ -31,6 +31,7 @@ d2 = D(sin(t)*cos(t))
 @test isequal(expand_derivatives(d1), t*cos(t)+sin(t))
 @test isequal(expand_derivatives(d2), simplify_constants(cos(t)*cos(t)+sin(t)*(-1*sin(t))))
 
+@test_broken begin
 eqs = [0 ~ σ*(y-x),
        0 ~ x*(ρ-z)-y,
        0 ~ x*y - β*z]
@@ -45,13 +46,17 @@ jac = calculate_jacobian(sys)
 @test isequal(jac[3,1], y)
 @test isequal(jac[3,2], x)
 @test isequal(jac[3,3], -1*β)
+end
 
 # Variable dependence checking in differentiation
 @variables a(t) b(a)
 @test !isequal(D(b), 0)
+@test isequal(expand_derivatives(D(t)), 1)
+@test isequal(expand_derivatives(Dx(x)), 1)
 
 @test isequal(expand_derivatives(D(x * y)), simplify_constants(y*D(x) + x*D(y)))
 @test_broken isequal(expand_derivatives(D(x * y)), simplify_constants(D(x)*y + x*D(y)))
 
 @test isequal(expand_derivatives(D(2t)), 2)
 @test isequal(expand_derivatives(D(2x)), 2D(x))
+@test_broken isequal(expand_derivatives(D(x^2)), simplify_constants(2 * x * D(x)))

test/simplify.jl

@@ -1,7 +1,7 @@
 using ModelingToolkit
 using Test
 
-@parameters t
+@parameters t()
 @variables x(t) y(t) z(t)
 
 null_op = 0*t