From c436ae179c11f65cbb1ac83ff0d322ef565bc9b8 Mon Sep 17 00:00:00 2001 From: ChrisRackauckas Date: Fri, 17 Oct 2025 08:41:38 -0400 Subject: [PATCH] WIP: Fix Radau methods to preserve jacobian prototype sparsity for KLU - Added is_sparse and nonzeros imports to OrdinaryDiffEqFIRK - Added SparseArrays as dependency - Modified RadauIIA3, RadauIIA5, RadauIIA9, and AdaptiveRadau cache initialization to preserve sparsity pattern from jacobian prototype for complex matrices - Added test cases to reproduce the issue Addresses #2892 where KLUFactorization fails with jacobian prototypes. The fix preserves the sparsity structure but further work may be needed to ensure KLU can properly initialize with the complex sparse matrices. Issue: Complex sparse matrices initialized with zeros prevent KLU from creating proper factorization cache, resulting in 'type Nothing has no field colptr' error. --- Project.toml | 6 ++- lib/OrdinaryDiffEqFIRK/Project.toml | 2 + .../src/OrdinaryDiffEqFIRK.jl | 3 +- lib/OrdinaryDiffEqFIRK/src/firk_caches.jl | 47 +++++++++++++++---- test_issue_2892.jl | 41 ++++++++++++++++ test_radauiia5_klu.jl | 28 +++++++++++ 6 files changed, 115 insertions(+), 12 deletions(-) create mode 100644 test_issue_2892.jl create mode 100644 test_radauiia5_klu.jl diff --git a/Project.toml b/Project.toml index 1ffdffe5a4..1dfa814726 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "OrdinaryDiffEq" uuid = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed" -authors = ["Chris Rackauckas ", "Yingbo Ma "] version = "6.102.1" +authors = ["Chris Rackauckas ", "Yingbo Ma "] [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" @@ -69,6 +69,8 @@ SciMLStructures = "53ae85a6-f571-4167-b2af-e1d143709226" SimpleNonlinearSolve = "727e6d20-b764-4bd8-a329-72de5adea6c7" SimpleUnPack = "ce78b400-467f-4804-87d8-8f486da07d0a" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" +SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5" +Sparspak = "e56a9233-b9d6-4f03-8d0f-1825330902ac" Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3" StaticArrayInterface = "0d7ed370-da01-4f52-bd93-41d350b8b718" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" @@ -173,6 +175,8 @@ SciMLOperators = "1.8" SciMLStructures = "1.7" SimpleNonlinearSolve = "2.7" SimpleUnPack = "1.1" +SparseConnectivityTracer = "1.1.1" +Sparspak = "0.3.14" Static = "1.2" StaticArrayInterface = "1.8" StaticArrays = "1.9.14" diff --git a/lib/OrdinaryDiffEqFIRK/Project.toml b/lib/OrdinaryDiffEqFIRK/Project.toml index 64249c963d..94b0923fc1 100644 --- a/lib/OrdinaryDiffEqFIRK/Project.toml +++ b/lib/OrdinaryDiffEqFIRK/Project.toml @@ -10,6 +10,7 @@ MuladdMacro = "46d2c3a1-f734-5fdb-9937-b9b9aeba4221" LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae" Polyester = "f517fe37-dbe3-4b94-8317-1923a5111588" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" OrdinaryDiffEqDifferentiation = "4302a76b-040a-498a-8c04-15b101fed76b" SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" OrdinaryDiffEqCore = "bbf590c4-e513-4bbe-9b18-05decba2e5d8" @@ -42,6 +43,7 @@ MuladdMacro = "0.2" LinearSolve = "3.26" Polyester = "0.7" LinearAlgebra = "1.10" +SparseArrays = "1.10" OrdinaryDiffEqDifferentiation = "1.12.0" SciMLBase = "2.99" OrdinaryDiffEqCore = "1.29.0" diff --git a/lib/OrdinaryDiffEqFIRK/src/OrdinaryDiffEqFIRK.jl b/lib/OrdinaryDiffEqFIRK/src/OrdinaryDiffEqFIRK.jl index 7ab3cffccf..4391e80cce 100644 --- a/lib/OrdinaryDiffEqFIRK/src/OrdinaryDiffEqFIRK.jl +++ b/lib/OrdinaryDiffEqFIRK/src/OrdinaryDiffEqFIRK.jl @@ -23,6 +23,7 @@ using MuladdMacro, DiffEqBase, RecursiveArrayTools, Polyester isfirk, generic_solver_docstring using SciMLOperators: AbstractSciMLOperator using LinearAlgebra: I, UniformScaling, mul!, lu +using SparseArrays: nonzeros import LinearSolve import FastBroadcast: @.. import OrdinaryDiffEqCore @@ -30,7 +31,7 @@ import OrdinaryDiffEqCore: _ode_interpolant, _ode_interpolant!, has_stiff_interp import FastPower: fastpower using OrdinaryDiffEqDifferentiation: UJacobianWrapper, build_J_W, build_jac_config, UDerivativeWrapper, calc_J!, dolinsolve, calc_J, - islinearfunction + islinearfunction, is_sparse using OrdinaryDiffEqNonlinearSolve: du_alias_or_new, Convergence, FastConvergence, NLStatus, VerySlowConvergence, Divergence, get_new_W_γdt_cutoff diff --git a/lib/OrdinaryDiffEqFIRK/src/firk_caches.jl b/lib/OrdinaryDiffEqFIRK/src/firk_caches.jl index 6f99a62cf3..1b3adb9e6e 100644 --- a/lib/OrdinaryDiffEqFIRK/src/firk_caches.jl +++ b/lib/OrdinaryDiffEqFIRK/src/firk_caches.jl @@ -98,9 +98,15 @@ function alg_cache(alg::RadauIIA3, u, rate_prototype, ::Type{uEltypeNoUnits}, recursivefill!(atmp, false) jac_config = build_jac_config(alg, f, uf, du1, uprev, u, tmp, dw12) - J, W1 = build_J_W(alg, u, uprev, p, t, dt, f, jac_config, uEltypeNoUnits, Val(true)) - W1 = similar(J, Complex{eltype(W1)}) - recursivefill!(W1, false) + J, W1_temp = build_J_W(alg, u, uprev, p, t, dt, f, jac_config, uEltypeNoUnits, Val(true)) + # For sparse matrices, preserve sparsity pattern for KLU compatibility + if is_sparse(J) + W1 = similar(J, Complex{eltype(W1_temp)}) + fill!(nonzeros(W1), false) + else + W1 = similar(J, Complex{eltype(W1_temp)}) + recursivefill!(W1, false) + end linprob = LinearProblem(W1, _vec(cubuff); u0 = _vec(dw12)) linsolve = init( @@ -239,8 +245,14 @@ function alg_cache(alg::RadauIIA5, u, rate_prototype, ::Type{uEltypeNoUnits}, if J isa AbstractSciMLOperator error("Non-concrete Jacobian not yet supported by RadauIIA5.") end - W2 = similar(J, Complex{eltype(W1)}) - recursivefill!(W2, false) + # For sparse matrices, preserve sparsity pattern for KLU compatibility + if is_sparse(J) + W2 = similar(J, Complex{eltype(W1)}) + fill!(nonzeros(W2), false) + else + W2 = similar(J, Complex{eltype(W1)}) + recursivefill!(W2, false) + end linprob = LinearProblem(W1, _vec(ubuff); u0 = _vec(dw1)) linsolve1 = init( @@ -429,10 +441,18 @@ function alg_cache(alg::RadauIIA9, u, rate_prototype, ::Type{uEltypeNoUnits}, if J isa AbstractSciMLOperator error("Non-concrete Jacobian not yet supported by RadauIIA5.") end - W2 = similar(J, Complex{eltype(W1)}) - W3 = similar(J, Complex{eltype(W1)}) - recursivefill!(W2, false) - recursivefill!(W3, false) + # For sparse matrices, preserve sparsity pattern for KLU compatibility + if is_sparse(J) + W2 = similar(J, Complex{eltype(W1)}) + W3 = similar(J, Complex{eltype(W1)}) + fill!(nonzeros(W2), false) + fill!(nonzeros(W3), false) + else + W2 = similar(J, Complex{eltype(W1)}) + W3 = similar(J, Complex{eltype(W1)}) + recursivefill!(W2, false) + recursivefill!(W3, false) + end linprob = LinearProblem(W1, _vec(ubuff); u0 = _vec(dw1)) linsolve1 = init( @@ -638,8 +658,15 @@ function alg_cache(alg::AdaptiveRadau, u, rate_prototype, ::Type{uEltypeNoUnits} error("Non-concrete Jacobian not yet supported by AdaptiveRadau.") end + # For sparse matrices, preserve sparsity pattern for KLU compatibility W2 = [similar(J, Complex{eltype(W1)}) for _ in 1:((max_stages - 1) ÷ 2)] - recursivefill!.(W2, false) + if is_sparse(J) + for W in W2 + fill!(nonzeros(W), false) + end + else + recursivefill!.(W2, false) + end linprob = LinearProblem(W1, _vec(ubuff); u0 = _vec(dw1)) linsolve1 = init( diff --git a/test_issue_2892.jl b/test_issue_2892.jl new file mode 100644 index 0000000000..5d9e397cf8 --- /dev/null +++ b/test_issue_2892.jl @@ -0,0 +1,41 @@ +using OrdinaryDiffEq +using OrdinaryDiffEqFIRK +using ADTypes +using SparseConnectivityTracer +using LinearSolve +using Sparspak + +function test_sparse(ode_solver) + function f(du, u, p, t) + du .= [u[1], u[2]] + end + + u0 = [1.0, 2.0] + p = () + du0 = similar(u0) + jac_prototype = float.(ADTypes.jacobian_sparsity( + (du, u) -> f(du, u, p, 0.0), + du0, + u0, TracerSparsityDetector())) + + ode_fun = ODEFunction(f, jac_prototype=jac_prototype) + prob = ODEProblem(ode_fun, u0, (0, 10)) + sol = solve(prob, ode_solver) + return sol +end + +println("Testing AdaptiveRadau with LUFactorization...") +test_sparse(AdaptiveRadau(;linsolve=LUFactorization())) # Success +println("Success!") + +println("Testing AdaptiveRadau with SparspakFactorization...") +test_sparse(AdaptiveRadau(;linsolve=SparspakFactorization())) # Success +println("Success!") + +println("Testing QNDF with KLUFactorization...") +test_sparse(QNDF(;linsolve=KLUFactorization())) # Success +println("Success!") + +println("Testing AdaptiveRadau with KLUFactorization...") +test_sparse(AdaptiveRadau(;linsolve=KLUFactorization())) # Error +println("Success!") diff --git a/test_radauiia5_klu.jl b/test_radauiia5_klu.jl new file mode 100644 index 0000000000..0a52c29d78 --- /dev/null +++ b/test_radauiia5_klu.jl @@ -0,0 +1,28 @@ +using OrdinaryDiffEq +using OrdinaryDiffEqFIRK +using ADTypes +using SparseConnectivityTracer +using LinearSolve + +function test_sparse(ode_solver) + function f(du, u, p, t) + du .= [u[1], u[2]] + end + + u0 = [1.0, 2.0] + p = () + du0 = similar(u0) + jac_prototype = float.(ADTypes.jacobian_sparsity( + (du, u) -> f(du, u, p, 0.0), + du0, + u0, TracerSparsityDetector())) + + ode_fun = ODEFunction(f, jac_prototype=jac_prototype) + prob = ODEProblem(ode_fun, u0, (0, 10)) + sol = solve(prob, ode_solver) + return sol +end + +println("Testing RadauIIA5 with KLUFactorization...") +test_sparse(RadauIIA5(;linsolve=KLUFactorization())) +println("Success!")