diff --git a/benchmark/benchmarks.jl b/benchmark/benchmarks.jl index ae62f6ea0a..26e6e09838 100644 --- a/benchmark/benchmarks.jl +++ b/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))) diff --git a/src/systems/codegen.jl b/src/systems/codegen.jl index e0ceb729e2..fbbc9b6c9e 100644 --- a/src/systems/codegen.jl +++ b/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 diff --git a/test/jacobiansparsity.jl b/test/jacobiansparsity.jl index 9133417ecc..4423d534c1 100644 --- a/test/jacobiansparsity.jl +++ b/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)