Merge pull request #267 from SciML/calc

Fix calculated functions on components

Author: ChrisRackauckas

src/systems/diffeqs/abstractodesystem.jl

@@ -1,19 +1,19 @@
 function calculate_tgrad(sys::AbstractODESystem)
   isempty(sys.tgrad[]) || return sys.tgrad[]  # use cached tgrad, if possible
-  rhs = [detime_dvs(eq.rhs) for eq ∈ sys.eqs]
+  rhs = [detime_dvs(eq.rhs) for eq ∈ equations(sys)]
   iv = sys.iv()
   notime_tgrad = [expand_derivatives(ModelingToolkit.Differential(iv)(r)) for r in rhs]
-  tgrad = retime_dvs.(notime_tgrad,(sys.dvs,),iv)
+  tgrad = retime_dvs.(notime_tgrad,(states(sys),),iv)
   sys.tgrad[] = tgrad
   return tgrad
 end
 
 function calculate_jacobian(sys::AbstractODESystem)
     isempty(sys.jac[]) || return sys.jac[]  # use cached Jacobian, if possible
-    rhs = [eq.rhs for eq ∈ sys.eqs]
+    rhs = [eq.rhs for eq ∈ equations(sys)]
 
     iv = sys.iv()
-    dvs = [dv(iv) for dv ∈ sys.dvs]
+    dvs = [dv(iv) for dv ∈ states(sys)]
 
     jac = expand_derivatives.(calculate_jacobian(rhs, dvs))
     sys.jac[] = jac  # cache Jacobian
@@ -34,18 +34,18 @@ function (f::ODEToExpr)(O::Operation)
 end
 (f::ODEToExpr)(x) = convert(Expr, x)
 
-function generate_tgrad(sys::AbstractODESystem, dvs = sys.dvs, ps = sys.ps, expression = Val{true}; kwargs...)
+function generate_tgrad(sys::AbstractODESystem, dvs = states(sys), ps = parameters(sys), expression = Val{true}; kwargs...)
     tgrad = calculate_tgrad(sys)
     return build_function(tgrad, dvs, ps, (sys.iv.name,), ODEToExpr(sys), expression; kwargs...)
 end
 
-function generate_jacobian(sys::AbstractODESystem, dvs = sys.dvs, ps = sys.ps, expression = Val{true}; kwargs...)
+function generate_jacobian(sys::AbstractODESystem, dvs = states(sys), ps = parameters(sys), expression = Val{true}; kwargs...)
     jac = calculate_jacobian(sys)
     return build_function(jac, dvs, ps, (sys.iv.name,), ODEToExpr(sys), expression; kwargs...)
 end
 
-function generate_function(sys::AbstractODESystem, dvs = sys.dvs, ps = sys.ps, expression = Val{true}; kwargs...)
-    rhss = [deq.rhs for deq ∈ sys.eqs]
+function generate_function(sys::AbstractODESystem, dvs = states(sys), ps = parameters(sys), expression = Val{true}; kwargs...)
+    rhss = [deq.rhs for deq ∈ equations(sys)]
     dvs′ = [clean(dv) for dv ∈ dvs]
     ps′ = [clean(p) for p ∈ ps]
     return build_function(rhss, dvs′, ps′, (sys.iv.name,), ODEToExpr(sys), expression; kwargs...)
@@ -75,7 +75,7 @@ function calculate_factorized_W(sys::AbstractODESystem, simplify=true)
     (Wfact,Wfact_t)
 end
 
-function generate_factorized_W(sys::AbstractODESystem, vs = sys.dvs, ps = sys.ps, simplify=true, expression = Val{true}; kwargs...)
+function generate_factorized_W(sys::AbstractODESystem, vs = states(sys), ps = parameters(sys), simplify=true, expression = Val{true}; kwargs...)
     (Wfact,Wfact_t) = calculate_factorized_W(sys,simplify)
     siz = size(Wfact)
     constructor = :(x -> begin
@@ -90,13 +90,14 @@ function generate_factorized_W(sys::AbstractODESystem, vs = sys.dvs, ps = sys.ps
 end
 
 function calculate_massmatrix(sys::AbstractODESystem, simplify=true)
-    eqs = sys.eqs
+    eqs = equations(sys)
+    dvs = states(sys)
     M = zeros(length(eqs),length(eqs))
     for (i,eq) in enumerate(eqs)
         if eq.lhs isa Constant
             @assert eq.lhs.value == 0
         elseif eq.lhs.op isa Differential
-            j = findfirst(x->isequal(x.name,var_from_nested_derivative(eq.lhs)[1].name),sys.dvs)
+            j = findfirst(x->isequal(x.name,var_from_nested_derivative(eq.lhs)[1].name),dvs)
             M[i,j] = 1
         else
             error("Only semi-explicit mass matrices are currently supported")
@@ -113,7 +114,8 @@ Create an `ODEFunction` from the [`ODESystem`](@ref). The arguments `dvs` and `p
 are used to set the order of the dependent variable and parameter vectors,
 respectively.
 """
-function DiffEqBase.ODEFunction{iip}(sys::AbstractODESystem, dvs = sys.dvs, ps = sys.ps;
+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})
@@ -166,14 +168,14 @@ function DiffEqBase.ODEFunction(sys::AbstractODESystem, args...; kwargs...)
 end
 
 function namespace_variables(sys::AbstractODESystem)
-    [rename(x,renamespace(sys.name,x.name)) for x in sys.dvs]
+    [rename(x,renamespace(sys.name,x.name)) for x in states(sys)]
 end
 
 function namespace_parameters(sys::AbstractODESystem)
-    [rename(x,renamespace(sys.name,x.name)) for x in sys.ps]
+    [rename(x,renamespace(sys.name,x.name)) for x in parameters(sys)]
 end
 
-namespace_equations(sys::AbstractODESystem) = namespace_equation.(sys.eqs,sys.name,sys.iv.name)
+namespace_equations(sys::AbstractODESystem) = namespace_equation.(equations(sys),sys.name,sys.iv.name)
 
 function namespace_equation(eq::Equation,name,ivname)
     _lhs = namespace_operation(eq.lhs,name,ivname)
@@ -191,11 +193,11 @@ end
 namespace_operation(O::Constant,name,ivname) = O
 
 independent_variable(sys::AbstractODESystem) = sys.iv
-states(sys::AbstractODESystem) = [sys.dvs;reduce(vcat,namespace_variables.(sys.systems))]
-parameters(sys::AbstractODESystem) = [sys.ps;reduce(vcat,namespace_parameters.(sys.systems))]
+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))]
 
 function equations(sys::AbstractODESystem)
-    [sys.eqs;reduce(vcat,namespace_equations.(sys.systems))]
+    isempty(sys.systems) ? sys.eqs : [sys.eqs;reduce(vcat,namespace_equations.(sys.systems))]
 end
 
 function states(sys::AbstractODESystem,name::Symbol)
@@ -208,6 +210,18 @@ function parameters(sys::AbstractODESystem,name::Symbol)
     Variable(Symbol(string(sys.name)*"′"*string(x.name)),known=x.known)(sys.iv())
 end
 
+function states(sys::AbstractODESystem,args...)
+    name = last(args)
+    extra_names = reduce(*,["′$(x.name)" for x in args[1:end-1]])
+    Variable(Symbol(string(sys.name)*extra_names*"′"*string(name)))(sys.iv())
+end
+
+function parameters(sys::AbstractODESystem,args...)
+    name = last(args)
+    extra_names = reduce(*,["′$(x.name)" for x in args[1:end-1]])
+    Variable(Symbol(string(sys.name)*extra_names*"′"*string(name)))(sys.iv())
+end
+
 function _eq_unordered(a, b)
     length(a) === length(b) || return false
     n = length(a)

src/systems/nonlinear/nonlinear_system.jl

@@ -62,10 +62,8 @@ function _find_params(O, vs)
     return ps
 end
 
-
-independent_variables(sys::NonlinearSystem) = Set{Variable}()
-dependent_variables(sys::NonlinearSystem) = Set{Expression}(sys.vs)
-parameters(sys::NonlinearSystem) = Set{Variable}(sys.ps)
+states(sys::NonlinearSystem) = sys.vs
+parameters(sys::NonlinearSystem) = sys.ps
 
 
 function calculate_jacobian(sys::NonlinearSystem)

test/components.jl

@@ -18,7 +18,7 @@ lorenz2 = ODESystem(eqs,t,[x,y,z],[σ,ρ,β],name=:lorenz2)
 @variables a(t)
 connnectedeqs = [D(a) ~ a*states(lorenz1,:x)]
 
-connected = ODESystem(connnectedeqs,t,[a],[α],systems=[lorenz1,lorenz2],name=:connectedlorenz)
+connected1 = ODESystem(connnectedeqs,t,[a],[α],systems=[lorenz1,lorenz2],name=: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′β
@@ -31,6 +31,60 @@ eqs_flat = [D(a) ~ a*lorenz1′x,
             D(lorenz2′y) ~ lorenz2′x*(lorenz2′ρ-lorenz2′z)-lorenz2′y,
             0 ~ lorenz2′x + lorenz2′y + lorenz2′β*lorenz2′z]
 
-@test [x.name for x in states(connected)] == [:a,:lorenz1′x,:lorenz1′y,:lorenz1′z,:lorenz2′x,:lorenz2′y,:lorenz2′z]
-@test [x.name for x in parameters(connected)] == [:α,:lorenz1′σ,:lorenz1′ρ,:lorenz1′β,:lorenz2′σ,:lorenz2′ρ,:lorenz2′β]
-@test eqs_flat == equations(connected)
+@test [x.name for x in states(connected1)] == [:a,:lorenz1′x,:lorenz1′y,:lorenz1′z,:lorenz2′x,:lorenz2′y,:lorenz2′z]
+@test [x.name for x in parameters(connected1)] == [:α,:lorenz1′σ,:lorenz1′ρ,:lorenz1′β,:lorenz2′σ,:lorenz2′ρ,:lorenz2′β]
+@test eqs_flat == equations(connected1)
+
+connected2 = ODESystem(connnectedeqs,t,[a],[α],systems=[lorenz1,lorenz2],name=:connected2)
+
+@parameters γ
+@variables g(t)
+connnectedeqs2 = [D(g) ~ g*states(connected1,lorenz1,:x)]
+doublelevel = ODESystem(connnectedeqs2,t,[g],[γ],systems=[connected1,connected2],name=: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]
+@test [x.name for x in parameters(doublelevel)] == [:γ,
+                                                   :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′β]
+
+@variables connected1′a(t) connected1′lorenz1′x(t) connected1′lorenz1′y(t) connected1′lorenz1′z(t) connected1′lorenz2′x(t) connected1′lorenz2′y(t) connected1′lorenz2′z(t)
+@variables connected2′a(t) connected2′lorenz1′x(t) connected2′lorenz1′y(t) connected2′lorenz1′z(t) connected2′lorenz2′x(t) connected2′lorenz2′y(t) connected2′lorenz2′z(t)
+@parameters connected1′α connected1′lorenz1′σ connected1′lorenz1′ρ connected1′lorenz1′β connected1′lorenz2′σ connected1′lorenz2′ρ connected1′lorenz2′β
+@parameters connected2′α connected2′lorenz1′σ connected2′lorenz1′ρ connected2′lorenz1′β connected2′lorenz2′σ connected2′lorenz2′ρ connected2′lorenz2′β
+
+eqs_flat2 = [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_flat2 == equations(doublelevel)
+
+M = Array(I,15,15)
+M[5,5] = false
+M[8,8] = false
+M[12,12] = false
+M[15,15] = false
+@test calculate_massmatrix(doublelevel) == M
+
+jac = [connected1′lorenz1′x 0 g zeros(1,12)
+       zeros(7,1) calculate_jacobian(connected1) zeros(Expression,7,7)
+       zeros(Expression,7,8) calculate_jacobian(connected2)]
+
+jac2 = [connected1′lorenz1′x 0 g zeros(1,12)
+        zeros(7,1) ModelingToolkit.namespace_operation.(calculate_jacobian(connected1),connected1.name,:t) zeros(Expression,7,7)
+        zeros(Expression,7,8) ModelingToolkit.namespace_operation.(calculate_jacobian(connected2),connected2.name,:t)]
+
+@test all(isequal.(calculate_jacobian(doublelevel),jac2))

test/odesystem.jl

@@ -24,19 +24,19 @@ end
 _clean(x::Variable) = x
 function test_diffeq_inference(name, sys, iv, dvs, ps)
     @testset "ODESystem construction: $name" begin
-        @test independent_variables(sys) == Set([_clean(iv)])
-        @test dependent_variables(sys)   == Set(_clean.(dvs))
-        @test parameters(sys)            == Set(_clean.(ps))
+        @test independent_variable(sys) ==  _clean(iv)
+        @test Set(states(sys))          == Set(_clean.(dvs))
+        @test Set(parameters(sys))      == Set(_clean.(ps))
     end
 end
 function test_nlsys_inference(name, sys, vs, ps)
     @testset "NonlinearSystem construction: $name" begin
-        @test dependent_variables(sys) == Set(vs)
-        @test parameters(sys)          == Set(_clean.(ps))
+        @test Set(states(sys))     == Set(vs)
+        @test Set(parameters(sys)) == Set(_clean.(ps))
     end
 end
 
-test_diffeq_inference("standard", de, t, (x, y, z), (σ, ρ, β))
+test_diffeq_inference("standard", de, t, [x, y, z], [ρ, σ, β])
 generate_function(de, [x,y,z], [σ,ρ,β])
 jac_expr = generate_jacobian(de)
 jac = calculate_jacobian(de)