Merge pull request #4037 from hersle/sparse_performance

Improve performance with sparse Jacobians

Author: AayushSabharwal

benchmark/benchmarks.jl

@@ -76,3 +76,18 @@ prob = ODEProblem(model, u0, tspan)
 large_param_init["ODEProblem"] = @benchmarkable ODEProblem($model, $u0, $tspan)
 
 large_param_init["init"] = @benchmarkable init($prob)
+
+sparse_analytical_jacobian = SUITE["sparse_analytical_jacobian"]
+
+eqs = [D(x[i]) ~ prod(x[j] for j in 1:N if (i+j) % 3 == 0) for i in 1:N]
+@mtkcompile model = System(eqs, t)
+u0 = collect(x .=> 1.0)
+tspan = (0.0, 1.0)
+jac = true
+sparse = true
+prob = ODEProblem(model, u0, tspan; jac, sparse)
+out = similar(prob.f.jac_prototype)
+
+sparse_analytical_jacobian["ODEProblem"] = @benchmarkable ODEProblem($model, $u0, $tspan; jac, sparse)
+sparse_analytical_jacobian["f_oop"] = @benchmarkable $(prob.f.jac.f_oop)($(prob.u0), $(prob.p), $(first(tspan)))
+sparse_analytical_jacobian["f_iip"] = @benchmarkable $(prob.f.jac.f_iip)($out, $(prob.u0), $(prob.p), $(first(tspan)))

src/systems/codegen.jl

@@ -186,16 +186,14 @@ function calculate_jacobian(sys::System;
     if sparse
         jac = sparsejacobian(rhs, dvs; simplify)
         if get_iv(sys) !== nothing
-            W_s = W_sparsity(sys)
-            (Is, Js, Vs) = findnz(W_s)
-            # Add nonzeros of W as non-structural zeros of the Jacobian (to ensure equal
-            # results for oop and iip Jacobian)
-            for (i, j) in zip(Is, Js)
-                iszero(jac[i, j]) && begin
-                    jac[i, j] = 1
-                    jac[i, j] = 0
-                end
-            end
+            # Add nonzeros of W as non-structural zeros of the Jacobian
+            # (to ensure equal results for oop and iip Jacobian)
+            JIs, JJs, JVs = findnz(jac)
+            WIs, WJs, _ = findnz(W_sparsity(sys))
+            append!(JIs, WIs) # explicitly put all W's indices also in J,
+            append!(JJs, WJs) # even if it duplicates some indices
+            append!(JVs, zeros(eltype(JVs), length(WIs))) # add zero
+            jac = SparseArrays.sparse(JIs, JJs, JVs) # values at duplicate indices are summed; not overwritten
         end
     else
         jac = jacobian(rhs, dvs; simplify)
@@ -213,21 +211,23 @@ Generate the jacobian function for the equations of a [`System`](@ref).
 
 $GENERATE_X_KWARGS
 - `simplify`, `sparse`: Forwarded to [`calculate_jacobian`](@ref).
+- `checkbounds`: Whether to check correctness of indices at runtime if `sparse`.
+  Also forwarded to `build_function_wrapper`.
 
 All other keyword arguments are forwarded to [`build_function_wrapper`](@ref).
 """
 function generate_jacobian(sys::System;
         simplify = false, sparse = false, eval_expression = false,
         eval_module = @__MODULE__, expression = Val{true}, wrap_gfw = Val{false},
-        kwargs...)
+        checkbounds = false, kwargs...)
     dvs = unknowns(sys)
     jac = calculate_jacobian(sys; simplify, sparse, dvs)
     p = reorder_parameters(sys)
     t = get_iv(sys)
-    if t === nothing
-        wrap_code = (identity, identity)
+    if t !== nothing && sparse && checkbounds
+        wrap_code = assert_jac_length_header(sys) # checking sparse J indices at runtime is expensive for large systems
     else
-        wrap_code = sparse ? assert_jac_length_header(sys) : (identity, identity)
+        wrap_code = (identity, identity)
     end
     args = (dvs, p...)
     nargs = 2
@@ -236,7 +236,7 @@ function generate_jacobian(sys::System;
         nargs = 3
     end
     res = build_function_wrapper(sys, jac, args...; wrap_code, expression = Val{true},
-        expression_module = eval_module, kwargs...)
+        expression_module = eval_module, checkbounds, kwargs...)
     return maybe_compile_function(
         expression, wrap_gfw, (2, nargs, is_split(sys)), res; eval_expression, eval_module)
 end
@@ -328,12 +328,14 @@ Generate the `W = γ * M + J` function for the equations of a [`System`](@ref).
 
 $GENERATE_X_KWARGS
 - `simplify`, `sparse`: Forwarded to [`calculate_jacobian`](@ref).
+- `checkbounds`: Whether to check correctness of indices at runtime if `sparse`.
+  Also forwarded to `build_function_wrapper`.
 
 All other keyword arguments are forwarded to [`build_function_wrapper`](@ref).
 """
 function generate_W(sys::System;
         simplify = false, sparse = false, expression = Val{true}, wrap_gfw = Val{false},
-        eval_expression = false, eval_module = @__MODULE__, kwargs...)
+        eval_expression = false, eval_module = @__MODULE__, checkbounds = false, kwargs...)
     dvs = unknowns(sys)
     ps = parameters(sys; initial_parameters = true)
     M = calculate_massmatrix(sys; simplify)
@@ -343,13 +345,15 @@ function generate_W(sys::System;
     J = calculate_jacobian(sys; simplify, sparse, dvs)
     W = W_GAMMA * M + J
     t = get_iv(sys)
-    if t !== nothing
-        wrap_code = sparse ? assert_jac_length_header(sys) : (identity, identity)
+    if t !== nothing && sparse && checkbounds
+        wrap_code = assert_jac_length_header(sys)
+    else
+        wrap_code = (identity, identity)
     end
 
     p = reorder_parameters(sys, ps)
     res = build_function_wrapper(sys, W, dvs, p..., W_GAMMA, t; wrap_code,
-        p_end = 1 + length(p), kwargs...)
+        p_end = 1 + length(p), checkbounds, kwargs...)
     return maybe_compile_function(
         expression, wrap_gfw, (2, 4, is_split(sys)), res; eval_expression, eval_module)
 end

test/jacobiansparsity.jl

@@ -1,4 +1,4 @@
-using ModelingToolkit, SparseArrays, OrdinaryDiffEq, DiffEqBase
+using ModelingToolkit, SparseArrays, OrdinaryDiffEq, DiffEqBase, BenchmarkTools
 
 N = 3
 xyd_brusselator = range(0, stop = 1, length = N)
@@ -58,6 +58,8 @@ prob = ODEProblem(sys, u0, (0, 11.5), sparse = true, jac = true)
 #@test_nowarn solve(prob, Rosenbrock23())
 @test findnz(calculate_jacobian(sys, sparse = true))[1:2] ==
       findnz(prob.f.jac_prototype)[1:2]
+out = similar(prob.f.jac_prototype)
+@test (@ballocated $(prob.f.jac.f_iip)($out, $(prob.u0), $(prob.p), 0.0)) == 0 # should not allocate
 
 # test when not sparse
 prob = ODEProblem(sys, u0, (0, 11.5), sparse = false, jac = true)
@@ -75,6 +77,12 @@ f = DiffEqBase.ODEFunction(sys, u0 = nothing, sparse = true, jac = false)
 @test findnz(f.jac_prototype)[1:2] == findnz(JP)[1:2]
 @test eltype(f.jac_prototype) == Float64
 
+# test sparsity index pattern checking
+f = DiffEqBase.ODEFunction(sys, u0 = nothing, sparse = true, jac = true, checkbounds = true)
+out = sparse([1.0 0.0; 0.0 1.0]) # choose a wrong size on purpose
+@test size(out) != size(f.jac_prototype) # check that the size is indeed wrong
+@test_throws AssertionError f.jac.f_iip(out, u0, p, 0.0) # check that we get an error
+
 # test when u0 is not Float64
 u0 = similar(init_brusselator_2d(xyd_brusselator), Float32)
 prob_ode_brusselator_2d = ODEProblem(brusselator_2d_loop,
@@ -100,7 +108,7 @@ prob = ODEProblem(sys, u0, (0, 11.5), sparse = true, jac = true)
     u0 = [x => 1, y => 0]
     prob = ODEProblem(
         pend, [u0; [g => 1]], (0, 11.5), guesses = [λ => 1], sparse = true, jac = true)
-    jac, jac! = generate_jacobian(pend; expression = Val{false}, sparse = true)
+    jac, jac! = generate_jacobian(pend; expression = Val{false}, sparse = true, checkbounds = true)
     jac_prototype = ModelingToolkit.jacobian_sparsity(pend)
     W_prototype = ModelingToolkit.W_sparsity(pend)
     @test nnz(W_prototype) == nnz(jac_prototype) + 2
@@ -113,7 +121,7 @@ prob = ODEProblem(sys, u0, (0, 11.5), sparse = true, jac = true)
     t = 0.0
     @test_throws AssertionError jac!(similar(jac_prototype, Float64), u, p, t)
 
-    W, W! = generate_W(pend; expression = Val{false}, sparse = true)
+    W, W! = generate_W(pend; expression = Val{false}, sparse = true, checkbounds = true)
     γ = 0.1
     M = sparse(calculate_massmatrix(pend))
     @test_throws AssertionError W!(similar(jac_prototype, Float64), u, p, γ, t)