From 0d34b8a79cf89da8dd0b1cf1aba853b30f7d8d86 Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Wed, 29 Jul 2026 16:34:18 -0400 Subject: [PATCH] Preserve structured Jacobian cache storage Co-Authored-By: Chris Rackauckas --- lib/DiffEqBase/test/downstream/Project.toml | 2 + .../downstream/default_linsolve_structure.jl | 60 ++++++++++++------- .../src/OrdinaryDiffEqDifferentiation.jl | 24 +++++++- .../src/derivative_utils.jl | 4 ++ 4 files changed, 67 insertions(+), 23 deletions(-) diff --git a/lib/DiffEqBase/test/downstream/Project.toml b/lib/DiffEqBase/test/downstream/Project.toml index d1b907b8c2a..1333fab34d0 100644 --- a/lib/DiffEqBase/test/downstream/Project.toml +++ b/lib/DiffEqBase/test/downstream/Project.toml @@ -9,6 +9,7 @@ FlexUnits = "76e01b6b-c995-4ce6-8559-91e72a3d4e95" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" GTPSA = "b27dd330-f138-47c5-815b-40db9dd9b6e8" LabelledArrays = "2ee39098-c373-598a-b85f-a56591580800" +LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae" Measurements = "eff96d63-e80a-5855-80a2-b1b0885c5ab7" MultiScaleArrays = "f9640e96-87f6-5992-9c3b-0743c6a49ffa" ODEProblemLibrary = "fdc4e326-1af4-4b90-96e7-779fcce2daa5" @@ -46,6 +47,7 @@ DifferentiationInterface = "0.7" Enzyme = "0.13.100" ForwardDiff = "0.10, 1" GTPSA = "1.4" +LinearSolve = "5.1" MultiScaleArrays = "1.8" OrdinaryDiffEq = "7" SciMLSensitivity = "7" diff --git a/lib/DiffEqBase/test/downstream/default_linsolve_structure.jl b/lib/DiffEqBase/test/downstream/default_linsolve_structure.jl index a45e7074507..a70210532a2 100644 --- a/lib/DiffEqBase/test/downstream/default_linsolve_structure.jl +++ b/lib/DiffEqBase/test/downstream/default_linsolve_structure.jl @@ -1,33 +1,49 @@ using LinearAlgebra, OrdinaryDiffEq, Test, ADTypes +using LinearSolve: LUFactorization using OrdinaryDiffEqRosenbrock: Rosenbrock23 f = (du, u, p, t) -> du .= u ./ t -jac = (J, u, p, t) -> (J[1, 1] = 1 / t; J[2, 2] = 1 / t; J[1, 2] = 0; J[2, 1] = 0) +function jac(J, u, p, t) + fill!(J, zero(eltype(J))) + for i in axes(J, 1) + J[i, i] = 1 / t + end + return nothing +end -jp_diag = Diagonal(zeros(2)) -fun = ODEFunction(f; jac = jac, jac_prototype = jp_diag) -prob = ODEProblem(fun, ones(2), (1.0, 10.0)) -sol = solve(prob, Rosenbrock23()) -@test sol.u[end] ≈ [10.0, 10.0] -@test length(sol.t) < 60 +jp_diag = Diagonal(zeros(3)) +structured_prototypes = ( + (jp_diag, Diagonal(ones(3))), + (Bidiagonal(zeros(3), zeros(2), :U), Bidiagonal(ones(3), ones(2), :U)), + (Tridiagonal(jp_diag), Tridiagonal(ones(2), ones(3), ones(2))), + (SymTridiagonal(zeros(3), zeros(2)), SymTridiagonal(ones(3), ones(2))), + (UpperTriangular(zeros(3, 3)), UpperTriangular(ones(3, 3))), +) -sol = solve(prob, Rosenbrock23(autodiff = AutoFiniteDiff())) -@test sol.u[end] ≈ [10.0, 10.0] -@test length(sol.t) < 60 +for (jp, expected_cache) in structured_prototypes + fun = ODEFunction(f; jac = jac, jac_prototype = jp) + prob = ODEProblem(fun, ones(3), (1.0, 10.0)) -jp = Tridiagonal(jp_diag) -fun = ODEFunction(f; jac = jac, jac_prototype = jp) -prob = ODEProblem(fun, ones(2), (1.0, 10.0)) + alg = jp isa UpperTriangular ? Rosenbrock23(linsolve = LUFactorization()) : Rosenbrock23() + integ = init(prob, alg) + @test integ.cache.J isa typeof(jp) + @test integ.cache.W isa typeof(jp) + @test integ.cache.J == expected_cache + @test integ.cache.W == expected_cache -sol = solve(prob, Rosenbrock23()) -@test sol.u[end] ≈ [10.0, 10.0] -@test length(sol.t) < 60 + if jp isa Union{Diagonal, Tridiagonal} + sol = solve!(integ) + @test sol.u[end] ≈ [10.0, 10.0, 10.0] + @test length(sol.t) < 60 -sol = solve(prob, Rosenbrock23(autodiff = AutoFiniteDiff())) -@test sol.u[end] ≈ [10.0, 10.0] -@test length(sol.t) < 60 + integ = init(prob, Rosenbrock23(autodiff = AutoFiniteDiff())) + sol = solve!(integ) + @test sol.u[end] ≈ [10.0, 10.0, 10.0] + @test length(sol.t) < 60 + end +end #= -jp = SymTridiagonal(jp_diag) +jp = SymTridiagonal(Diagonal(zeros(2))) fun = ODEFunction(f; jac=jac, jac_prototype=jp) prob = ODEProblem(fun,ones(2),(1.0,10.0)) sol = solve(prob,Rosenbrock23()) @@ -40,7 +56,7 @@ sol = solve(prob,Rosenbrock23()) # LoadError: ArgumentError: broadcasted assignment breaks symmetry between locations (1, 2) and (2, 1) @test_broken begin - local jp = Hermitian(jp_diag) + local jp = Hermitian(Diagonal(zeros(2))) local fun = ODEFunction(f; jac = jac, jac_prototype = jp) local prob = ODEProblem(fun, ones(2), (1.0, 10.0)) local sol = solve(prob, Rosenbrock23(autodiff = AutoFiniteDiff())) @@ -49,7 +65,7 @@ sol = solve(prob,Rosenbrock23()) end @test_broken begin - local jp = Symmetric(jp_diag) + local jp = Symmetric(Diagonal(zeros(2))) local fun = ODEFunction(f; jac = jac, jac_prototype = jp) local prob = ODEProblem(fun, ones(2), (1.0, 10.0)) local sol = solve(prob, Rosenbrock23(autodiff = AutoFiniteDiff())) diff --git a/lib/OrdinaryDiffEqDifferentiation/src/OrdinaryDiffEqDifferentiation.jl b/lib/OrdinaryDiffEqDifferentiation/src/OrdinaryDiffEqDifferentiation.jl index 67721a69bf6..0415031934a 100644 --- a/lib/OrdinaryDiffEqDifferentiation/src/OrdinaryDiffEqDifferentiation.jl +++ b/lib/OrdinaryDiffEqDifferentiation/src/OrdinaryDiffEqDifferentiation.jl @@ -11,7 +11,8 @@ import FunctionWrappersWrappers import DiffEqBase import LinearAlgebra -import LinearAlgebra: Diagonal, I, UniformScaling, diagind, opnorm +import LinearAlgebra: AbstractTriangular, Diagonal, I, SymTridiagonal, UniformScaling, + UpperHessenberg, diagind, opnorm import LinearAlgebra: LowerTriangular import ArrayInterface @@ -71,6 +72,27 @@ spzeros(args...) = error("SparseArrays extension not loaded. Please load SparseA get_nzval(A) = error("SparseArrays extension not loaded. Please load SparseArrays to use sparse matrix functionality.") set_all_nzval!(A, val) = error("SparseArrays extension not loaded. Please load SparseArrays to use sparse matrix functionality.") +function fill_structural_nonzeros!(A, value) + storage = parent(A) + if storage !== A + fill!(storage, value) + return A + end + + rows, cols = ArrayInterface.findstructralnz(A) + for i in 1:length(rows) + ArrayInterface.allowed_setindex!(A, value, rows[i], cols[i]) + end + return A +end + +# SymTridiagonal does not permit off-diagonal setindex!, even at stored positions. +function fill_structural_nonzeros!(A::SymTridiagonal, value) + n = size(A, 1) + source = SymTridiagonal(fill(value, n), fill(value, max(n - 1, 0))) + return copyto!(A, source) +end + include("alg_utils.jl") include("linsolve_utils.jl") include("derivative_utils.jl") diff --git a/lib/OrdinaryDiffEqDifferentiation/src/derivative_utils.jl b/lib/OrdinaryDiffEqDifferentiation/src/derivative_utils.jl index ce072d3f8df..501794d0c7e 100644 --- a/lib/OrdinaryDiffEqDifferentiation/src/derivative_utils.jl +++ b/lib/OrdinaryDiffEqDifferentiation/src/derivative_utils.jl @@ -1227,6 +1227,10 @@ function build_J_W( if is_sparse(J) set_all_nzval!(J, one(eltype(J))) set_all_nzval!(W, one(eltype(W))) + elseif J isa Union{AbstractTriangular, UpperHessenberg} || + ArrayInterface.isstructured(J) || ArrayInterface.has_sparsestruct(J) + fill_structural_nonzeros!(J, one(eltype(J))) + fill_structural_nonzeros!(W, one(eltype(W))) else fill!(J, one(eltype(J))) fill!(W, one(eltype(W)))