Merge pull request #1030 from SciML/nonlinear_jac

Cache nonlinear solve jacobians

Author: ChrisRackauckas

src/systems/nonlinear/nonlinearsystem.jl

@@ -27,6 +27,11 @@ struct NonlinearSystem <: AbstractSystem
     ps::Vector
     observed::Vector{Equation}
     """
+    Jacobian matrix. Note: this field will not be defined until
+    [`calculate_jacobian`](@ref) is called on the system.
+    """
+    jac::RefValue{Any}
+    """
     Name: the name of the system
     """
     name::Symbol
@@ -61,9 +66,10 @@ function NonlinearSystem(eqs, states, ps;
     if !(isempty(default_u0) && isempty(default_p))
         Base.depwarn("`default_u0` and `default_p` are deprecated. Use `defaults` instead.", :NonlinearSystem, force=true)
     end
+    jac = RefValue{Any}(Matrix{Num}(undef, 0, 0))
     defaults = todict(defaults)
     defaults = Dict(value(k) => value(v) for (k, v) in pairs(defaults))
-    NonlinearSystem(eqs, value.(states), value.(ps), observed, name, systems, defaults, nothing, connection_type)
+    NonlinearSystem(eqs, value.(states), value.(ps), observed, jac, name, systems, defaults, nothing, connection_type)
 end
 
 function calculate_jacobian(sys::NonlinearSystem;sparse=false,simplify=false)
@@ -74,6 +80,7 @@ function calculate_jacobian(sys::NonlinearSystem;sparse=false,simplify=false)
     else
         jac = jacobian(rhs, vals, simplify=simplify)
     end
+    get_jac(sys)[] = jac
     return jac
 end
 

test/nonlinearsystem.jl

@@ -107,3 +107,14 @@ prob = ODEProblem(sys, [subsys.x => 1, subsys.z => 2.0], (0, 1.0), [subsys.σ=>1
 sol = solve(prob, Rodas5())
 @test sol[subsys.x] + sol[subsys.y] - sol[subsys.z] ≈ sol[subsys.u]
 @test_throws ArgumentError convert_system(ODESystem, sys, t)
+
+@parameters t σ ρ β
+@variables x y z
+
+# Define a nonlinear system
+eqs = [0 ~ σ*(y-x),
+       0 ~ x*(ρ-z)-y,
+       0 ~ x*y - β*z]
+ns = NonlinearSystem(eqs, [x,y,z], [σ,ρ,β])
+np = NonlinearProblem(ns, [0,0,0], [1,2,3], jac=true, sparse=true)
+@test ModelingToolkit.get_jac(ns)[] isa SparseMatrixCSC