Merge pull request #296 from SciML/expression

fix composed expression generation

Author: ChrisRackauckas

README.md

@@ -42,7 +42,7 @@ sys = ODESystem(eqs)
 sys = ode_order_lowering(sys)
 
 u0 = [D(x) => 2.0,
-	  x => 1.0,
+      x => 1.0,
       y => 0.0,
       z => 0.0]
 
@@ -53,6 +53,7 @@ p  = [σ => 28.0,
 tspan = (0.0,100.0)
 prob = ODEProblem(sys,u0,tspan,p,jac=true)
 sol = solve(prob,Tsit5())
+using Plots; plot(sol,vars=(:x,:y))
 ```
 
 ![Lorenz2](https://user-images.githubusercontent.com/1814174/79118645-744eb580-7d5c-11ea-9c37-13c4efd585ca.png)
@@ -77,26 +78,28 @@ lorenz2 = ODESystem(eqs,name=:lorenz2)
 
 @variables α
 @parameters γ
-connections = [0 ~ lorenz1.x + lorenz2.y + sin(α*γ)]
+connections = [0 ~ lorenz1.x + lorenz2.y + α*γ]
 connected = ODESystem(connections,t,[α],[γ],systems=[lorenz1,lorenz2])
 
 u0 = [lorenz1.x => 1.0,
       lorenz1.y => 0.0,
       lorenz1.z => 0.0,
-	lorenz2.x => 1.0,
-	lorenz2.y => 0.0,
-	lorenz2.z => 0.0,
-	α => 2.0]
+	  lorenz2.x => 0.0,
+	  lorenz2.y => 1.0,
+	  lorenz2.z => 0.0,
+	  α => 2.0]
 
-p  = [lorenz1.σ => 28.0,
-      lorenz1.ρ => 10.0,
+p  = [lorenz1.σ => 10.0,
+      lorenz1.ρ => 28.0,
       lorenz1.β => 8/3,
-	lorenz2.σ => 28.0,
-	lorenz2.ρ => 10.0,
-	lorenz2.β => 8/3,
-	γ => 2.0]
+	  lorenz2.σ => 10.0,
+	  lorenz2.ρ => 28.0,
+	  lorenz2.β => 8/3,
+	  γ => 2.0]
 
 tspan = (0.0,100.0)
 prob = ODEProblem(connected,u0,tspan,p)
 sol = solve(prob,Rodas5())
+
+using Plots; plot(sol,vars=(:α,Symbol(lorenz1.x),Symbol(lorenz2.y)))
 ```

src/systems/abstractsystem.jl

@@ -83,3 +83,18 @@ function parameters(sys::AbstractSystem,args...)
     newname = renamespace(extra_names,name)
     rename(x,renamespace(sys.name,newname))()
 end
+
+struct AbstractSysToExpr
+    sys::AbstractSystem
+    states::Vector{Variable}
+end
+AbstractSysToExpr(sys) = AbstractSysToExpr(sys,states(sys))
+function (f::AbstractSysToExpr)(O::Operation)
+    any(isequal(O), f.states) && return O.op.name  # variables
+    if isa(O.op, Variable)
+        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::AbstractSysToExpr)(x) = convert(Expr, x)

src/systems/diffeqs/abstractodesystem.jl

@@ -22,11 +22,13 @@ end
 
 struct ODEToExpr
     sys::AbstractODESystem
+    states::Vector{Variable}
 end
+ODEToExpr(sys) = ODEToExpr(sys,states(sys))
 function (f::ODEToExpr)(O::Operation)
     if isa(O.op, Variable)
         isequal(O.op, f.sys.iv) && return O.op.name  # independent variable
-        O.op ∈ f.sys.states        && return O.op.name  # dependent variables
+        O.op ∈ f.states         && 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

src/systems/nonlinear/nonlinearsystem.jl

@@ -53,27 +53,14 @@ function generate_jacobian(sys::NonlinearSystem, vs = states(sys), ps = paramete
     if sparse
         jac = SparseArrays.sparse(jac)
     end
-    return build_function(jac, convert.(Variable,vs), convert.(Variable,ps), (), NLSysToExpr(sys))
+    return build_function(jac, convert.(Variable,vs), convert.(Variable,ps), (), AbstractSysToExpr(sys))
 end
 
-struct NLSysToExpr
-    sys::NonlinearSystem
-end
-function (f::NLSysToExpr)(O::Operation)
-    any(isequal(O), f.sys.states) && return O.op.name  # variables
-    if isa(O.op, Variable)
-        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::NLSysToExpr)(x) = convert(Expr, x)
-
 function generate_function(sys::NonlinearSystem, vs = states(sys), ps = parameters(sys), expression = Val{true}; kwargs...)
     rhss = [eq.rhs for eq ∈ sys.eqs]
     vs′ = convert.(Variable,vs)
     ps′ = convert.(Variable,ps)
-    return build_function(rhss, vs′, ps′, (), NLSysToExpr(sys), expression; kwargs...)
+    return build_function(rhss, vs′, ps′, (), AbstractSysToExpr(sys), expression; kwargs...)
 end
 
 """

src/systems/optimization/optimizationsystem.jl

@@ -45,7 +45,7 @@ end
 
 function generate_gradient(sys::OptimizationSystem, vs = states(sys), ps = parameters(sys), expression = Val{true}; kwargs...)
     grad = calculate_gradient(sys)
-    return build_function(grad, convert.(Variable,vs), convert.(Variable,ps), (), x->convert(Expr, x); kwargs...)
+    return build_function(grad, convert.(Variable,vs), convert.(Variable,ps), (), AbstractSysToExpr(sys); kwargs...)
 end
 
 function calculate_hessian(sys::OptimizationSystem)
@@ -58,13 +58,13 @@ function generate_hessian(sys::OptimizationSystem, vs = states(sys), ps = parame
     if sparse
         hes = sparse(hes)
     end
-    return build_function(hes, convert.(Variable,vs), convert.(Variable,ps), (), x->convert(Expr, x); kwargs...)
+    return build_function(hes, convert.(Variable,vs), convert.(Variable,ps), (), AbstractSysToExpr(sys); kwargs...)
 end
 
 function generate_function(sys::OptimizationSystem, vs = states(sys), ps = parameters(sys), expression = Val{true}; kwargs...)
     vs′ = convert.(Variable,vs)
     ps′ = convert.(Variable,ps)
-    return build_function(equations(sys), vs′, ps′, (), x->convert(Expr, x), expression; kwargs...)
+    return build_function(equations(sys), vs′, ps′, (), AbstractSysToExpr(sys), expression; kwargs...)
 end
 
 equations(sys::OptimizationSystem) = isempty(sys.systems) ? sys.op : sys.op + reduce(+,namespace_operation.(sys.systems))