Support caching

dlfivefifty · dlfivefifty · commit 186a218f2f5b · 2025-02-03T22:13:33.000Z
diff --git a/Project.toml b/Project.toml
@@ -26,7 +26,7 @@ FillArrays = "1.0"
 InfiniteArrays = "0.15"
 InfiniteRandomArrays = "0.2"
 Infinities = "0.1"
-LazyArrays = "2.3"
+LazyArrays = "2.5"
 LazyBandedMatrices = "0.11"
 LinearAlgebra = "1"
 MatrixFactorizations = "3.0"
diff --git a/src/InfiniteLinearAlgebra.jl b/src/InfiniteLinearAlgebra.jl
@@ -14,7 +14,7 @@ import Base.Broadcast: BroadcastStyle, Broadcasted, broadcasted
 
 import ArrayLayouts: AbstractBandedLayout, AbstractQLayout, AdjQRPackedQLayout, CNoPivot, DenseColumnMajor, FillLayout,
                      MatLdivVec, MatLmulMat, MatLmulVec, MemoryLayout, QRPackedQLayout, RangeCumsum, TriangularLayout,
-                     TridiagonalLayout, __qr, _factorize, _qr, check_mul_axes, colsupport,
+                     TridiagonalLayout, _qr_layout, factorize_layout, qr_layout, check_mul_axes, colsupport,
                      diagonaldata, ldiv!, lmul!, mul, mulreduce, reflector!, reflectorApply!,
                      rowsupport, sub_materialize, subdiagonaldata, sublayout, supdiagonaldata, transposelayout,
                      triangulardata, triangularlayout, zero!, materialize!
@@ -40,7 +40,7 @@ import LazyArrays: AbstractCachedMatrix, AbstractCachedVector, AbstractLazyLayou
                    CachedArray, CachedLayout, CachedMatrix, CachedVector, LazyArrayStyle, LazyLayout,
                    LazyLayouts, LazyMatrix, LazyVector, AbstractPaddedLayout, PaddedColumns, _broadcast_sub_arguments,
                    applybroadcaststyle, applylayout, arguments, cacheddata, paddeddata, resizedata!, simplifiable,
-                   simplify, islazy, islazy_layout, cache_getindex
+                   simplify, islazy, islazy_layout, cache_getindex, cache_layout
 
 import LazyBandedMatrices: AbstractLazyBandedBlockBandedLayout, AbstractLazyBandedLayout, ApplyBandedLayout, BlockVec,
                            BroadcastBandedLayout, KronTravBandedBlockBandedLayout, LazyBandedLayout,
diff --git a/src/banded/tridiagonalconjugation.jl b/src/banded/tridiagonalconjugation.jl
@@ -28,7 +28,11 @@ function upper_mul_tri_triview!(UX::Tridiagonal, U::BandedMatrix, X::Tridiagonal
 end
 
 # populate first row of UX with UX
-function initiate_upper_mul_tri_triview!(UX, U, X)
+
+initiate_upper_mul_tri_triview!(UX, U::UpperTriangular, X) = initiate_upper_mul_tri_triview!(UX, parent(U), X)
+initiate_upper_mul_tri_triview!(UX, U::CachedMatrix, X) = initiate_upper_mul_tri_triview!(UX, U.data, X)
+
+function initiate_upper_mul_tri_triview!(UX, U::BandedMatrix, X)
     Xdl, Xd, Xdu = X.dl, X.d, X.du
     UXdl, UXd, UXdu = UX.dl, UX.d, UX.du
     Udat = U.data
@@ -46,7 +50,12 @@ function initiate_upper_mul_tri_triview!(UX, U, X)
 end
 
 # fills in the rows kr of UX
-function main_upper_mul_tri_triview!(UX, U, X, kr, bₖ=X.du[kr[1]-1], aₖ=X.d[kr[1]], cₖ=X.dl[kr[1]], cₖ₋₁=X.du[kr[1]-1])
+function main_upper_mul_tri_triview!(UX, U::CachedMatrix, X, kr, kwds...)
+    resizedata!(U, kr[end], kr[end]+2)
+    main_upper_mul_tri_triview!(UX, U.data, X, kr, kwds...)
+end
+
+function main_upper_mul_tri_triview!(UX, U::BandedMatrix, X, kr, bₖ=X.du[kr[1]-1], aₖ=X.d[kr[1]], cₖ=X.dl[kr[1]], cₖ₋₁=X.dl[kr[1]-1])
     Xdl, Xd, Xdu = X.dl, X.d, X.du
     UXdl, UXd, UXdu = UX.dl, UX.d, UX.du
     Udat = U.data
@@ -109,8 +118,14 @@ function tri_mul_invupper_triview!(Y::Tridiagonal, X::Tridiagonal, R::BandedMatr
     finalize_tri_mul_invupper_triview!(Y, X, R, n, Rₖₖ, Rₖₖ₊₁)
 end
 
-# populate first row of X/R
-function initiate_tri_mul_invupper_triview!(Y, X, R)
+# partially-populate first row of X/R
+# Ydu[k] is updated below
+function initiate_tri_mul_invupper_triview!(Y, X, R::CachedMatrix)
+    resizedata!(R, 1, 2)
+    initiate_tri_mul_invupper_triview!(Y, X, R.data)
+end
+
+function initiate_tri_mul_invupper_triview!(Y, X, R::BandedMatrix)
     Xdl, Xd, Xdu = X.dl, X.d, X.du
     Ydl, Yd, Ydu = Y.dl, Y.d, Y.du
     Rdat = R.data
@@ -120,15 +135,21 @@ function initiate_tri_mul_invupper_triview!(Y, X, R)
     k = 1
     aₖ,bₖ = Xd[k], Xdu[k]
     Rₖₖ,Rₖₖ₊₁ = Rdat[u+1,k], Rdat[u,k+1] # R[1,1], R[1,2]
+    
     Yd[k] = aₖ/Rₖₖ
     Ydu[k] = bₖ - aₖ * Rₖₖ₊₁/Rₖₖ
 
     Y, Rₖₖ, Rₖₖ₊₁
 end
 
 
-# populate rows kr of X/R
-function main_tri_mul_invupper_triview!(Y::Tridiagonal, X::Tridiagonal, R::BandedMatrix, kr, Rₖₖ=R[first(kr),first(kr)], Rₖₖ₊₁=R[first(kr),first(kr)+1])
+# populate rows kr of X/R. Ydu[k] is wrong until next run.
+function main_tri_mul_invupper_triview!(Y::Tridiagonal, X::Tridiagonal, R::CachedMatrix, kr, kwds...)
+    resizedata!(R, kr[end], kr[end]+1)
+    main_tri_mul_invupper_triview!(Y, X, R.data, kr, kwds...)
+end
+
+function main_tri_mul_invupper_triview!(Y::Tridiagonal, X::Tridiagonal, R::BandedMatrix, kr, Rₖₖ=R[first(kr)-1,first(kr)-1], Rₖₖ₊₁=R[first(kr)-1,first(kr)])
     Xdl, Xd, Xdu = X.dl, X.d, X.du
     Ydl, Yd, Ydu = Y.dl, Y.d, Y.du
     Rdat = R.data
@@ -187,11 +208,18 @@ function TridiagonalConjugationData(U, X, V)
     n_init = 100
     UX = Tridiagonal(Vector{T}(undef, n_init-1), Vector{T}(undef, n_init), Vector{T}(undef, n_init-1))
     Y = Tridiagonal(Vector{T}(undef, n_init-1), Vector{T}(undef, n_init), Vector{T}(undef, n_init-1))
+    resizedata!(U, n_init, n_init)
+    resizedata!(V, n_init, n_init)
     initiate_upper_mul_tri_triview!(UX, U, X) # fill-in 1st row
-    return TridiagonalConjugationData(U, X, V, UX, Y, 1)
+    initiate_tri_mul_invupper_triview!(Y, UX, V)
+    return TridiagonalConjugationData(U, X, V, UX, Y, 0)
 end
 
-TridiagonalConjugationData(U, X) = TridiagonalConjugationData(U, X, U)
+
+function TridiagonalConjugationData(U, X)
+    C = cache(U)
+    TridiagonalConjugationData(C, X, C)
+end
 
 copy(data::TridiagonalConjugationData) = TridiagonalConjugationData(copy(data.U), copy(data.X), copy(data.V), copy(data.UX), copy(data.Y), data.datasize)
 
@@ -209,8 +237,13 @@ function resizedata!(data::TridiagonalConjugationData, n)
         resize!(data.Y.du, 2n)
     end
 
-    main_upper_mul_tri_triview!(data.UX, data.U, data.X, data.datasize+1:n)
 
-    data.datasize = n
+    if n > data.datasize
+        main_upper_mul_tri_triview!(data.UX, data.U, data.X, data.datasize+2:n+1)
+        main_tri_mul_invupper_triview!(data.Y, data.UX, data.U, data.datasize+2:n+1) # need one extra as it updates first row
+        data.datasize = n
+    end
+
+    data
 end
 
diff --git a/src/infqr.jl b/src/infqr.jl
@@ -149,11 +149,14 @@ function adaptiveqr(A)
     QR(AdaptiveQRFactors(data), AdaptiveQRTau(data))
 end
 
-_qr(::AbstractBandedLayout, ::NTuple{2,OneToInf{Int}}, A) = adaptiveqr(A)
-_qr(::AbstractAlmostBandedLayout, ::NTuple{2,OneToInf{Int}}, A) = adaptiveqr(A)
-__qr(_, ::NTuple{2,InfiniteCardinal{0}}, A) = adaptiveqr(A)
-_qr(::AbstractBlockBandedLayout, ::NTuple{2,InfiniteCardinal{0}}, A) = adaptiveqr(A)
-_factorize(::AbstractBandedLayout, ::NTuple{2,OneToInf{Int}}, A) = qr(A)
+qr_layout(::AbstractBandedLayout, ::NTuple{2,OneToInf{Int}}, A) = adaptiveqr(A)
+qr_layout(::AbstractAlmostBandedLayout, ::NTuple{2,OneToInf{Int}}, A) = adaptiveqr(A)
+__qr_layout(_, ::NTuple{2,InfiniteCardinal{0}}, A) = adaptiveqr(A)
+qr_layout(::AbstractBlockBandedLayout, ::NTuple{2,InfiniteCardinal{0}}, A) = adaptiveqr(A)
+factorize_layout(::AbstractBandedLayout, ::NTuple{2,OneToInf{Int}}, A) = qr(A)
+
+
+cache_layout(::TriangularLayout{UPLO, UNIT, <:AdaptiveLayout}, A::AbstractMatrix) where {UPLO, UNIT} = A # already cached
 
 partialqr!(F::QR, n) = partialqr!(F.factors, n)
 partialqr!(F::AdaptiveQRFactors, n) = partialqr!(F.data, n)
diff --git a/test/test_bidiagonalconjugation.jl b/test/test_bidiagonalconjugation.jl
@@ -3,6 +3,7 @@ using InfiniteLinearAlgebra: BidiagonalConjugation, OneToInf
 using ArrayLayouts: supdiagonaldata, subdiagonaldata, diagonaldata
 using LinearAlgebra
 using LazyArrays: LazyLayout
+using BandedMatrices: _BandedMatrix
 
 @testset "BidiagonalConjugation" begin
     @test InfiniteLinearAlgebra._to_uplo('U') == 'U'
@@ -102,11 +103,25 @@ end
 
 
 @testset "TridiagonalConjugation" begin
-    @testset "T -> U" begin
-        R = BandedMatrices._BandedMatrix(Vcat(-Ones(1,∞)/2,
-                                        Zeros(1,∞),
-                                        Hcat(Ones(1,1),Ones(1,∞)/2)), ℵ₀, 0,2)
-        X_T = LazyBandedMatrices.Tridiagonal(Vcat(1.0, Fill(1/2,∞)), Zeros(∞), Fill(1/2,∞))
+    for (R,X_T) in (
+            # T -> U
+            (_BandedMatrix(Vcat(-Ones(1,∞)/2, Zeros(1,∞), Hcat(Ones(1,1),Ones(1,∞)/2)), ℵ₀, 0,2),
+            LazyBandedMatrices.Tridiagonal(Vcat(1.0, Fill(1/2,∞)), Zeros(∞), Fill(1/2,∞))),
+            # P -> C^(3/2)
+            (_BandedMatrix(Vcat((-1 ./ (1:2:∞))',
+                                         Zeros(1,∞),
+                                         (1 ./ (1:2:∞))'), ℵ₀, 0,2),
+             LazyBandedMatrices.Tridiagonal((1:∞) ./ (1:2:∞), Zeros(∞), (1:∞) ./ (3:2:∞))),
+             # P^(1,0) -> P^(2,0)
+            (_BandedMatrix(Vcat(Zeros(1,∞), # extra band since code assumes two bands
+             (-(0:∞) ./ (2:2:∞))',
+             ((2:∞) ./ (2:2:∞))'), ℵ₀, 0,2),
+             LazyBandedMatrices.Tridiagonal((2:∞) ./ (3:2:∞), -1 ./ ((1:2:∞) .* (3:2:∞)), (1:∞) ./ (3:2:∞))),
+            # P -> C^(5/2)
+            (_BandedMatrix(Vcat((-3 ./ (3:2:∞))', Zeros(1,∞), (3 ./ (3:2:∞))'), ℵ₀, 0,2) *
+            _BandedMatrix(Vcat((-1 ./ (1:2:∞))', Zeros(1,∞), (1 ./ (1:2:∞))'), ℵ₀, 0,2),
+            LazyBandedMatrices.Tridiagonal((1:∞) ./ (1:2:∞), Zeros(∞), (1:∞) ./ (3:2:∞)))
+            )
         n = 1000
         @time U = V = R[1:n,1:n];
         @time X = Tridiagonal(Vector(X_T.dl[1:n-1]), Vector(X_T.d[1:n]), Vector(X_T.du[1:n-1]));
@@ -116,50 +131,31 @@ end
         @time Y = InfiniteLinearAlgebra.tri_mul_invupper_triview(UX, U)
         @test Tridiagonal(U*X / U) ≈ Tridiagonal(UX / U) ≈ Y
 
-        InfiniteLinearAlgebra.TridiagonalConjugationData(U, X, U)
-    end
-    @testset "P -> Ultraspherical(3/2)" begin
-        R = BandedMatrices._BandedMatrix(Vcat((-1 ./ (1:2:∞))',
-                                         Zeros(1,∞),
-                                         (1 ./ (1:2:∞))'), ℵ₀, 0,2)
-        X_P = LazyBandedMatrices.Tridiagonal((1:∞) ./ (1:2:∞), Zeros(∞), (1:∞) ./ (3:2:∞))
-        n = 1000
-        @time U = V = R[1:n,1:n]
-        @time X = Tridiagonal(Vector(X_P.dl[1:n-1]), Vector(X_P.d[1:n]), Vector(X_P.du[1:n-1]))
-        @time UX = InfiniteLinearAlgebra.upper_mul_tri_triview(U, X)
-        @test Tridiagonal(U*X) ≈  UX
-        # U*X*inv(U) only depends on Tridiagonal(U*X)
-        @time Y = InfiniteLinearAlgebra.tri_mul_invupper_triview(UX, U)
-        @test Tridiagonal(U*X / U) ≈ Tridiagonal(UX / U) ≈ Y
-    end
+        data = InfiniteLinearAlgebra.TridiagonalConjugationData(R, X_T)
+        @test data.UX[1,:] ≈ UX[1,1:100]
+        InfiniteLinearAlgebra.resizedata!(data, 1)
+        @test data.UX[1:2,:] ≈ UX[1:2,1:100]
+        @test data.Y[1,:] ≈ Y[1,1:100]
+        InfiniteLinearAlgebra.resizedata!(data, 2)
+        @test data.UX[1:3,:] ≈ UX[1:3,1:100]
+        @test data.Y[1:2,:] ≈ Y[1:2,1:100]
+        InfiniteLinearAlgebra.resizedata!(data, 3)
+        @test data.UX[1:4,:] ≈ UX[1:4,1:100]
+        @test data.Y[1:3,:] ≈ Y[1:3,1:100]
+        InfiniteLinearAlgebra.resizedata!(data, 1000)
+        @test data.UX[1:999,1:999] ≈ UX[1:999,1:999]
+        @test data.Y[1:999,1:999] ≈ Y[1:999,1:999]
 
-    @testset "Jacobi(1,0) -> Jacobi(2,0)" begin
-        R = BandedMatrices._BandedMatrix(Vcat(Zeros(1,∞), # extra band since code assumes two bands
-                                         (-(0:∞) ./ (2:2:∞))',
-                                         ((2:∞) ./ (2:2:∞))'), ℵ₀, 0,2)
-        X_P = LazyBandedMatrices.Tridiagonal((2:∞) ./ (3:2:∞), -1 ./ ((1:2:∞) .* (3:2:∞)), (1:∞) ./ (3:2:∞))
-        n = 1000
-        @time U = V = R[1:n,1:n]
-        @time X = Tridiagonal(Vector(X_P.dl[1:n-1]), Vector(X_P.d[1:n]), Vector(X_P.du[1:n-1]))
-        @time UX = InfiniteLinearAlgebra.upper_mul_tri_triview(U, X)
-        @test Tridiagonal(U*X) ≈  UX
-        # U*X*inv(U) only depends on Tridiagonal(U*X)
-        @time Y = InfiniteLinearAlgebra.tri_mul_invupper_triview(UX, U)
-        @test Tridiagonal(U*X / U) ≈ Tridiagonal(UX / U) ≈ Y
+        data = InfiniteLinearAlgebra.TridiagonalConjugationData(R, X_T)
+        InfiniteLinearAlgebra.resizedata!(data, 1000)
+        @test data.UX[1:999,1:999] ≈ UX[1:999,1:999]
+        @test data.Y[1:999,1:999] ≈ Y[1:999,1:999]
     end
-
-    @testset "Legendre() -> Jacobi(5/2)" begin
-        R = BandedMatrices._BandedMatrix(Vcat((-3 ./ (3:2:∞))', Zeros(1,∞), (3 ./ (3:2:∞))'), ℵ₀, 0,2) *
-            BandedMatrices._BandedMatrix(Vcat((-1 ./ (1:2:∞))', Zeros(1,∞), (1 ./ (1:2:∞))'), ℵ₀, 0,2)
-        X_P = LazyBandedMatrices.Tridiagonal((1:∞) ./ (1:2:∞), Zeros(∞), (1:∞) ./ (3:2:∞))
-
-        n = 1000
-        @time U = V = R[1:n,1:n]
-        @time X = Tridiagonal(Vector(X_P.dl[1:n-1]), Vector(X_P.d[1:n]), Vector(X_P.du[1:n-1]))
-        @time UX = InfiniteLinearAlgebra.upper_mul_tri_triview(U, X)
-        @test Tridiagonal(U*X) ≈  UX
-        # U*X*inv(U) only depends on Tridiagonal(U*X)
-        @time Y = InfiniteLinearAlgebra.tri_mul_invupper_triview(UX, U)
-        @test Tridiagonal(U*X / U) ≈ Tridiagonal(UX / U) ≈ Y
+    
+    @testset "Cholesky" begin
+        M = Symmetric(_BandedMatrix(Vcat(Hcat(Fill(-1/(2sqrt(2)),1,3), Fill(-1/4,1,∞)), Zeros(1,∞), Hcat([0.5 0.25], Fill(0.5,1,∞))), ∞, 0, 2))
+        R = cholesky(M).U
+        X_T = LazyBandedMatrices.Tridiagonal(Vcat(1/sqrt(2), Fill(1/2,∞)), Zeros(∞), Vcat(1/sqrt(2), Fill(1/2,∞)))
+        data = InfiniteLinearAlgebra.TridiagonalConjugationData(R, X_T)
     end
 end
