Merge pull request #122 from JuliaDiffEq/hg/fix/toexpr

Variable function fixes

Author: ChrisRackauckas

README.md

@@ -28,7 +28,7 @@ and parameters. Therefore we label them as follows:
 using ModelingToolkit
 
 # Define some variables
-@parameters t σ ρ β
+@parameters t() σ() ρ() β()
 @variables x(t) y(t) z(t)
 @derivatives D'~t
 ```
@@ -80,22 +80,22 @@ state of the previous ODE. This is the nonlinear system defined by where the
 derivatives are zero. We use (unknown) variables for our nonlinear system.
 
 ```julia
-@variables x y z
-@parameters σ ρ β
+@variables x() y() z()
+@parameters σ() ρ() β()
 
 # Define a nonlinear system
 eqs = [0 ~ σ*(y-x),
        0 ~ x*(ρ-z)-y,
        0 ~ x*y - β*z]
 ns = NonlinearSystem(eqs, [x,y,z])
-nlsys_func = generate_function(ns, [x,y,z], [ρ,σ,β])
+nlsys_func = generate_function(ns, [x,y,z], [σ,ρ,β])
 ```
 
 which generates:
 
 ```julia
 :((##364, u, p)->begin
-          let (x, z, y, ρ, σ, β) = (u[1], u[2], u[3], p[1], p[2], p[3])
+          let (x, y, z, σ, ρ, β) = (u[1], u[2], u[3], p[1], p[2], p[3])
               ##364[1] = σ * (y - x)
               ##364[2] = x * (ρ - z) - y
               ##364[3] = x * y - β * z

src/systems/diffeqs/diffeqsystem.jl

@@ -94,11 +94,6 @@ function calculate_jacobian(sys::ODESystem)
     return jac
 end
 
-function generate_jacobian(sys::ODESystem; version::FunctionVersion = ArrayFunction)
-    jac = calculate_jacobian(sys)
-    return build_function(jac, sys.dvs, sys.ps, (sys.iv.name,); version = version)
-end
-
 struct ODEToExpr
     sys::ODESystem
 end
@@ -113,6 +108,11 @@ function (f::ODEToExpr)(O::Operation)
 end
 (f::ODEToExpr)(x) = convert(Expr, x)
 
+function generate_jacobian(sys::ODESystem; version::FunctionVersion = ArrayFunction)
+    jac = calculate_jacobian(sys)
+    return build_function(jac, sys.dvs, sys.ps, (sys.iv.name,), ODEToExpr(sys); version = version)
+end
+
 function generate_function(sys::ODESystem, dvs, ps; version::FunctionVersion = ArrayFunction)
     rhss = [deq.rhs for deq ∈ sys.eqs]
     dvs′ = [clean(dv) for dv ∈ dvs]

src/systems/nonlinear/nonlinear_system.jl

@@ -45,12 +45,26 @@ end
 
 function generate_jacobian(sys::NonlinearSystem; version::FunctionVersion = ArrayFunction)
     jac = calculate_jacobian(sys)
-    return build_function(jac, clean.(sys.vs), sys.ps; version = version)
+    return build_function(jac, clean.(sys.vs), sys.ps, (), NLSysToExpr(sys); version = version)
 end
 
+struct NLSysToExpr
+    sys::NonlinearSystem
+end
+function (f::NLSysToExpr)(O::Operation)
+    any(isequal(O), f.sys.vs) && 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, ps; version::FunctionVersion = ArrayFunction)
     rhss = [eq.rhs for eq ∈ sys.eqs]
     vs′ = [clean(v) for v ∈ vs]
     ps′ = [clean(p) for p ∈ ps]
-    return build_function(rhss, vs′, ps′; version = version)
+    return build_function(rhss, vs′, ps′, (), NLSysToExpr(sys); version = version)
 end

test/system_construction.jl

@@ -119,8 +119,12 @@ eqs = [0 ~ σ*(y-x),
        0 ~ x*y - β*z]
 ns = NonlinearSystem(eqs, [x,y,z])
 test_nlsys_inference("standard", ns, (x, y, z), (σ, ρ, β))
-
-generate_function(ns, [x,y,z], [σ,ρ,β])
+@test begin
+    f = eval(generate_function(ns, [x,y,z], [σ,ρ,β]))
+    du = [0.0, 0.0, 0.0]
+    f(du, [1,2,3], [1,2,3])
+    du ≈ [1, -3, -7]
+end
 
 @derivatives D'~t
 @parameters A() B() C()