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,170 @@
+"""
+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
+    bands = CholeskyJacobiBands(W, P)
+    return SymTridiagonal(bands[1,:],bands[2,:])
+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.")
+    bands = CholeskyJacobiBands(W, P)
+    return SymTridiagonal(bands[1,:],bands[2,:])
+end
+
+# The generated Jacobi operators are symmetric tridiagonal, so we store their data as two bands
+mutable struct CholeskyJacobiBands{T} <: AbstractCachedMatrix{T}
+    data::Matrix{T}         # store so-far computed block of entries 
+    U::UpperTriangular      # store upper triangular conversion matrix (needed to extend available entries)
+    X::SymTridiagonal{T}    # store the Jacobi matrix of the original basis (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(W::Symmetric{T}, P::OrthogonalPolynomial) where T
+    U = cholesky(W).U
+    X = jacobimatrix(P)
+    dat = zeros(T,2,10)
+    UX = U*X
+    dat[1,1] = (UX * (U \ [1; zeros(∞)]))[1]
+    dat[2,1] = (UX * (U \ [zeros(1); 1; zeros(∞)]))[1]
+    @inbounds for k in 2:10
+        yk = view(UX,k,k-1:k+1)
+        dat[1,k] = dot(yk[1:2], U[k-1:k,k-1:k] \ [0; 1])
+        dat[2,k] = dot(yk, U[k-1:k+1,k-1:k+1] \ [zeros(2); 1])
+    end
+    return CholeskyJacobiBands{T}(dat, U, X, 10)
+end
+
+size(::CholeskyJacobiBands) = (2,ℵ₀) # Stored as two infinite bands
+
+# Resize and filling functions for cached implementation
+function resizedata!(K::CholeskyJacobiBands, nm::Integer)
+    νμ = K.datasize
+    if nm > νμ
+        olddata = copy(K.data)
+        K.data = similar(K.data, 2, nm)
+        K.data[axes(olddata)...] = olddata
+        inds = νμ:nm
+        cache_filldata!(K, inds)
+        K.datasize = nm
+    end
+    K
+end
+function cache_filldata!(J::CholeskyJacobiBands, inds::UnitRange{Int})
+    UX = J.U*J.X
+    @inbounds for k in inds
+        yk = view(UX,k,k-1:k+1)
+        J.data[1,k] = dot(yk[1:2], J.U[k-1:k,k-1:k] \ [0; 1])
+        J.data[2,k] = dot(yk, J.U[k-1:k+1,k-1:k+1] \ [zeros(2); 1])
+    end
+end
+
+function getindex(K::CholeskyJacobiBands, k::Integer, j::Integer)
+    resizedata!(K, max(k,j))
+    K.data[k, j]
+end
+function getindex(K::CholeskyJacobiBands, kr::Integer, jr::UnitRange{Int})
+    resizedata!(K, maximum(jr))
+    K.data[kr, jr]
+end
+function getindex(K::CholeskyJacobiBands, I::Vararg{Int,2})
+    resizedata!(K,maximum(I))
+    getindex(K.data,I[1],I[2])
+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
+    bands = QRJacobiBands(sqrtW,P)
+    K = SymTridiagonal(bands[1,:],bands[2,:])
+    return K
+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.")
+    bands = QRJacobiBands(sqrtW,P)
+    K = SymTridiagonal(bands[1,:],bands[2,:])
+    return K
+end
+
+# The generated Jacobi operators are symmetric tridiagonal, so we store their data as two bands
+mutable struct QRJacobiBands{T} <: AbstractCachedMatrix{T}
+    data::Matrix{T}         # store so-far computed block of entries 
+    U::UpperTriangular      # store upper triangular conversion matrix (needed to extend available entries)
+    X::SymTridiagonal{T}    # store the Jacobi matrix of the original basis (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(sqrtW, P::OrthogonalPolynomial{T}) where T
+    U = qr(sqrtW).R
+    U = ApplyArray(*,Diagonal(sign.(view(U,band(0)))),U)
+    X = jacobimatrix(P)
+    UX = U*X
+    dat = zeros(T,2,10)
+    dat[1,1] = (UX * (U \ [1; zeros(∞)]))[1]
+    dat[2,1] = (UX * (U \ [zeros(1); 1; zeros(∞)]))[1]
+    @inbounds for k in 2:10
+        yk = view(UX,k,k-1:k+1)
+        dat[1,k] = dot(yk[1:2], U[k-1:k,k-1:k] \ [0; 1])
+        dat[2,k] = dot(yk, U[k-1:k+1,k-1:k+1] \ [zeros(2); 1])
+    end
+    return QRJacobiBands{T}(dat, UpperTriangular(U), X, 10)
+end
+
+size(::QRJacobiBands) = (2,ℵ₀) # Stored as two infinite bands
+
+# Resize and filling functions for cached implementation
+function resizedata!(K::QRJacobiBands, nm::Integer)
+    νμ = K.datasize
+    if nm > νμ
+        olddata = copy(K.data)
+        K.data = similar(K.data, 2, nm)
+        K.data[axes(olddata)...] = olddata
+        inds = νμ:nm
+        cache_filldata!(K, inds)
+        K.datasize = nm
+    end
+    K
+end
+function cache_filldata!(J::QRJacobiBands, inds::UnitRange{Int})
+    UX = J.U*J.X
+    @inbounds for k in inds
+        yk = view(UX,k,k-1:k+1)
+        J.data[1,k] = dot(yk[1:2], J.U[k-1:k,k-1:k] \ [0; 1])
+        J.data[2,k] = dot(yk, J.U[k-1:k+1,k-1:k+1] \ [zeros(2); 1])
+    end
+end
+
+function getindex(K::QRJacobiBands, k::Integer, j::Integer)
+    resizedata!(K, max(k,j))
+    K.data[k, j]
+end
+function getindex(K::QRJacobiBands, kr::Integer, jr::UnitRange{Int})
+    resizedata!(K, maximum(jr))
+    K.data[kr, jr]
+end
+function getindex(K::QRJacobiBands, I::Vararg{Int,2})
+    resizedata!(K,maximum(I))
+    getindex(K.data,I[1],I[2])
+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,178 @@
+using Test, ClassicalOrthogonalPolynomials, BandedMatrices, LinearAlgebra, LazyArrays, ContinuumArrays, LazyBandedMatrices, InfiniteLinearAlgebra
+import ClassicalOrthogonalPolynomials: CholeskyJacobiBands, cholesky_jacobimatrix, qr_jacobimatrix, QRJacobiBands
+import LazyArrays: AbstractCachedMatrix
+
+@testset "Cholesky Jacobi Matrix - Basic properties" begin
+    # basis
+    P = Normalized(legendre(0..1))
+    x = axes(P,1)
+    J = jacobimatrix(P)
+    # example weight
+    w(x) = (1 - x^2)
+    W = Symmetric(P \ (w.(x) .* P))
+    # banded cholesky for symmetric-tagged W
+    @test cholesky(W).U isa UpperTriangular
+    # compute Jacobi matrix via cholesky
+    Jchol = cholesky_jacobimatrix(w,P)
+    Jcholalt = cholesky_jacobimatrix(W,P)
+    @test Jchol[1:200,1:200] ≈ Jcholalt[1:200,1:200]
+    # CholeskyJacobiBands object
+    Cbands = CholeskyJacobiBands(W,P)
+    @test Cbands isa CholeskyJacobiBands
+    @test Cbands isa AbstractCachedMatrix
+    @test getindex(Cbands,1,100) == getindex(Cbands,1,1:100)[100]
+    @test getindex(Cbands,2,2) == getindex(Cbands,1:2,1:2)[2,2]
+end
+
+@testset "Comparison 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 "QR version" begin
+    @testset "basic properties" begin
+        # basis
+        P = Normalized(legendre(0..1))
+        x = axes(P,1)
+        J = jacobimatrix(P)
+        # example weight
+        sqrtw(x) = (1 - x)
+        sqrtW = Symmetric(P \ (sqrtw.(x) .* P))
+        # bands test
+        QRbands = QRJacobiBands(sqrtW,P)
+        @test QRbands isa QRJacobiBands
+        @test QRbands isa AbstractCachedMatrix
+        @test getindex(QRbands,1,100) == getindex(QRbands,1,1:100)[100]
+        @test getindex(QRbands,2,2) == getindex(QRbands,1:2,1:2)[2,2]
+    end
+
+    @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