Support Jacobi matrix computation via Cholesky and QR

README.md

@@ -69,8 +69,8 @@ julia> associatedlegendre(2)[0.1,1:10]
 ## p-Finite Element Method
 
 The language of quasi-arrays gives a natural framework for constructing p-finite element methods. The convention
-is that adjoint-products are understood as inner products over the axes with uniform weight. Thus to solve Poisson's equation
-using its weak formulation with Dirichlet conditions we can expand in a weighted Jacobi basis:
+is that adjoint-products are understood as inner products over the axes with uniform weight. To solve Poisson's equation
+using its weak formulation with Dirichlet conditions we can thus expand in a weighted Jacobi basis:
 ```julia
 julia> P¹¹ = Jacobi(1.0,1.0); # Quasi-matrix of Jacobi polynomials
 

src/ClassicalOrthogonalPolynomials.jl

@@ -100,8 +100,8 @@ __sum(::MappedOPLayout, A, dims) = __sum(MappedBasisLayout(), A, dims)
 ContinuumArrays.transform_ldiv_if_columns(L::Ldiv{MappedOPLayout,Lay}, ax::OneTo) where Lay = ContinuumArrays.transform_ldiv_if_columns(Ldiv{MappedBasisLayout,Lay}(L.A,L.B), ax)
 ContinuumArrays.transform_ldiv_if_columns(L::Ldiv{MappedOPLayout,ApplyLayout{typeof(hcat)}}, ax::OneTo) = ContinuumArrays.transform_ldiv_if_columns(Ldiv{MappedBasisLayout,UnknownLayout}(L.A,L.B), ax)
 
-_equals(::AbstractOPLayout, ::AbstractWeightedBasisLayout, _, _) = false # Weighedt-Legendre doesn't exist
-_equals(::AbstractWeightedBasisLayout, ::AbstractOPLayout, _, _) = false # Weighedt-Legendre doesn't exist
+_equals(::AbstractOPLayout, ::AbstractWeightedBasisLayout, _, _) = false # Weighted-Legendre doesn't exist
+_equals(::AbstractWeightedBasisLayout, ::AbstractOPLayout, _, _) = false # Weighted-Legendre doesn't exist
 
 _equals(::WeightedOPLayout, ::WeightedOPLayout, wP, wQ) = unweighted(wP) == unweighted(wQ)
 _equals(::WeightedOPLayout, ::WeightedBasisLayout, wP, wQ) = unweighted(wP) == unweighted(wQ) && weight(wP) == weight(wQ)
@@ -322,6 +322,7 @@ include("classical/chebyshev.jl")
 include("classical/ultraspherical.jl")
 include("classical/laguerre.jl")
 include("classical/fourier.jl")
+include("choleskyQR.jl")
 include("roots.jl")
 
 end # module

src/choleskyQR.jl

@@ -0,0 +1,167 @@
+"""
+cholesky_jacobimatrix(w, P)
+
+returns the Jacobi matrix `X` associated to a quasi-matrix of polynomials
+orthogonal with respect to `w(x) w_p(x)` where `w_p(x)` is the weight of the polynomials in `P`.
+
+The resulting polynomials are orthonormal on the same domain as `P`. The supplied `P` must be normalized. Accepted inputs are `w` as a function or `W` as an infinite matrix representing multiplication with the function `w` on the basis `P`.
+
+An optional bool can be supplied, i.e. `cholesky_jacobimatrix(sqrtw, P, false)` to disable checks of symmetry for the weight multiplication matrix and orthonormality for the basis (use with caution).
+"""
+function cholesky_jacobimatrix(w::Function, P::OrthogonalPolynomial, checks::Bool = true)
+    checks && !(P isa Normalized) && error("Polynomials must be orthonormal.")
+    W = Symmetric(P \ (w.(axes(P,1)) .* P)) # Compute weight multiplication via Clenshaw
+    return cholesky_jacobimatrix(W, P, false)     # At this point checks already passed or were entered as false, no need to recheck
+end
+function cholesky_jacobimatrix(W::AbstractMatrix, P::OrthogonalPolynomial, checks::Bool = true)
+    checks && !(P isa Normalized) && error("Polynomials must be orthonormal.")
+    checks && !(W isa Symmetric) && error("Weight modification matrix must be symmetric.")
+    return SymTridiagonal(CholeskyJacobiBands{:dv}(W,P),CholeskyJacobiBands{:ev}(W,P))
+end
+
+# The generated Jacobi operators are symmetric tridiagonal, so we store their data in cached bands
+mutable struct CholeskyJacobiBands{dv,T} <: AbstractCachedVector{T}
+    data::Vector{T}       # store band entries, :dv for diagonal, :ev for off-diagonal
+    U::UpperTriangular{T} # store upper triangular conversion matrix (needed to extend available entries)
+    UX::ApplyArray{T}     # store U*X, where X is the Jacobi matrix of the original P (needed to extend available entries)
+    datasize::Int         # size of so-far computed block 
+end
+
+# Computes the initial data for the Jacobi operator bands
+function CholeskyJacobiBands{:dv}(W, P::OrthogonalPolynomial{T}) where T
+    U = cholesky(W).U
+    X = jacobimatrix(P)
+    UX = ApplyArray(*,U,X)
+    dv = zeros(T,2) # compute a length 2 vector on first go
+    dv[1] = dot(view(UX,1,1), U[1,1] \ [one(T)])
+    dv[2] = dot(view(UX,2,1:2), U[1:2,1:2] \ [zero(T); one(T)])
+    return CholeskyJacobiBands{:dv,T}(dv, U, UX, 2)
+end
+function CholeskyJacobiBands{:ev}(W, P::OrthogonalPolynomial{T}) where T
+    U = cholesky(W).U
+    X = jacobimatrix(P)
+    UX = ApplyArray(*,U,X)
+    ev = zeros(T,2) # compute a length 2 vector on first go
+    ev[1] = dot(view(UX,1,1:2), U[1:2,1:2] \ [zero(T); one(T)])
+    ev[2] = dot(view(UX,2,1:3), U[1:3,1:3] \ [zeros(T,2); one(T)])
+    return CholeskyJacobiBands{:ev,T}(ev, U, UX, 2)
+end
+
+size(::CholeskyJacobiBands) = (ℵ₀,) # Stored as an infinite cached vector
+
+# Resize and filling functions for cached implementation
+function resizedata!(K::CholeskyJacobiBands, nm::Integer)
+    νμ = K.datasize
+    if nm > νμ
+        resize!(K.data,nm)
+        cache_filldata!(K, νμ:nm)
+        K.datasize = nm
+    end
+    K
+end
+function cache_filldata!(J::CholeskyJacobiBands{:dv,T}, inds::UnitRange{Int}) where T
+    # pre-fill U and UX to prevent expensive step-by-step filling in of cached U and UX in the loop
+    getindex(J.U,inds[end]+1,inds[end]+1)
+    getindex(J.UX,inds[end]+1,inds[end]+1)
+
+    ek = [zero(T); one(T)]
+    @inbounds for k in inds
+        J.data[k] = dot(view(J.UX,k,k-1:k), J.U[k-1:k,k-1:k] \ ek)
+    end
+end
+function cache_filldata!(J::CholeskyJacobiBands{:ev, T}, inds::UnitRange{Int}) where T
+    # pre-fill U and UX to prevent expensive step-by-step filling in of cached U and UX in the loop
+    getindex(J.U,inds[end]+1,inds[end]+1)
+    getindex(J.UX,inds[end]+1,inds[end]+1)
+
+    ek = [zeros(T,2); one(T)]
+    @inbounds for k in inds
+        J.data[k] = dot(view(J.UX,k,k-1:k+1), J.U[k-1:k+1,k-1:k+1] \ ek)
+    end
+end
+
+
+"""
+qr_jacobimatrix(sqrtw, P)
+
+returns the Jacobi matrix `X` associated to a quasi-matrix of polynomials
+orthogonal with respect to `w(x) w_p(x)` where `w_p(x)` is the weight of the polynomials in `P`.
+
+The resulting polynomials are orthonormal on the same domain as `P`. The supplied `P` must be normalized. Accepted inputs for `sqrtw` are the square root of the weight modification as a function or `sqrtW` as an infinite matrix representing multiplication with the function `sqrt(w)` on the basis `P`.
+
+An optional bool can be supplied, i.e. `qr_jacobimatrix(sqrtw, P, false)` to disable checks of symmetry for the weight multiplication matrix and orthonormality for the basis (use with caution).
+"""
+function qr_jacobimatrix(sqrtw::Function, P::OrthogonalPolynomial, checks::Bool = true)
+    checks && !(P isa Normalized) && error("Polynomials must be orthonormal.")
+    sqrtW = (P \ (sqrtw.(axes(P,1)) .* P))  # Compute weight multiplication via Clenshaw
+    return qr_jacobimatrix(sqrtW, P, false) # At this point checks already passed or were entered as false, no need to recheck
+end
+function qr_jacobimatrix(sqrtW::AbstractMatrix, P::OrthogonalPolynomial, checks::Bool = true)
+    checks && !(P isa Normalized) && error("Polynomials must be orthonormal.")
+    checks && !(sqrtW isa Symmetric) && error("Weight modification matrix must be symmetric.")
+    K = SymTridiagonal(QRJacobiBands{:dv}(sqrtW,P),QRJacobiBands{:ev}(sqrtW,P))
+    return K
+end
+
+# The generated Jacobi operators are symmetric tridiagonal, so we store their data in cached bands
+mutable struct QRJacobiBands{dv,T} <: AbstractCachedVector{T}
+    data::Vector{T}       # store band entries, :dv for diagonal, :ev for off-diagonal
+    U::ApplyArray{T}      # store upper triangular conversion matrix (needed to extend available entries)
+    UX::ApplyArray{T}     # store U*X, where X is the Jacobi matrix of the original P (needed to extend available entries)
+    datasize::Int         # size of so-far computed block 
+end
+
+# Computes the initial data for the Jacobi operator bands
+function QRJacobiBands{:dv}(sqrtW, P::OrthogonalPolynomial{T}) where T
+    U = qr(sqrtW).R
+    U = ApplyArray(*,Diagonal(sign.(view(U,band(0)))),U)
+    X = jacobimatrix(P)
+    UX = ApplyArray(*,U,X)
+    dv = zeros(T,2) # compute a length 2 vector on first go
+    dv[1] = dot(view(UX,1,1), U[1,1] \ [one(T)])
+    dv[2] = dot(view(UX,2,1:2), U[1:2,1:2] \ [zero(T); one(T)])
+    return QRJacobiBands{:dv,T}(dv, U, UX, 2)
+end
+function QRJacobiBands{:ev}(sqrtW, P::OrthogonalPolynomial{T}) where T
+    U = qr(sqrtW).R
+    U = ApplyArray(*,Diagonal(sign.(view(U,band(0)))),U)
+    X = jacobimatrix(P)
+    UX = ApplyArray(*,U,X)
+    ev = zeros(T,2) # compute a length 2 vector on first go
+    ev[1] = dot(view(UX,1,1:2), U[1:2,1:2] \ [zero(T); one(T)])
+    ev[2] = dot(view(UX,2,1:3), U[1:3,1:3] \ [zeros(T,2); one(T)])
+    return QRJacobiBands{:ev,T}(ev, U, UX, 2)
+end
+
+size(::QRJacobiBands) = (ℵ₀,) # Stored as an infinite cached vector
+
+# Resize and filling functions for cached implementation
+function resizedata!(K::QRJacobiBands, nm::Integer)
+    νμ = K.datasize
+    if nm > νμ
+        resize!(K.data,nm)
+        cache_filldata!(K, νμ:nm)
+        K.datasize = nm
+    end
+    K
+end
+function cache_filldata!(J::QRJacobiBands{:dv,T}, inds::UnitRange{Int}) where T
+    # pre-fill U and UX to prevent expensive step-by-step filling in of cached U and UX in the loop
+    getindex(J.U,inds[end]+1,inds[end]+1)
+    getindex(J.UX,inds[end]+1,inds[end]+1)
+
+    ek = [zero(T); one(T)]
+    @inbounds for k in inds
+        J.data[k] = dot(view(J.UX,k,k-1:k), J.U[k-1:k,k-1:k] \ ek)
+    end
+end
+function cache_filldata!(J::QRJacobiBands{:ev, T}, inds::UnitRange{Int}) where T
+    # pre-fill U and UX to prevent expensive step-by-step filling in of cached U and UX in the loop
+    getindex(J.U,inds[end]+1,inds[end]+1)
+    getindex(J.UX,inds[end]+1,inds[end]+1)
+
+    ek = [zeros(T,2); one(T)]
+    @inbounds for k in inds
+        J.data[k] = dot(view(J.UX,k,k-1:k+1), J.U[k-1:k+1,k-1:k+1] \ ek)
+    end
+end

test/runtests.jl

@@ -41,6 +41,7 @@ include("test_ratios.jl")
 include("test_normalized.jl")
 include("test_lanczos.jl")
 include("test_interlace.jl")
+include("test_choleskyQR.jl")
 include("test_roots.jl")
 
 @testset "Auto-diff" begin

test/test_choleskyQR.jl

@@ -0,0 +1,140 @@
+using Test, ClassicalOrthogonalPolynomials, BandedMatrices, LinearAlgebra, LazyArrays, ContinuumArrays, LazyBandedMatrices, InfiniteLinearAlgebra
+import ClassicalOrthogonalPolynomials: CholeskyJacobiBands, cholesky_jacobimatrix, qr_jacobimatrix, QRJacobiBands
+import LazyArrays: AbstractCachedMatrix
+
+@testset "Comparison of Cholesky with Lanczos and Classical" begin
+    @testset "Using Clenshaw for polynomial weights" begin
+        @testset "w(x) = x^2*(1-x)" begin
+            P = Normalized(legendre(0..1))
+            x = axes(P,1)
+            J = jacobimatrix(P)
+            wf(x) = x^2*(1-x)
+            # compute Jacobi matrix via cholesky
+            Jchol = cholesky_jacobimatrix(wf, P)
+            # compute Jacobi matrix via classical recurrence
+            Q = Normalized(Jacobi(1,2)[affine(0..1,Inclusion(-1..1)),:])
+            Jclass = jacobimatrix(Q)
+            # compute Jacobi matrix via Lanczos
+            Jlanc = jacobimatrix(LanczosPolynomial(@.(wf.(x)),Normalized(legendre(0..1))))
+            # Comparison with Lanczos
+            @test Jchol[1:500,1:500] ≈ Jlanc[1:500,1:500]
+            # Comparison with Classical
+            @test Jchol[1:500,1:500] ≈ Jclass[1:500,1:500]
+        end
+
+        @testset "w(x) = (1-x^2)" begin
+            P = Normalized(Legendre()[affine(0..1,Inclusion(-1..1)),:])
+            x = axes(P,1)
+            J = jacobimatrix(P)
+            wf(x) = (1-x^2)
+            # compute Jacobi matrix via cholesky
+            Jchol = cholesky_jacobimatrix(wf, P)
+            # compute Jacobi matrix via Lanczos
+            Jlanc = jacobimatrix(LanczosPolynomial(@.(wf.(x)),Normalized(legendre(0..1))))
+            # Comparison with Lanczos
+            @test Jchol[1:500,1:500] ≈ Jlanc[1:500,1:500]
+        end
+
+        @testset "w(x) = (1-x^4)" begin
+            P = Normalized(legendre(0..1))
+            x = axes(P,1)
+            J = jacobimatrix(P)
+            wf(x) = (1-x^4)
+            # compute Jacobi matrix via cholesky
+            Jchol = cholesky_jacobimatrix(wf, P)
+            # compute Jacobi matrix via Lanczos
+            Jlanc = jacobimatrix(LanczosPolynomial(@.(wf.(x)),Normalized(legendre(0..1))))
+            # Comparison with Lanczos
+            @test Jchol[1:500,1:500] ≈ Jlanc[1:500,1:500]
+        end
+
+        @testset "w(x) = (1.014-x^3)" begin
+            P = Normalized(Legendre()[affine(0..1,Inclusion(-1..1)),:])
+            x = axes(P,1)
+            J = jacobimatrix(P)
+            wf(x) = 1.014-x^4
+            # compute Jacobi matrix via cholesky
+            Jchol = cholesky_jacobimatrix(wf, P)
+            # compute Jacobi matrix via Lanczos
+            Jlanc = jacobimatrix(LanczosPolynomial(@.(wf.(x)),Normalized(legendre(0..1))))
+            # Comparison with Lanczos
+            @test Jchol[1:500,1:500] ≈ Jlanc[1:500,1:500]
+        end
+    end
+
+    @testset "Using Cholesky with exponential weights" begin
+        @testset "w(x) = exp(x)" begin
+            P = Normalized(Legendre()[affine(0..1,Inclusion(-1..1)),:])
+            x = axes(P,1)
+            J = jacobimatrix(P)
+            wf(x) = exp(x)
+            # compute Jacobi matrix via cholesky
+            Jchol = cholesky_jacobimatrix(wf, P)
+            # compute Jacobi matrix via Lanczos
+            Jlanc = jacobimatrix(LanczosPolynomial(@.(wf.(x)),Normalized(legendre(0..1))))
+            # Comparison with Lanczos
+            @test Jchol[1:500,1:500] ≈ Jlanc[1:500,1:500]
+        end
+
+        @testset "w(x) = (1-x)*exp(x)" begin
+            P = Normalized(Legendre()[affine(0..1,Inclusion(-1..1)),:])
+            x = axes(P,1)
+            J = jacobimatrix(P)
+            wf(x) = (1-x)*exp(x)
+            # compute Jacobi matrix via cholesky
+            Jchol = cholesky_jacobimatrix(wf, P)
+            # compute Jacobi matrix via Lanczos
+            Jlanc = jacobimatrix(LanczosPolynomial(@.(wf.(x)),Normalized(Legendre()[affine(0..1,Inclusion(-1..1)),:])))
+            # Comparison with Lanczos
+            @test Jchol[1:500,1:500] ≈ Jlanc[1:500,1:500]
+        end
+
+        @testset "w(x) = (1-x^2)*exp(x^2)" begin
+            P = Normalized(legendre(0..1))
+            x = axes(P,1)
+            J = jacobimatrix(P)
+            wf(x) = (1-x)^2*exp(x^2)
+            # compute Jacobi matrix via decomp
+            Jchol = cholesky_jacobimatrix(wf, P)
+            # compute Jacobi matrix via Lanczos
+            Jlanc = jacobimatrix(LanczosPolynomial(@.(wf.(x)),Normalized(legendre(0..1))))
+            # Comparison with Lanczos
+            @test Jchol[1:500,1:500] ≈ Jlanc[1:500,1:500]
+        end
+
+        @testset "w(x) = x*(1-x^2)*exp(-x^2)" begin
+            P = Normalized(Legendre()[affine(0..1,Inclusion(-1..1)),:])
+            x = axes(P,1)
+            J = jacobimatrix(P)
+            wf(x) = x*(1-x^2)*exp(-x^2)
+            # compute Jacobi matrix via cholesky
+            Jchol = cholesky_jacobimatrix(wf, P)
+            # compute Jacobi matrix via Lanczos
+            Jlanc = jacobimatrix(LanczosPolynomial(@.(wf.(x)),Normalized(legendre(0..1))))
+            # Comparison with Lanczos
+            @test Jchol[1:500,1:500] ≈ Jlanc[1:500,1:500]
+        end
+    end
+end
+
+@testset "Comparison of QR with Lanczos" begin
+    @testset "QR case, w(x) = (1-x)^2" begin
+        P = Normalized(legendre(0..1))
+        x = axes(P,1)
+        J = jacobimatrix(P)
+        wf(x) = (1-x)^2
+        sqrtwf(x) = (1-x)
+        # compute Jacobi matrix via decomp
+        Jchol = cholesky_jacobimatrix(wf, P)
+        Jqr = qr_jacobimatrix(sqrtwf, P)
+        # use alternative inputs
+        sqrtW = (P \ (sqrtwf.(x) .* P))
+        Jqralt = qr_jacobimatrix(sqrtW, P, false)
+        # compute Jacobi matrix via Lanczos
+        Jlanc = jacobimatrix(LanczosPolynomial(@.(wf.(x)),Normalized(legendre(0..1))))
+        # Comparison with Lanczos
+        @test Jchol[1:500,1:500] ≈ Jlanc[1:500,1:500]
+        @test Jqr[1:500,1:500] ≈ Jlanc[1:500,1:500]
+        @test Jqralt[1:500,1:500] ≈ Jlanc[1:500,1:500]
+    end
+end