working ODE jacobians

Author: ChrisRackauckas

src/systems/diffeqs/diffeqsystem.jl

@@ -55,12 +55,11 @@ function build_equals_expr(eq)
     end
 end
 
-function generate_ode_jacobian(sys::DiffEqSystem,simplify=true)
-    var_exprs = [:($(sys.dvs[i].name) = u[$i]) for i in 1:length(sys.dvs)]
-    param_exprs = [:($(sys.ps[i].name) = p[$i]) for i in 1:length(sys.ps)]
+function calculate_jacobian(sys::DiffEqSystem,simplify=true)
     diff_idxs = map(eq->eq.args[1].diff !=nothing,sys.eqs)
     diff_exprs = sys.eqs[diff_idxs]
     rhs = [eq.args[2] for eq in diff_exprs]
+    # Handle intermediate calculations by substitution
     calcs = sys.eqs[.!(diff_idxs)]
     for i in 1:length(calcs)
         find_replace!.(rhs,calcs[i].args[1],calcs[i].args[2])
@@ -73,6 +72,18 @@ function generate_ode_jacobian(sys::DiffEqSystem,simplify=true)
     sys_exprs
 end
 
+function generate_ode_jacobian(sys::DiffEqSystem,simplify=true)
+    var_exprs = [:($(sys.dvs[i].name) = u[$i]) for i in 1:length(sys.dvs)]
+    param_exprs = [:($(sys.ps[i].name) = p[$i]) for i in 1:length(sys.ps)]
+    diff_idxs = map(eq->eq.args[1].diff !=nothing,sys.eqs)
+    diff_exprs = sys.eqs[diff_idxs]
+    jac = calculate_jacobian(sys,simplify)
+    jac_exprs = [:(J[$i,$j] = $(Expr(jac[i,j]))) for i in 1:size(jac,1), j in 1:size(jac,2)]
+    exprs = vcat(var_exprs,param_exprs,vec(jac_exprs))
+    block = expr_arr_to_block(exprs)
+    :((J,u,p,t)->$(block))
+end
+
 function DiffEqBase.DiffEqFunction(sys::DiffEqSystem)
     expr = generate_ode_function(sys)
     DiffEqFunction{true}(eval(expr))

test/system_construction.jl

@@ -15,16 +15,15 @@ eqs = [D*x ~ σ*(y-x),
        D*z ~ x*y - β*z]
 de = DiffEqSystem(eqs,[t],[x,y,z],Variable[],[σ,ρ,β])
 SciCompDSL.generate_ode_function(de)
-jac = SciCompDSL.generate_ode_jacobian(de,false)
-jac = SciCompDSL.generate_ode_jacobian(de)
+jac_expr = SciCompDSL.generate_ode_jacobian(de)
+jac = SciCompDSL.calculate_jacobian(de)
 f = DiffEqFunction(de)
 I - jac
-@test_broken inv(jac)
+@test_broken inv(I - jac)
 
 # Differential equation with automatic extraction of variables on rhs
 de2 = DiffEqSystem(eqs, [t])
 
-
 function test_vars_extraction(de, de2)
     for el in (:ivs, :dvs, :vs, :ps)
         names2 = sort(collect(var.name for var in getfield(de2,el)))
@@ -67,8 +66,7 @@ eqs = [a ~ y-x,
        D*z ~ x*y - β*z]
 de = DiffEqSystem(eqs,[t],[x,y,z],[a],[σ,ρ,β])
 SciCompDSL.generate_ode_function(de)
-jac = SciCompDSL.generate_ode_jacobian(de,false)
-jac = SciCompDSL.generate_ode_jacobian(de)
+jac = SciCompDSL.calculate_jacobian(de)
 f = DiffEqFunction(de)
 
 # Define a nonlinear system
@@ -95,7 +93,6 @@ eqs = [0 ~ σ*(y-x),
        0 ~ x*y - β*z]
 ns = NonlinearSystem(eqs)
 nlsys_func = SciCompDSL.generate_nlsys_function(ns)
-jac = SciCompDSL.generate_nlsys_jacobian(ns,false)
 jac = SciCompDSL.generate_nlsys_jacobian(ns)
 f = @eval eval(nlsys_func)
 
@@ -107,5 +104,4 @@ eqs = [a ~ y-x,
        0 ~ x*y - β*z]
 ns = NonlinearSystem(eqs,[x,y,z],[σ,ρ,β])
 nlsys_func = SciCompDSL.generate_nlsys_function(ns)
-jac = SciCompDSL.generate_nlsys_jacobian(ns,false)
 jac = SciCompDSL.generate_nlsys_jacobian(ns)