Fix function generation

Author: HarrisonGrodin

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()