Preserve structured Jacobian cache storage during initialization

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"

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()))

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")

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)))