Refactor system storage

README.md

@@ -46,7 +46,7 @@ Each operation builds an `Operation` type, and thus `eqs` is an array of
 analyzed by other programs. We can turn this into a `DiffEqSystem` via:
 
 ```julia
-de = DiffEqSystem(eqs,[t],[x,y,z],[σ,ρ,β])
+de = DiffEqSystem(eqs,t,[x,y,z],[σ,ρ,β])
 de = DiffEqSystem(eqs)
 ```
 

src/ModelingToolkit.jl

@@ -20,7 +20,7 @@ Base.convert(::Type{Variable},x::Int64) = Constant(x)
 
 function caclulate_jacobian end
 
-@enum FunctionVersions ArrayFunction=1 SArrayFunction=2
+@enum FunctionVersion ArrayFunction=1 SArrayFunction=2
 
 include("operations.jl")
 include("differentials.jl")

src/equations.jl

@@ -1,19 +1,20 @@
 export Equation
 
 
-mutable struct Equation
+struct Equation
     lhs::Expression
     rhs::Expression
 end
 Base.broadcastable(eq::Equation) = Ref(eq)
+Base.:(==)(a::Equation, b::Equation) = (a.lhs, a.rhs) == (b.lhs, b.rhs)
 
 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.subtype === :Unknown && !isempty(x.dependents)
-_is_parameter(ivs) = x -> x.subtype === :Parameter && x ∉ ivs
+_is_parameter(iv) = x -> x.subtype === :Parameter && x ≠ iv
 _subtype(subtype::Symbol) = x -> x.subtype === subtype
 
 function extract_elements(eqs, predicates)
@@ -29,10 +30,10 @@ function extract_elements(eqs, predicates)
     return result
 end
 
+get_args(O::Operation) = O.args
+get_args(eq::Equation) = Expression[eq.lhs, eq.rhs]
 function vars!(vars, op)
-    args = isa(op, Equation) ? Expression[op.lhs, op.rhs] : op.args
-
-    for arg ∈ args
+    for arg ∈ get_args(op)
         if isa(arg, Operation)
             vars!(vars, arg)
         elseif isa(arg, Variable)

src/systems/diffeqs/diffeqsystem.jl

@@ -1,45 +1,58 @@
-mutable struct DiffEqSystem <: AbstractSystem
-    eqs::Vector{Equation}
-    ivs::Vector{Variable}
+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
+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)
+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]
+
+struct DiffEqSystem <: AbstractSystem
+    eqs::Vector{DiffEq}
+    iv::Variable
     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)
+    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
 
-DiffEqSystem(eqs, ivs, dvs, ps) = DiffEqSystem(eqs, ivs, dvs, ps, Matrix{Expression}(undef,0,0))
-
 function DiffEqSystem(eqs)
     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))
+    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
 
-function DiffEqSystem(eqs, ivs)
-    dvs, ps = extract_elements(eqs, [_is_dependent, _is_parameter(ivs)])
-    DiffEqSystem(eqs, ivs, dvs, ps, Matrix{Expression}(undef,0,0))
+function DiffEqSystem(eqs, iv)
+    dvs, ps = extract_elements(eqs, [_is_dependent, _is_parameter(iv)])
+    DiffEqSystem(eqs, iv, dvs, ps)
 end
 
-isintermediate(eq::Equation) = !(isa(eq.lhs, Operation) && isa(eq.lhs.op, Differential))
-
 
-function generate_ode_function(sys::DiffEqSystem;version = ArrayFunction)
+function generate_ode_function(sys::DiffEqSystem; version::FunctionVersion = 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)]
     sys_exprs = build_equals_expr.(sys.eqs)
-    if version == ArrayFunction
-        dvar_exprs = [:(du[$i] = $(Symbol("$(sys.dvs[i].name)_$(sys.ivs[1].name)"))) for i in eachindex(sys.dvs)]
+    if version === ArrayFunction
+        dvar_exprs = [:(du[$i] = $(Symbol("$(sys.dvs[i].name)_$(sys.iv.name)"))) for i in eachindex(sys.dvs)]
         exprs = vcat(var_exprs,param_exprs,sys_exprs,dvar_exprs)
         block = expr_arr_to_block(exprs)
         :((du,u,p,t)->$(toexpr(block)))
-    elseif version == SArrayFunction
-        dvar_exprs = [:($(Symbol("$(sys.dvs[i].name)_$(sys.ivs[1].name)"))) for i in eachindex(sys.dvs)]
+    elseif version === SArrayFunction
+        dvar_exprs = [:($(Symbol("$(sys.dvs[i].name)_$(sys.iv.name)"))) for i in eachindex(sys.dvs)]
         svector_expr = quote
             E = eltype(tuple($(dvar_exprs...)))
             T = StaticArrays.similar_type(typeof(u), E)
@@ -51,26 +64,24 @@ function generate_ode_function(sys::DiffEqSystem;version = ArrayFunction)
     end
 end
 
-function build_equals_expr(eq::Equation)
-    @assert !isintermediate(eq)
-
-    lhs = Symbol(eq.lhs.args[1].name, :_, eq.lhs.op.x.name)
+function build_equals_expr(eq::DiffEq)
+    lhs = Symbol(eq.var.name, :_, eq.D.x.name)
     return :($lhs = $(convert(Expr, eq.rhs)))
 end
 
 function calculate_jacobian(sys::DiffEqSystem, simplify=true)
+    isempty(sys.jac[]) || return sys.jac[]  # use cached Jacobian, if possible
     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]
-    sys_exprs
+    jac = expand_derivatives.(calculate_jacobian(rhs, sys.dvs))
+    sys.jac[] = jac  # cache Jacobian
+    return jac
 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)]
     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)]
     exprs = vcat(var_exprs,param_exprs,vec(jac_exprs))
     block = expr_arr_to_block(exprs)
@@ -80,7 +91,7 @@ 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)]
-    jac = sys.jac
+    jac = calculate_jacobian(sys, simplify)
 
     gam = Parameter(:gam)
 
@@ -109,12 +120,12 @@ function generate_ode_iW(sys::DiffEqSystem, simplify=true)
     :((iW,u,p,gam,t)->$(block)),:((iW,u,p,gam,t)->$(block2))
 end
 
-function DiffEqBase.ODEFunction(sys::DiffEqSystem;version = ArrayFunction,kwargs...)
-    expr = generate_ode_function(sys;version=version,kwargs...)
-    if version == ArrayFunction
-      ODEFunction{true}(eval(expr))
-    elseif version == SArrayFunction
-      ODEFunction{false}(eval(expr))
+function DiffEqBase.ODEFunction(sys::DiffEqSystem; version::FunctionVersion = ArrayFunction)
+    expr = generate_ode_function(sys; version = version)
+    if version === ArrayFunction
+        ODEFunction{true}(eval(expr))
+    elseif version === SArrayFunction
+        ODEFunction{false}(eval(expr))
     end
 end
 

src/systems/diffeqs/first_order_transform.jl

@@ -1,80 +1,56 @@
-extract_idv(eq::Equation) = eq.lhs.op.x
+extract_idv(eq::DiffEq) = eq.D.x
 
-function lower_varname(O::Operation, naming_scheme; lower=false)
-    @assert isa(O.op, Differential)
-
-    D, x = O.op, O.args[1]
+function lower_varname(D::Differential, x; lower=false)
     order = lower ? D.order-1 : D.order
-
-    lower_varname(x, D.x, order, naming_scheme)
+    return lower_varname(x, D.x, order)
 end
-function lower_varname(var::Variable, idv, order::Int, naming_scheme)
+function lower_varname(var::Variable, idv, order::Int)
     sym = var.name
-    name = order == 0 ? sym : Symbol(sym, naming_scheme, string(idv.name)^order)
+    name = order == 0 ? sym : Symbol(sym, :_, string(idv.name)^order)
     return Variable(name, var.subtype, var.dependents)
 end
 
-function ode_order_lowering(sys::DiffEqSystem; kwargs...)
-    eqs = sys.eqs
-    ivs = sys.ivs
-    eqs_lowered = ode_order_lowering(eqs; kwargs...)
-    DiffEqSystem(eqs_lowered, ivs)
+function ode_order_lowering(sys::DiffEqSystem)
+    eqs_lowered = ode_order_lowering(sys.eqs, sys.iv)
+    DiffEqSystem(eqs_lowered, sys.iv)
 end
-ode_order_lowering(eqs; naming_scheme = "_") = ode_order_lowering!(deepcopy(eqs), naming_scheme)
-function ode_order_lowering!(eqs, naming_scheme)
-    idv = extract_idv(eqs[1])
-    D   = Differential(idv, 1)
+function ode_order_lowering(eqs, iv)
+    D = Differential(iv, 1)
     var_order = Dict{Variable,Int}()
     vars = Variable[]
-    dv_name = eqs[1].lhs.args[1].subtype
+    new_eqs = similar(eqs, DiffEq)
 
-    for eq in eqs
+    for (i, eq) ∈ enumerate(eqs)
         var, maxorder = extract_var_order(eq)
         maxorder == 1 && continue # fast pass
         if maxorder > get(var_order, var, 0)
             var_order[var] = maxorder
             var ∈ vars || push!(vars, var)
         end
-        lhs_renaming!(eq, D, naming_scheme)
-        rhs_renaming!(eq, naming_scheme)
+        var′ = lower_varname(eq.D, eq.var, lower = true)
+        rhs′ = rename(eq.rhs)
+        new_eqs[i] = DiffEq(D, var′, rhs′)
     end
 
     for var ∈ vars
         order = var_order[var]
         for o in (order-1):-1:1
-            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)
+            lvar = lower_varname(var, iv, o-1)
+            rhs = lower_varname(var, iv, o)
+            eq = DiffEq(D, lvar, rhs)
+            push!(new_eqs, eq)
         end
     end
 
-    return eqs
+    return new_eqs
 end
 
-function lhs_renaming!(eq, D, naming_scheme)
-    eq.lhs = D(lower_varname(eq.lhs, naming_scheme, lower=true))
-    return eq
+function rename(O::Expression)
+    isa(O, Operation) || return O
+    isa(O.op, Differential) && return lower_varname(O.op, O.args[1])
+    return Operation(O.op, rename.(O.args))
 end
-rhs_renaming!(eq, naming_scheme) = _rec_renaming!(eq.rhs, naming_scheme)
 
-function _rec_renaming!(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)
-            args[i] = _rec_renaming!(args[i], naming_scheme)
-        end
-    end
-    rhs
-end
-
-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
-    var = dv.args[1]
-    order = dv.op.order
-    return (var, order)
-end
+extract_var_order(eq::DiffEq) = (eq.var, eq.D.order)
 
 export ode_order_lowering

test/system_construction.jl

@@ -11,7 +11,7 @@ using Test
 eqs = [D(x) ~ σ*(y-x),
        D(y) ~ x*(ρ-z)-y,
        D(z) ~ x*y - β*z]
-de = DiffEqSystem(eqs,[t],[x,y,z],[σ,ρ,β])
+de = DiffEqSystem(eqs,t,[x,y,z],[σ,ρ,β])
 ModelingToolkit.generate_ode_function(de)
 ModelingToolkit.generate_ode_function(de;version=ModelingToolkit.SArrayFunction)
 jac_expr = ModelingToolkit.generate_ode_jacobian(de)
@@ -20,10 +20,11 @@ f = ODEFunction(de)
 ModelingToolkit.generate_ode_iW(de)
 
 # Differential equation with automatic extraction of variables
-de2 = DiffEqSystem(eqs, [t])
+de2 = DiffEqSystem(eqs, t)
 
 function test_vars_extraction(de, de2)
-    for el in (:ivs, :dvs, :ps)
+    @test de.iv == de2.iv
+    for el in (:dvs, :ps)
         names2 = sort(collect(var.name for var in getfield(de2,el)))
         names = sort(collect(var.name for var in getfield(de,el)))
         @test names2 == names
@@ -37,34 +38,21 @@ test_vars_extraction(de, de2)
 @Unknown 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 = DiffEqSystem(eqs, [t])
+de = DiffEqSystem(eqs, t)
 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]
-function test_eqs(eqs1, eqs2)
-    length(eqs1) == length(eqs2) || return false
-    eq = true
-    for (eq1, eq2) in zip(eqs1, eqs2)
-        lhs1, lhs2 = eq1.lhs, eq2.lhs
-        typeof(lhs1) === typeof(lhs2) || return false
-        for f in fieldnames(typeof(lhs1))
-            eq = eq & isequal(getfield(lhs1, f), getfield(lhs2, f))
-        end
-        eq = eq & isequal(eq1.rhs, eq2.rhs)
-    end
-    eq
-end
-@test test_eqs(de1.eqs, lowered_eqs)
+@test de1.eqs == convert.(ModelingToolkit.DiffEq, lowered_eqs)
 
 # Internal calculations
 a = y - x
 eqs = [D(x) ~ σ*a,
        D(y) ~ x*(ρ-z)-y,
        D(z) ~ x*y - β*z]
-de = DiffEqSystem(eqs,[t],[x,y,z],[σ,ρ,β])
+de = DiffEqSystem(eqs,t,[x,y,z],[σ,ρ,β])
 ModelingToolkit.generate_ode_function(de)
 jac = ModelingToolkit.calculate_jacobian(de)
 f = ODEFunction(de)
@@ -88,8 +76,8 @@ ModelingToolkit.generate_nlsys_function(ns)
 _x = y / C
 eqs = [D(x) ~ -A*x,
        D(y) ~ A*x - B*_x]
-de = DiffEqSystem(eqs,[t],[x,y],[A,B,C])
-test_vars_extraction(de, DiffEqSystem(eqs,[t]))
+de = DiffEqSystem(eqs,t,[x,y],[A,B,C])
+test_vars_extraction(de, DiffEqSystem(eqs,t))
 test_vars_extraction(de, DiffEqSystem(eqs))
 @test eval(ModelingToolkit.generate_ode_function(de))([0.0,0.0],[1.0,2.0],[1,2,3],0.0) ≈ -1/3