Merge pull request #277 from SciML/dot

add dot syntax, variable searching in scope, and variable->operation

Author: ChrisRackauckas

src/operations.jl

@@ -66,6 +66,9 @@ Operation(x) = convert(Operation, x)
 #convert to Expr
 Base.Expr(op::Operation) = simplified_expr(op)
 Base.convert(::Type{Expr},x::Operation) = Expr(x)
+function Base.convert(::Type{Variable},x::Operation)
+    x.op isa Variable ? x.op : throw(error("This Operation is not a Variable"))
+end
 
 # promotion
 Base.promote_rule(::Type{<:Constant}, ::Type{<:Operation}) = Operation

src/systems/diffeqs/abstractodesystem.jl

@@ -107,60 +107,6 @@ function calculate_massmatrix(sys::AbstractODESystem, simplify=true)
     M == I ? I : M
 end
 
-"""
-$(SIGNATURES)
-
-Create an `ODEFunction` from the [`ODESystem`](@ref). The arguments `dvs` and `ps`
-are used to set the order of the dependent variable and parameter vectors,
-respectively.
-"""
-function DiffEqBase.ODEFunction{iip}(sys::AbstractODESystem, dvs = states(sys),
-                                     ps = parameters(sys);
-                                     version = nothing, tgrad=false,
-                                     jac = false, Wfact = false) where {iip}
-    f_oop,f_iip = generate_function(sys, dvs, ps, Val{false})
-
-    f(u,p,t) = f_oop(u,p,t)
-    f(du,u,p,t) = f_iip(du,u,p,t)
-
-    if tgrad
-        tgrad_oop,tgrad_iip = generate_tgrad(sys, dvs, ps, Val{false})
-        _tgrad(u,p,t) = tgrad_oop(u,p,t)
-        _tgrad(J,u,p,t) = tgrad_iip(J,u,p,t)
-    else
-        _tgrad = nothing
-    end
-
-    if jac
-        jac_oop,jac_iip = generate_jacobian(sys, dvs, ps, Val{false})
-        _jac(u,p,t) = jac_oop(u,p,t)
-        _jac(J,u,p,t) = jac_iip(J,u,p,t)
-    else
-        _jac = nothing
-    end
-
-    if Wfact
-        tmp_Wfact,tmp_Wfact_t = generate_factorized_W(sys, dvs, ps, true, Val{false})
-        Wfact_oop, Wfact_iip = tmp_Wfact
-        Wfact_oop_t, Wfact_iip_t = tmp_Wfact_t
-        _Wfact(u,p,dtgamma,t) = Wfact_oop(u,p,dtgamma,t)
-        _Wfact(W,u,p,dtgamma,t) = Wfact_iip(W,u,p,dtgamma,t)
-        _Wfact_t(u,p,dtgamma,t) = Wfact_oop_t(u,p,dtgamma,t)
-        _Wfact_t(W,u,p,dtgamma,t) = Wfact_iip_t(W,u,p,dtgamma,t)
-    else
-        _Wfact,_Wfact_t = nothing,nothing
-    end
-
-    M = calculate_massmatrix(sys)
-
-    ODEFunction{iip}(f,jac=_jac,
-                      tgrad = _tgrad,
-                      Wfact = _Wfact,
-                      Wfact_t = _Wfact_t,
-                      mass_matrix = M,
-                      syms = Symbol.(sys.dvs))
-end
-
 renamespace(namespace,name) = Symbol(string(namespace)*"′"*string(name))
 
 function DiffEqBase.ODEFunction(sys::AbstractODESystem, args...; kwargs...)
@@ -192,6 +138,8 @@ function namespace_operation(O::Operation,name,ivname)
 end
 namespace_operation(O::Constant,name,ivname) = O
 
+
+
 independent_variable(sys::AbstractODESystem) = sys.iv
 states(sys::AbstractODESystem) = isempty(sys.systems) ? sys.dvs : [sys.dvs;reduce(vcat,namespace_variables.(sys.systems))]
 parameters(sys::AbstractODESystem) = isempty(sys.systems) ? sys.ps : [sys.ps;reduce(vcat,namespace_parameters.(sys.systems))]
@@ -202,12 +150,12 @@ end
 
 function states(sys::AbstractODESystem,name::Symbol)
     x = sys.dvs[findfirst(x->x.name==name,sys.dvs)]
-    Variable(Symbol(string(sys.name)*"′"*string(x.name)),known=x.known)(sys.iv())
+    Variable(Symbol(string(sys.name)*"′"*string(x.name)))(sys.iv())
 end
 
 function parameters(sys::AbstractODESystem,name::Symbol)
     x = sys.ps[findfirst(x->x.name==name,sys.ps)]
-    Variable(Symbol(string(sys.name)*"′"*string(x.name)),known=x.known)(sys.iv())
+    Variable(Symbol(string(sys.name)*"′"*string(x.name)))(sys.iv())
 end
 
 function states(sys::AbstractODESystem,args...)
@@ -234,3 +182,77 @@ function _eq_unordered(a, b)
     end
     return true
 end
+
+"""
+$(SIGNATURES)
+
+Create an `ODEFunction` from the [`ODESystem`](@ref). The arguments `dvs` and `ps`
+are used to set the order of the dependent variable and parameter vectors,
+respectively.
+"""
+function DiffEqBase.ODEFunction{iip}(sys::AbstractODESystem, dvs = states(sys),
+                                     ps = parameters(sys);
+                                     version = nothing, tgrad=false,
+                                     jac = false, Wfact = false) where {iip}
+    f_oop,f_iip = generate_function(sys, dvs, ps, Val{false})
+
+    f(u,p,t) = f_oop(u,p,t)
+    f(du,u,p,t) = f_iip(du,u,p,t)
+
+    if tgrad
+        tgrad_oop,tgrad_iip = generate_tgrad(sys, dvs, ps, Val{false})
+        _tgrad(u,p,t) = tgrad_oop(u,p,t)
+        _tgrad(J,u,p,t) = tgrad_iip(J,u,p,t)
+    else
+        _tgrad = nothing
+    end
+
+    if jac
+        jac_oop,jac_iip = generate_jacobian(sys, dvs, ps, Val{false})
+        _jac(u,p,t) = jac_oop(u,p,t)
+        _jac(J,u,p,t) = jac_iip(J,u,p,t)
+    else
+        _jac = nothing
+    end
+
+    if Wfact
+        tmp_Wfact,tmp_Wfact_t = generate_factorized_W(sys, dvs, ps, true, Val{false})
+        Wfact_oop, Wfact_iip = tmp_Wfact
+        Wfact_oop_t, Wfact_iip_t = tmp_Wfact_t
+        _Wfact(u,p,dtgamma,t) = Wfact_oop(u,p,dtgamma,t)
+        _Wfact(W,u,p,dtgamma,t) = Wfact_iip(W,u,p,dtgamma,t)
+        _Wfact_t(u,p,dtgamma,t) = Wfact_oop_t(u,p,dtgamma,t)
+        _Wfact_t(W,u,p,dtgamma,t) = Wfact_iip_t(W,u,p,dtgamma,t)
+    else
+        _Wfact,_Wfact_t = nothing,nothing
+    end
+
+    M = calculate_massmatrix(sys)
+
+    ODEFunction{iip}(f,jac=_jac,
+                      tgrad = _tgrad,
+                      Wfact = _Wfact,
+                      Wfact_t = _Wfact_t,
+                      mass_matrix = M,
+                      syms = Symbol.(sys.dvs))
+end
+
+function Base.getproperty(sys::AbstractODESystem, name::Symbol)
+    if name ∈ fieldnames(typeof(sys))
+        return getfield(sys,name)
+    elseif !isempty(sys.systems)
+        i = findfirst(x->x.name==name,sys.systems)
+        if i !== nothing
+            return rename(sys.systems[i],renamespace(sys.name,name))
+        end
+    end
+    i = findfirst(x->x.name==name,sys.dvs)
+    if i !== nothing
+        return rename(sys.dvs[i],renamespace(sys.name,name))(getfield(sys,:iv)())
+    end
+    i = findfirst(x->x.name==name,sys.ps)
+    if i !== nothing
+        return rename(sys.ps[i],renamespace(sys.name,name))()
+    end
+    throw(error("Variable name does not exist"))
+end

src/systems/diffeqs/first_order_transform.jl

@@ -30,7 +30,7 @@ function ode_order_lowering(eqs, iv)
             any(isequal(var), vars) || push!(vars, var)
         end
         var′ = lower_varname(var, iv, maxorder - 1)
-        rhs′ = rename(eq.rhs)
+        rhs′ = rename_lower_order(eq.rhs)
         push!(new_eqs,Differential(iv())(var′(iv())) ~ rhs′)
     end
 
@@ -50,11 +50,11 @@ function ode_order_lowering(eqs, iv)
     return (new_eqs, new_vars)
 end
 
-function rename(O::Expression)
+function rename_lower_order(O::Expression)
     isa(O, Operation) || return O
     if is_derivative(O)
         (x, t, order) = flatten_differential(O)
         return lower_varname(x.op, t.op, order)(x.args...)
     end
-    return Operation(O.op, rename.(O.args))
+    return Operation(O.op, rename_lower_order.(O.args))
 end

src/systems/diffeqs/odesystem.jl

@@ -94,3 +94,7 @@ 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
+
+function rename(sys::ODESystem,name)
+    ODESystem(sys.eqs, sys.iv, sys.dvs, sys.ps, sys.tgrad, sys.jac, sys.Wfact, sys.Wfact_t, name, sys.systems)
+end

src/systems/diffeqs/sdesystem.jl

@@ -117,3 +117,7 @@ end
 function DiffEqBase.SDEFunction(sys::SDESystem, args...; kwargs...)
     SDEFunction{true}(sys, args...; kwargs...)
 end
+
+function rename(sys::SDESystem,name)
+    ODESystem(sys.eqs, sys.noiseeqs, sys.iv, sys.dvs, sys.ps, sys.tgrad, sys.jac, sys.Wfact, sys.Wfact_t, name, sys.systems)
+end

src/variables.jl

@@ -40,6 +40,7 @@ end
 
 vartype(::Variable{T}) where T = T
 (x::Variable)(args...) = Operation(x, collect(Expression, args))
+rename(x::Variable{T},name) where T = Variable{T}(name)
 
 Base.isequal(x::Variable, y::Variable) = x.name == y.name
 Base.print(io::IO, x::Variable) = show(io, x)

test/components.jl

@@ -16,14 +16,26 @@ lorenz2 = ODESystem(eqs,t,[x,y,z],[σ,ρ,β],name=:lorenz2)
 
 @parameters α
 @variables a(t)
-connnectedeqs = [D(a) ~ a*states(lorenz1,:x)]
+connnectedeqs = [D(a) ~ a*lorenz1.x]
 
 connected1 = ODESystem(connnectedeqs,t,[a],[α],systems=[lorenz1,lorenz2],name=:connected1)
 
+eqs_flat = [D(a) ~ a*lorenz1.x,
+            D(lorenz1.x) ~ lorenz1.σ*(lorenz1.y-lorenz1.x),
+            D(lorenz1.y) ~ lorenz1.x*(lorenz1.ρ-lorenz1.z)-lorenz1.y,
+            0 ~ lorenz1.x + lorenz1.y + lorenz1.β*lorenz1.z,
+            D(lorenz2.x) ~ lorenz2.σ*(lorenz2.y-lorenz2.x),
+            D(lorenz2.y) ~ lorenz2.x*(lorenz2.ρ-lorenz2.z)-lorenz2.y,
+            0 ~ lorenz2.x + lorenz2.y + lorenz2.β*lorenz2.z]
+
+@test states(connected1) == convert.(Variable,[a,lorenz1.x,lorenz1.y,lorenz1.z,lorenz2.x,lorenz2.y,lorenz2.z])
+@test parameters(connected1) == convert.(Variable,[α,lorenz1.σ,lorenz1.ρ,lorenz1.β,lorenz2.σ,lorenz2.ρ,lorenz2.β])
+@test eqs_flat == equations(connected1)
+
 @variables lorenz1′x(t) lorenz1′y(t) lorenz1′z(t) lorenz2′x(t) lorenz2′y(t) lorenz2′z(t)
 @parameters lorenz1′σ lorenz1′ρ lorenz1′β lorenz2′σ lorenz2′ρ lorenz2′β
 
-eqs_flat = [D(a) ~ a*lorenz1′x,
+eqs_flat2 = [D(a) ~ a*lorenz1′x,
             D(lorenz1′x) ~ lorenz1′σ*(lorenz1′y-lorenz1′x),
             D(lorenz1′y) ~ lorenz1′x*(lorenz1′ρ-lorenz1′z)-lorenz1′y,
             0 ~ lorenz1′x + lorenz1′y + lorenz1′β*lorenz1′z,
@@ -39,9 +51,34 @@ connected2 = ODESystem(connnectedeqs,t,[a],[α],systems=[lorenz1,lorenz2],name=:
 
 @parameters γ
 @variables g(t)
-connnectedeqs2 = [D(g) ~ g*states(connected1,lorenz1,:x)]
+connnectedeqs2 = [D(g) ~ g*connected1.lorenz1.x]
 doublelevel = ODESystem(connnectedeqs2,t,[g],[γ],systems=[connected1,connected2],name=:doublelevel)
 
+@test states(doublelevel) == convert.(Variable,[g,connected1.a,connected1.lorenz1.x,connected1.lorenz1.y,connected1.lorenz1.z,connected1.lorenz2.x,connected1.lorenz2.y,connected1.lorenz2.z,
+                                          connected2.a,connected2.lorenz1.x,connected2.lorenz1.y,connected2.lorenz1.z,connected2.lorenz2.x,connected2.lorenz2.y,connected2.lorenz2.z])
+
+@test parameters(doublelevel) == convert.(Variable,[γ,
+                                             connected1.α,connected1.lorenz1.σ,connected1.lorenz1.ρ,connected1.lorenz1.β,connected1.lorenz2.σ,connected1.lorenz2.ρ,connected1.lorenz2.β,
+                                             connected2.α,connected2.lorenz1.σ,connected2.lorenz1.ρ,connected2.lorenz1.β,connected2.lorenz2.σ,connected2.lorenz2.ρ,connected2.lorenz2.β])
+
+eqs_flat = [D(g) ~ g*connected1.lorenz1.x,
+             D(connected1.a) ~ connected1.a*connected1.lorenz1.x,
+             D(connected1.lorenz1.x) ~ connected1.lorenz1.σ*(connected1.lorenz1.y-connected1.lorenz1.x),
+             D(connected1.lorenz1.y) ~ connected1.lorenz1.x*(connected1.lorenz1.ρ-connected1.lorenz1.z)-connected1.lorenz1.y,
+             0 ~ connected1.lorenz1.x + connected1.lorenz1.y + connected1.lorenz1.β*connected1.lorenz1.z,
+             D(connected1.lorenz2.x) ~ connected1.lorenz2.σ*(connected1.lorenz2.y-connected1.lorenz2.x),
+             D(connected1.lorenz2.y) ~ connected1.lorenz2.x*(connected1.lorenz2.ρ-connected1.lorenz2.z)-connected1.lorenz2.y,
+             0 ~ connected1.lorenz2.x + connected1.lorenz2.y + connected1.lorenz2.β*connected1.lorenz2.z,
+             D(connected2.a) ~ connected2.a*connected2.lorenz1.x,
+             D(connected2.lorenz1.x) ~ connected2.lorenz1.σ*(connected2.lorenz1.y-connected2.lorenz1.x),
+             D(connected2.lorenz1.y) ~ connected2.lorenz1.x*(connected2.lorenz1.ρ-connected2.lorenz1.z)-connected2.lorenz1.y,
+             0 ~ connected2.lorenz1.x + connected2.lorenz1.y + connected2.lorenz1.β*connected2.lorenz1.z,
+             D(connected2.lorenz2.x) ~ connected2.lorenz2.σ*(connected2.lorenz2.y-connected2.lorenz2.x),
+             D(connected2.lorenz2.y) ~ connected2.lorenz2.x*(connected2.lorenz2.ρ-connected2.lorenz2.z)-connected2.lorenz2.y,
+             0 ~ connected2.lorenz2.x + connected2.lorenz2.y + connected2.lorenz2.β*connected2.lorenz2.z]
+
+@test eqs_flat == equations(doublelevel)
+
 @test [x.name for x in states(doublelevel)] == [:g,
                                                 :connected1′a,:connected1′lorenz1′x,:connected1′lorenz1′y,:connected1′lorenz1′z,:connected1′lorenz2′x,:connected1′lorenz2′y,:connected1′lorenz2′z,
                                                 :connected2′a,:connected2′lorenz1′x,:connected2′lorenz1′y,:connected2′lorenz1′z,:connected2′lorenz2′x,:connected2′lorenz2′y,:connected2′lorenz2′z]