Merge pull request #1036 from SciML/myb/jac

Compute Jacobian sparsity when u0 is not provided

Author: YingboMa

src/systems/diffeqs/abstractodesystem.jl

@@ -19,7 +19,10 @@ end
 
 function calculate_jacobian(sys::AbstractODESystem;
                             sparse=false, simplify=false)
-    isempty(get_jac(sys)[]) || return get_jac(sys)[]  # use cached Jacobian, if possible
+    cache = get_jac(sys)[]
+    if cache isa Tuple && cache[2] == (sparse, simplify)
+        return cache[1]
+    end
     rhs = [eq.rhs for eq ∈ equations(sys)]
 
     iv = get_iv(sys)
@@ -31,7 +34,7 @@ function calculate_jacobian(sys::AbstractODESystem;
         jac = jacobian(rhs, dvs, simplify=simplify)
     end
 
-    get_jac(sys)[] = jac  # cache Jacobian
+    get_jac(sys)[] = jac, (sparse, simplify) # cache Jacobian
     return jac
 end
 
@@ -217,13 +220,22 @@ function DiffEqBase.ODEFunction{iip}(sys::AbstractODESystem, dvs = states(sys),
         end
     end
 
-    uElType = eltype(u0)
+    jac_prototype = if sparse
+        uElType = u0 === nothing ? Float64 : eltype(u0)
+        if jac
+            similar(calculate_jacobian(sys, sparse=sparse), uElType)
+        else
+            similar(jacobian_sparsity(sys), uElType)
+        end
+    else
+        nothing
+    end
     ODEFunction{iip}(
                      f,
                      jac = _jac === nothing ? nothing : _jac,
                      tgrad = _tgrad === nothing ? nothing : _tgrad,
                      mass_matrix = _M,
-                     jac_prototype = (!isnothing(u0) && sparse) ? (!jac ? similar(jacobian_sparsity(sys),uElType) : SparseArrays.sparse(similar(get_jac(sys)[],uElType))) : nothing,
+                     jac_prototype = jac_prototype,
                      syms = Symbol.(states(sys)),
                      indepsym = Symbol(independent_variable(sys)),
                      observed = observedfun,

src/systems/nonlinear/nonlinearsystem.jl

@@ -76,15 +76,20 @@ function NonlinearSystem(eqs, states, ps;
     NonlinearSystem(eqs, value.(states), value.(ps), observed, jac, name, systems, defaults, nothing, connection_type)
 end
 
-function calculate_jacobian(sys::NonlinearSystem;sparse=false,simplify=false)
+function calculate_jacobian(sys::NonlinearSystem; sparse=false, simplify=false)
+    cache = get_jac(sys)[]
+    if cache isa Tuple && cache[2] == (sparse, simplify)
+        return cache[1]
+    end
+
     rhs = [eq.rhs for eq ∈ equations(sys)]
     vals = [dv for dv in states(sys)]
     if sparse
         jac = sparsejacobian(rhs, vals, simplify=simplify)
     else
         jac = jacobian(rhs, vals, simplify=simplify)
     end
-    get_jac(sys)[] = jac
+    get_jac(sys)[] = jac, (sparse, simplify)
     return jac
 end
 
@@ -170,7 +175,7 @@ function DiffEqBase.NonlinearFunction{iip}(sys::NonlinearSystem, dvs = states(sy
 
     NonlinearFunction{iip}(f,
                      jac = _jac === nothing ? nothing : _jac,
-                     jac_prototype = sparse ? similar(sys.jac[],Float64) : nothing,
+                     jac_prototype = sparse ? similar(calculate_jacobian(sys, sparse=sparse),Float64) : nothing,
                      syms = Symbol.(states(sys)), observed = observedfun)
 end
 

test/jacobiansparsity.jl

@@ -1,7 +1,7 @@
 using OrdinaryDiffEq, ModelingToolkit, Test, SparseArrays
 
-N = 32
-const xyd_brusselator = range(0,stop=1,length=N)
+N = 3
+xyd_brusselator = range(0,stop=1,length=N)
 brusselator_f(x, y, t) = (((x-0.3)^2 + (y-0.6)^2) <= 0.1^2) * (t >= 1.1) * 5.
 limit(a, N) = ModelingToolkit.ifelse(a == N+1, 1, ModelingToolkit.ifelse(a == 0, N, a))
 function brusselator_2d_loop(du, u, p, t)
@@ -40,11 +40,12 @@ sys = modelingtoolkitize(prob_ode_brusselator_2d)
 
 # test sparse jacobian pattern only.
 prob = ODEProblem(sys, u0, (0, 11.5), sparse=true, jac=false)
-@test findnz(Symbolics.jacobian_sparsity(map(x->x.rhs, equations(sys)), states(sys)))[1:2] == findnz(prob.f.jac_prototype)[1:2]
+JP = prob.f.jac_prototype
+@test findnz(Symbolics.jacobian_sparsity(map(x->x.rhs, equations(sys)), states(sys)))[1:2] == findnz(JP)[1:2]
 
 # test sparse jacobian
 prob = ODEProblem(sys, u0, (0, 11.5), sparse=true, jac=true)
-@test findnz(calculate_jacobian(sys))[1:2] == findnz(prob.f.jac_prototype)[1:2]
+@test findnz(calculate_jacobian(sys, sparse=true))[1:2] == findnz(prob.f.jac_prototype)[1:2]
 
 # test when not sparse
 prob = ODEProblem(sys, u0, (0, 11.5), sparse=false, jac=true)
@@ -55,10 +56,12 @@ prob = ODEProblem(sys, u0, (0, 11.5), sparse=false, jac=false)
 
 # test when u0 is nothing
 f = DiffEqBase.ODEFunction(sys, u0=nothing, sparse=true, jac=true)
-@test f.jac_prototype == nothing
+@test findnz(f.jac_prototype)[1:2] == findnz(JP)[1:2]
+@test eltype(f.jac_prototype) == Float64
 
 f = DiffEqBase.ODEFunction(sys, u0=nothing, sparse=true, jac=false)
-@test f.jac_prototype == nothing
+@test findnz(f.jac_prototype)[1:2] == findnz(JP)[1:2]
+@test eltype(f.jac_prototype) == Float64
 
 # test when u0 is not Float64
 u0 = similar(init_brusselator_2d(xyd_brusselator), Float32)

test/nonlinearsystem.jl

@@ -117,21 +117,21 @@ eqs = [0 ~ σ*(y-x),
        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
+@test calculate_jacobian(ns, sparse=true) isa SparseMatrixCSC
 
 # issue #819
 @testset "Combined system name collisions" begin
        function makesys(name)
            @parameters a
            @variables x f
-   
+
            NonlinearSystem([0 ~ -a*x + f],[x,f],[a], name=name)
        end
-   
+
        function issue819()
            sys1 = makesys(:sys1)
            sys2 = makesys(:sys1)
            @test_throws ArgumentError NonlinearSystem([sys2.f ~ sys1.x, sys1.f ~ 0],[],[], systems=[sys1, sys2])
        end
        issue819()
-end
+end