From 1762fd93299d51a375b7fbdac319b49af58a3a5a Mon Sep 17 00:00:00 2001 From: Herman Sletmoen Date: Fri, 28 Nov 2025 17:42:10 +0100 Subject: [PATCH 1/8] Add W indices to sparse J to avoid iszero check --- src/systems/codegen.jl | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/src/systems/codegen.jl b/src/systems/codegen.jl index e0ceb729e2..57614d5550 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) From 4d7e17831f2c6d576edeaa4a01209ea6760c092e Mon Sep 17 00:00:00 2001 From: Herman Sletmoen Date: Sat, 29 Nov 2025 11:21:21 +0100 Subject: [PATCH 2/8] Check indices in sparse iip Jacobian function only if checkbounds --- src/systems/codegen.jl | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/src/systems/codegen.jl b/src/systems/codegen.jl index 57614d5550..e797a9911e 100644 --- a/src/systems/codegen.jl +++ b/src/systems/codegen.jl @@ -217,15 +217,15 @@ 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 = true, 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 @@ -331,7 +331,7 @@ 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 = true, kwargs...) dvs = unknowns(sys) ps = parameters(sys; initial_parameters = true) M = calculate_massmatrix(sys; simplify) @@ -341,8 +341,10 @@ 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) From ac61c8a7fcecc6544c2eb2ae9d3032e8c84d0547 Mon Sep 17 00:00:00 2001 From: Herman Sletmoen Date: Sat, 29 Nov 2025 12:37:54 +0100 Subject: [PATCH 3/8] Check that checkbounds checks sparse iip Jacobian indices --- test/jacobiansparsity.jl | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/test/jacobiansparsity.jl b/test/jacobiansparsity.jl index 9133417ecc..5d53c4160a 100644 --- a/test/jacobiansparsity.jl +++ b/test/jacobiansparsity.jl @@ -75,6 +75,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, From e14cc64a12cc81b1ca8c3f6a69dd2321f7ce02c2 Mon Sep 17 00:00:00 2001 From: Herman Sletmoen Date: Sat, 29 Nov 2025 13:14:55 +0100 Subject: [PATCH 4/8] Test that sparse+analytical iip Jacobian does not allocate --- test/jacobiansparsity.jl | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/test/jacobiansparsity.jl b/test/jacobiansparsity.jl index 5d53c4160a..4e30693b59 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) From 57340c32b3a660894cae46cbc707f6aa56b94e30 Mon Sep 17 00:00:00 2001 From: Herman Sletmoen Date: Mon, 1 Dec 2025 14:30:51 +0100 Subject: [PATCH 5/8] Generate J and W with default checkbounds=false; also forward checkbounds to build_function_wrapper --- src/systems/codegen.jl | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/systems/codegen.jl b/src/systems/codegen.jl index e797a9911e..66da47c838 100644 --- a/src/systems/codegen.jl +++ b/src/systems/codegen.jl @@ -217,7 +217,7 @@ 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}, - checkbounds = true, kwargs...) + checkbounds = false, kwargs...) dvs = unknowns(sys) jac = calculate_jacobian(sys; simplify, sparse, dvs) p = reorder_parameters(sys) @@ -234,7 +234,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 @@ -331,7 +331,7 @@ 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__, checkbounds = true, kwargs...) + eval_expression = false, eval_module = @__MODULE__, checkbounds = false, kwargs...) dvs = unknowns(sys) ps = parameters(sys; initial_parameters = true) M = calculate_massmatrix(sys; simplify) @@ -349,7 +349,7 @@ function generate_W(sys::System; 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 From 3db221b9346458449d13d5b8e8d6552983b57c8e Mon Sep 17 00:00:00 2001 From: Herman Sletmoen Date: Mon, 1 Dec 2025 14:57:07 +0100 Subject: [PATCH 6/8] Document checkbounds usage in generate_W and generate_jacobian --- src/systems/codegen.jl | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/systems/codegen.jl b/src/systems/codegen.jl index 66da47c838..fbbc9b6c9e 100644 --- a/src/systems/codegen.jl +++ b/src/systems/codegen.jl @@ -211,6 +211,8 @@ 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). """ @@ -326,6 +328,8 @@ 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). """ From 9ea59064e0cb413eaa7512a9ec9a7e389079f192 Mon Sep 17 00:00:00 2001 From: Herman Sletmoen Date: Mon, 1 Dec 2025 15:30:19 +0100 Subject: [PATCH 7/8] Benchmark ODEProblem/f_iip/f_oop with sparse+analytical Jacobian --- benchmark/benchmarks.jl | 15 +++++++++++++++ 1 file changed, 15 insertions(+) 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))) From 5ad86914a7551e497a6ddcd78cb0c2fba4605ae5 Mon Sep 17 00:00:00 2001 From: Herman Sletmoen Date: Mon, 1 Dec 2025 16:49:55 +0100 Subject: [PATCH 8/8] Set checkbounds=true in tests that call iip J and W with invalid sparse matrices --- test/jacobiansparsity.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/jacobiansparsity.jl b/test/jacobiansparsity.jl index 4e30693b59..4423d534c1 100644 --- a/test/jacobiansparsity.jl +++ b/test/jacobiansparsity.jl @@ -108,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 @@ -121,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)