remove ODEExpr

Author: ChrisRackauckas

src/build_function.jl

@@ -120,16 +120,16 @@ function numbered_expr(O::Operation,vars,parameters;
                         varname=varname,paramname=paramname) for x in O.args]...)
 end
 
-function numbered_expr(de::ModelingToolkit.ODEExpr,vars::Vector{Variable},parameters;
+function numbered_expr(de::ModelingToolkit.Equation,vars::Vector{Variable},parameters;
                        derivname=:du,varname=:u,paramname=:p)
-    i = findfirst(x->isequal(x.name,de.x.name),vars)
+    i = findfirst(x->isequal(x.name,var_from_nested_derivative(de.lhs)[1].name),vars)
     :($derivname[$i] = $(numbered_expr(de.rhs,vars,parameters;
                                      derivname=derivname,
                                      varname=varname,paramname=paramname)))
 end
-function numbered_expr(de::ModelingToolkit.ODEExpr,vars::Vector{Operation},parameters;
+function numbered_expr(de::ModelingToolkit.Equation,vars::Vector{Operation},parameters;
                        derivname=:du,varname=:u,paramname=:p)
-    i = findfirst(x->isequal(x.op.name,de.x.name),vars)
+    i = findfirst(x->isequal(x.op.name,var_from_nested_derivative(de.lhs)[1].name),vars)
     :($derivname[$i] = $(numbered_expr(de.rhs,vars,parameters;
                                      derivname=derivname,
                                      varname=varname,paramname=paramname)))

src/latexify_recipes.jl

@@ -22,26 +22,6 @@
     return lhs, rhs
 end
 
-@latexrecipe function f(eqs::Vector{ModelingToolkit.ODEExpr}; iv=:t)
-    # Set default option values.
-    env --> :align
-
-    # Convert both the left and right hand side to expressions of basic types
-    # that latexify can deal with.
-
-    rhs = getfield.(eqs, :rhs)
-    rhs = convert.(Expr, rhs)
-    rhs = [postwalk(x -> x isa ModelingToolkit.Constant ? x.value : x, eq) for eq in rhs]
-    rhs = [postwalk(x -> x isa Expr && length(x.args) == 1 ? x.args[1] : x, eq) for eq in rhs]
-    rhs = [postwalk(x -> x isa Expr && x.args[1] == :Differential && length(x.args[2].args) == 2 ? :($(Symbol(:d, x.args[2]))/($(Symbol(:d, x.args[2].args[2])))) : x, eq) for eq in rhs]
-    rhs = [postwalk(x -> x isa Expr && x.args[1] == :Differential ? "\\frac{d\\left($(Latexify.latexraw(x.args[2]))\\right)}{d$iv}" : x, eq) for eq in rhs]
-
-    var = getfield.(getfield.(eqs, :x),:name)
-    ns  = getfield.(eqs, :n)
-    lhs = [ns[i] == 1 ? Latexify.LaTeXString("\\frac{d$(Latexify.latexraw(var[i]))}{d$iv}") : LaTeXString("\\frac{d^{$(ns[i])}$(Latexify.latexraw(var[i]))}{d$iv^{$(ns[i])}}") for i in 1:length(var)]
-    return lhs, rhs
-end
-
 @latexrecipe function f(sys::ModelingToolkit.ODESystem; iv=:t)
     return latexify(sys.eqs; iv=iv)
 end

src/systems/diffeqs/abstractodesystem.jl

@@ -142,3 +142,20 @@ end
 function DiffEqBase.ODEFunction(sys::AbstractODESystem, args...; kwargs...)
     ODEFunction{true}(sys, args...; kwargs...)
 end
+
+independent_variables(sys::AbstractODESystem) = Set{Variable}([sys.iv])
+dependent_variables(sys::AbstractODESystem) = Set{Variable}(sys.dvs)
+parameters(sys::AbstractODESystem) = Set{Variable}(sys.ps)
+
+function _eq_unordered(a, b)
+    length(a) === length(b) || return false
+    n = length(a)
+    idxs = Set(1:n)
+    for x ∈ a
+        idx = findfirst(isequal(x), b)
+        idx === nothing && return false
+        idx ∈ idxs      || return false
+        delete!(idxs, idx)
+    end
+    return true
+end

src/systems/diffeqs/first_order_transform.jl

@@ -4,25 +4,34 @@ function lower_varname(var::Variable, idv, order)
     return Variable(name; known = var.known)
 end
 
+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])
+    isequal(t, O.op.x) || throw(ArgumentError("non-matching differentials on lhs: $t, $(O.op.x)"))
+    return (x, t, order + 1)
+end
+
 function ode_order_lowering(sys::ODESystem)
     eqs_lowered, _ = ode_order_lowering(sys.eqs, sys.iv)
-    ODESystem(eqs_lowered, sys.iv, [eq.x for eq in eqs_lowered], sys.ps)
+    ODESystem(eqs_lowered, sys.iv, [var_from_nested_derivative(eq.lhs)[1] for eq in eqs_lowered], sys.ps)
 end
+
 function ode_order_lowering(eqs, iv)
     var_order = Dict{Variable,Int}()
     vars = Variable[]
-    new_eqs = similar(eqs, ODEExpr)
+    new_eqs = Equation[]
     new_vars = Variable[]
 
     for (i, eq) ∈ enumerate(eqs)
-        var, maxorder = eq.x, eq.n
+        var, maxorder = var_from_nested_derivative(eq.lhs)
         if maxorder > get(var_order, var, 0)
             var_order[var] = maxorder
             any(isequal(var), vars) || push!(vars, var)
         end
-        var′ = lower_varname(eq.x, iv, eq.n - 1)
+        var′ = lower_varname(var, iv, maxorder - 1)
         rhs′ = rename(eq.rhs)
-        new_eqs[i] = ODEExpr(var′, 1, rhs′)
+        push!(new_eqs,Differential(iv())(var′(iv())) ~ rhs′)
     end
 
     for var ∈ vars
@@ -33,7 +42,7 @@ function ode_order_lowering(eqs, iv)
             push!(new_vars, rvar)
 
             rhs = rvar(iv())
-            eq = ODEExpr(lvar, 1, rhs)
+            eq = Differential(iv())(lvar(iv())) ~ rhs
             push!(new_eqs, eq)
         end
     end

src/systems/diffeqs/odesystem.jl

@@ -1,29 +1,3 @@
-isintermediate(eq::Equation) = !(isa(eq.lhs, Operation) && isa(eq.lhs.op, Differential))
-
-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])
-    isequal(t, O.op.x) || throw(ArgumentError("non-matching differentials on lhs: $t, $(O.op.x)"))
-    return (x, t, order + 1)
-end
-
-struct ODEExpr  # dⁿx/dtⁿ = rhs
-    x::Variable
-    n::Int
-    rhs::Expression
-end
-function Base.convert(::Type{ODEExpr},eq::Equation)
-    isintermediate(eq) && throw(ArgumentError("intermediate equation received"))
-    (x, t, n) = flatten_differential(eq.lhs)
-    (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, ODEExpr(x.op, n, eq.rhs)
-end
-Base.:(==)(a::ODEExpr, b::ODEExpr) = isequal((a.x, a.n, a.rhs), (b.x, b.n, b.rhs))
-
 """
 $(TYPEDEF)
 
@@ -50,7 +24,7 @@ de = ODESystem(eqs)
 """
 struct ODESystem <: AbstractODESystem
     """The ODEs defining the system."""
-    eqs::Vector{ODEExpr}
+    eqs::Vector{Equation}
     """Independent variable."""
     iv::Variable
     """Dependent (state) variables."""
@@ -83,55 +57,35 @@ struct ODESystem <: AbstractODESystem
     name::Symbol
 end
 
-function ODESystem(deqs::AbstractVector{ODEExpr}, iv, dvs, ps; name=gensym(:ODESystem))
+function ODESystem(deqs::AbstractVector{<:Equation}, iv, dvs, ps;
+                   name=gensym(:ODESystem))
+    iv′ = clean(iv)
+    dvs′ = [clean(dv) for dv ∈ dvs]
+    ps′ = [clean(p) for p ∈ ps]
     tgrad = RefValue(Vector{Expression}(undef, 0))
     jac = RefValue(Matrix{Expression}(undef, 0, 0))
     Wfact   = RefValue(Matrix{Expression}(undef, 0, 0))
     Wfact_t = RefValue(Matrix{Expression}(undef, 0, 0))
-    ODESystem(deqs, iv, dvs, ps, tgrad, jac, Wfact, Wfact_t, name)
+    ODESystem(deqs, iv′, dvs′, ps′, tgrad, jac, Wfact, Wfact_t, name)
 end
 
-function ODESystem(eqs; kwargs...)
-    reformatted = convert.(ODEExpr,eqs)
+var_from_nested_derivative(x) = var_from_nested_derivative(x,0)
+var_from_nested_derivative(x,i) = x.op isa Differential ? var_from_nested_derivative(x.args[1],i+1) : (x.op,i)
+iv_from_nested_derivative(x) = x.op isa Differential ? iv_from_nested_derivative(x.args[1]) : x.args[1].op
 
-    ivs = unique(r[1] for r ∈ reformatted)
+function ODESystem(eqs; kwargs...)
+    ivs = unique(iv_from_nested_derivative(eq.lhs) for eq ∈ eqs)
     length(ivs) == 1 || throw(ArgumentError("one independent variable currently supported"))
     iv = first(ivs)
 
-    deqs = [r[2] for r ∈ reformatted]
-
-    dvs = unique(deq.x for deq ∈ deqs)
-    ps = filter(vars(deq.rhs for deq ∈ deqs)) do x
+    dvs = unique(var_from_nested_derivative(eq.lhs)[1] for eq ∈ eqs)
+    ps = filter(vars(eq.rhs for eq ∈ eqs)) do x
         x.known & !isequal(x, iv)
     end |> collect
-
-    ODESystem(deqs, iv, dvs, ps; kwargs...)
+    ODESystem(eqs, iv, dvs, ps; kwargs...)
 end
 
-function ODESystem(deqs::AbstractVector{<:Equation}, iv, dvs, ps; kwargs...)
-    _dvs = [deq.op for deq ∈ dvs]
-    _iv = iv.op
-    _ps = [p.op for p ∈ ps]
-    ODESystem(getindex.(convert.(ODEExpr,deqs),2), _iv, _dvs, _ps; kwargs...)
-end
-
-function _eq_unordered(a, b)
-    length(a) === length(b) || return false
-    n = length(a)
-    idxs = Set(1:n)
-    for x ∈ a
-        idx = findfirst(isequal(x), b)
-        idx === nothing && return false
-        idx ∈ idxs      || return false
-        delete!(idxs, idx)
-    end
-    return true
-end
 Base.:(==)(sys1::ODESystem, sys2::ODESystem) =
     _eq_unordered(sys1.eqs, sys2.eqs) && isequal(sys1.iv, sys2.iv) &&
     _eq_unordered(sys1.dvs, sys2.dvs) && _eq_unordered(sys1.ps, sys2.ps)
 # NOTE: equality does not check cached Jacobian
-
-independent_variables(sys::ODESystem) = Set{Variable}([sys.iv])
-dependent_variables(sys::ODESystem) = Set{Variable}(sys.dvs)
-parameters(sys::ODESystem) = Set{Variable}(sys.ps)

src/systems/diffeqs/sdesystem.jl

@@ -1,6 +1,6 @@
 struct SDESystem <: AbstractODESystem
     """The expressions defining the drift term."""
-    eqs::Vector{ODEExpr}
+    eqs::Vector{Equation}
     """The expressions defining the diffusion term."""
     noiseeqs
     """Independent variable."""
@@ -35,19 +35,15 @@ struct SDESystem <: AbstractODESystem
     name::Symbol
 end
 
-function SDESystem(deqs::AbstractVector{ODEExpr}, neqs, iv, dvs, ps; name = gensym(:SDESystem))
+function SDESystem(deqs::AbstractVector{<:Equation}, neqs, iv, dvs, ps; name = gensym(:SDESystem))
+    dvs′ = [clean(dv) for dv ∈ dvs]
+    ps′ = [clean(p) for p ∈ ps]
+    iv′ = clean(iv)
     tgrad = RefValue(Vector{Expression}(undef, 0))
     jac = RefValue(Matrix{Expression}(undef, 0, 0))
     Wfact   = RefValue(Matrix{Expression}(undef, 0, 0))
     Wfact_t = RefValue(Matrix{Expression}(undef, 0, 0))
-    SDESystem(deqs, neqs, iv, dvs, ps, tgrad, jac, Wfact, Wfact_t, name)
-end
-
-function SDESystem(deqs::AbstractVector{<:Equation}, neqs, iv, dvs, ps; kwargs...)
-    _dvs = [deq.op for deq ∈ dvs]
-    _iv = iv.op
-    _ps = [p.op for p ∈ ps]
-    SDESystem(getindex.(convert.(ODEExpr,deqs),2), neqs, _iv, _dvs, _ps; kwargs...)
+    SDESystem(deqs, neqs, iv′, dvs′, ps′, tgrad, jac, Wfact, Wfact_t, name)
 end
 
 function generate_diffusion_function(sys::SDESystem, dvs = sys.dvs, ps = sys.ps, expression = Val{true}; kwargs...)

src/systems/diffeqs/validation.jl

@@ -0,0 +1,17 @@
+isintermediate(eq::Equation) = !(isa(eq.lhs, Operation) && isa(eq.lhs.op, Differential))
+
+struct ODEExpr  # dⁿx/dtⁿ = rhs
+    x::Variable
+    n::Int
+    rhs::Expression
+end
+function Base.convert(::Type{ODEExpr},eq::Equation)
+    isintermediate(eq) && throw(ArgumentError("intermediate equation received"))
+    (x, t, n) = flatten_differential(eq.lhs)
+    (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, ODEExpr(x.op, n, eq.rhs)
+end
+Base.:(==)(a::ODEExpr, b::ODEExpr) = isequal((a.x, a.n, a.rhs), (b.x, b.n, b.rhs))

test/build_targets.jl

@@ -4,7 +4,7 @@ using ModelingToolkit, Test
 @derivatives D'~t
 eqs = [D(x) ~ a*x - x*y,
        D(y) ~ -3y + x*y]
-@test ModelingToolkit.build_function(eqs,[x,y],[a],t,target = ModelingToolkit.StanTarget()) ==
+@test ModelingToolkit.build_function(eqs,ModelingToolkit.clean.([x,y]),ModelingToolkit.clean.([a]),t,target = ModelingToolkit.StanTarget()) ==
     """
     real[] diffeqf(real t,real[] internal_var___u,real[] internal_var___p,real[] x_r,int[] x_i) {
       real internal_var___du[2];

test/components.jl

@@ -0,0 +1,37 @@
+using ModelingToolkit, StaticArrays, LinearAlgebra
+using DiffEqBase
+using Test
+
+# Define some variables
+@parameters t σ ρ β
+@variables x(t) y(t) z(t)
+@derivatives D'~t
+
+eqs = [D(x) ~ σ*(y-x),
+       D(y) ~ x*(ρ-z)-y,
+       0 ~ x + y + z]
+
+de1 = ODESystem(eqs,t,[x,y,z],[σ,ρ,β],name=:lorenz1)
+de2 = ODESystem(eqs,t,[x,y,z],[σ,ρ,β],name=:lorenz2)
+
+@parameters α
+@variables a(t)
+for
+
+convert(DAESystem,sys)
+convert(ODESystem,sys)
+
+connected = ODESystems(connnectedeqs,[a],[α],systems=[de1,de2],name=:connectedlorenz)
+
+deleteat!(de2.eqs,3)
+de2
+push!(de2.eqs,convert(ModelingToolkit.ODEExpr,D(z) ~ z^2)[2])
+
+check_consistency(de2)
+
+b = a*z
+           [states(de1,:x) ~ α*states(de2,:z),
+            D(a) ~ b*states(de1,:x)]
+
+states(states) == [lorenz1′x,lorenz1′y,lorenz1′z,lorenz2′x,lorenz2′y,lorenz2′z,a]
+parameters(connected) == [lorenz1′σ,lorenz1′ρ,lorenz1′β,lorenz2′σ,lorenz2′ρ,lorenz2′β,α]

test/odesystem.jl

@@ -137,10 +137,11 @@ lowered_eqs = [D(u_tt) ~ 2u_tt + u_t + x_t + 1
                D(u_t)  ~ u_tt
                D(u)    ~ u_t
                D(x)    ~ x_t]
+
 @test de1 == ODESystem(lowered_eqs)
 
 # issue #219
-@test de1.dvs == [eq.x for eq in de1.eqs] == ODESystem(lowered_eqs).dvs
+@test de1.dvs == [ModelingToolkit.var_from_nested_derivative(eq.lhs)[1] for eq in de1.eqs] == ODESystem(lowered_eqs).dvs
 
 test_diffeq_inference("first-order transform", de1, t, [u_tt, x_t, u_t, u, x], [])
 du = zeros(5)