Add Zernike(0,a) Jacobi matrices

Project.toml

@@ -26,8 +26,8 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 [compat]
 ArrayLayouts = "0.7, 0.8"
 BandedMatrices = "0.16, 0.17"
-BlockArrays = "0.16.7"
-BlockBandedMatrices = "0.11"
+BlockArrays = "0.16.14"
+BlockBandedMatrices = "0.11.5"
 ClassicalOrthogonalPolynomials = "0.5.1, 0.6"
 ContinuumArrays = "0.10"
 DomainSets = "0.5"

src/MultivariateOrthogonalPolynomials.jl

@@ -30,7 +30,7 @@ export MultivariateOrthogonalPolynomial, BivariateOrthogonalPolynomial,
        DunklXuDisk, DunklXuDiskWeight, WeightedDunklXuDisk,
        Zernike, ZernikeWeight, zerniker, zernikez,
        PartialDerivative, Laplacian, AbsLaplacianPower, AngularMomentum,
-       RadialCoordinate, Weighted, Block
+       RadialCoordinate, Weighted, Block, jacobimatrix
 
 include("ModalInterlace.jl")
 include("disk.jl")

src/disk.jl

@@ -160,8 +160,141 @@ getindex(Z::Zernike, xy::StaticVector{2}, B::BlockIndex{1}) = Z[RadialCoordinate
 getindex(Z::Zernike, xy::StaticVector{2}, B::Block{1}) = [Z[xy, B[j]] for j=1:Int(B)]
 getindex(Z::Zernike, xy::StaticVector{2}, JR::BlockOneTo) = mortar([Z[xy,Block(J)] for J = 1:Int(JR[end])])
 
+###
+# Jacobi matrices
+###
+
+# Due to excessively long lazy typing, we wrap the bands of the Zernike Jacobi matrix in its own type
+# The bands for X and Y are distinct
+struct ZernikeJacobimatrixBandsX{T} <: AbstractBlockMatrix{T}
+    Z::Zernike{T}
+    data::AbstractBlockMatrix{T}
+    ZernikeJacobimatrixBandsX{T}(Z::Zernike{T}) where T = new{T}(Z, zernikejacobibandsX(Z))
+end
+struct ZernikeJacobimatrixBandsY{T} <: AbstractBlockMatrix{T}
+    Z::Zernike{T}
+    data::AbstractBlockMatrix{T}
+    ZernikeJacobimatrixBandsY{T}(Z::Zernike{T}) where T = new{T}(Z, zernikejacobibandsY(Z))
+end
+
+size(b::ZernikeJacobimatrixBandsX) = size(b.data)
+axes(b::ZernikeJacobimatrixBandsX) = axes(b.data)
+size(b::ZernikeJacobimatrixBandsY) = size(b.data)
+axes(b::ZernikeJacobimatrixBandsY) = axes(b.data)
+
+function zernikejacobibandsX(Z::Zernike)
+    α = Z.b     # extract second basis parameter
+
+    k = mortar(Base.OneTo.(oneto(∞)))     # k counts the the angular mode (+1)
+    n = mortar(Fill.(oneto(∞),oneto(∞)))  # n counts the block number which corresponds to the order (+1)
+
+    # repeatedly used for sorting
+    keven = iseven.(k)
+    kodd = isodd.(k)
+    neven = iseven.(n)
+    nodd = isodd.(n)
+
+    # h1-h5 are helpers for our different sorting scheme
+
+    ## Compute super diagonal of super diagonal blocks.
+    dufirst = neven .* (k .== 2) .* n .* (n .+ 2*α) ./ 2
+
+    ## Compute even-block entries in super diagonal of super diagonal blocks
+    h1 = n .- (n .+ 2 .- k .- keven .* (n .> 2)) .÷ 2
+    dueven = ((nodd .* k) .>= 2) .* h1 .* (h1 .+ α)
+
+    ## Compute odd-block entries in super diagonal of super diagonal blocks
+    h2 = (n .- 2 .+ k .+ kodd .* (n .> 2)) .÷ 2
+    duodd = (((neven .* k) .>= 2)  .- (neven .* (k .== 2))) .* h2 .* (h2 .+ α)
+
+    ## Compute even-block sub diagonal elements of super diagonal blocks
+    h3 = n .- (k .+ 1 .+ kodd .+ n) .÷ 2
+    dleven = ((k .<= (n .- 2)) .- ((k .< 2) .* nodd)) .* neven .* h3 .* (h3 .+ α)
+
+    ## Compute odd-block sub diagonal elements of super diagonal blocks
+    h4 = n .- (k .+ 1 .+ keven .+ n) .÷ 2
+    dlodd = (k .> 1) .* (n .> 3) .* nodd .* h4 .* (h4 .+ α)
 
+    ## Compute and add in special case odd-block sub diagonal elements of super diagonal blocks
+    h5 = n .- 2 .+ k .+ keven
+    dlspecial = nodd .* (k .== 1) .* h5 .* (h5 .+ 2*α) ./ 2
 
+    # finalize bands with explicit formula
+    quotient = 4 .* (n .+ (α-1)) .* (n .+ α)
+    du = sqrt.( (dufirst .+ dueven .+ duodd ) ./ quotient)
+    dl = sqrt.( (dleven .+ dlodd .+ dlspecial)  ./ quotient)
+
+    return BlockHcat(
+        BlockBroadcastArray(hcat, du, Zeros((axes(n,1),)), dl),
+        Zeros((axes(n,1),Base.OneTo(3))),
+        Zeros((axes(n,1),Base.OneTo(3)))                  
+        )
+end
+
+function zernikejacobibandsY(Z::Zernike)
+    α = Z.b     # extract second basis parameter
+
+    k = mortar(Base.OneTo.(oneto(∞)))     # k counts the the angular mode (+1)
+    n = mortar(Fill.(oneto(∞),oneto(∞)))  # n counts the block number which corresponds to the order (+1)
+
+    # repeatedly used for sorting
+    keven = iseven.(k)
+    kodd = isodd.(k)
+    neven = iseven.(n)
+    nodd = isodd.(n)
+
+    # h1-h4 are helpers for our different sorting scheme
+
+    # first entries for all blocks
+    h1 = (n .- nodd)
+    l1 = (k .== 1) .* (h1 .* (h1 .+ 2*α) ./ 2)
+
+    # Even blocks
+    h0 = (kodd .* ((k .÷ 2) .+ 1) .- (keven .* ((k .÷ 2) .- 1)))
+    h2 =  (k .>= 2) .* ((n .÷ 2 .- 1) .+ h0)
+    l2 = neven .* (h2 .* (h2 .+ α))
+
+    # Odd blocks
+    h3 = (n .> k .>= 2) .* (((n .+ 1) .÷ 2) .- h0)
+    l3 = nodd .* (h3 .* (h3 .+ α))
+    # Combine for diagonal of super diagonal block
+    d = sqrt.((l1 .+ l2 .+ l3) ./ (4 .* (n .+ (α-1)) .* (n .+ α)))
+
+    # The off-diagonals of the super diagonal block are negative, shifted versions of the diagonal with some entries skipped
+    dl = (-1) .* (nodd .* kodd .+ neven .* keven) .* Vcat(0 , d)
+    du = (-1) .* (nodd .* keven .+ neven .* kodd) .* d[2:end]
+
+    # zero blocks for banded blockarray structure
+    z = Zeros((axes(n,1),))
+    z5 = Zeros((axes(n,1),Base.OneTo(5)))
+
+    # generate and return bands
+    return dat = BlockHcat(BlockBroadcastArray(hcat, dl, z, d, z, du), z5, z5)  
+end
+
+function getindex(b::ZernikeJacobimatrixBandsX{T},i,j) where T
+    return BlockArrays.getindex(b.data,i,j)
+end
+function getindex(b::ZernikeJacobimatrixBandsY{T},i,j) where T
+    return BlockArrays.getindex(b.data,i,j)
+end
+
+function jacobimatrix(::Val{1}, Z::Zernike{T}) where T
+    if iszero(Z.a)
+        dat = ZernikeJacobimatrixBandsX{T}(Z)
+        return Symmetric(BlockBandedMatrices._BandedBlockBandedMatrix(dat', axes(dat,1), (1,1), (1,1)))
+    else
+        error("Implement for non-zero first basis parameter.")
+    end
+end
+function jacobimatrix(::Val{2}, Z::Zernike{T}) where T
+    if iszero(Z.a)
+        dat = ZernikeJacobimatrixBandsY{T}(Z)
+        return Symmetric(BlockBandedMatrices._BandedBlockBandedMatrix(dat', axes(dat,1), (1,1), (2,2)))
+    else
+        error("Implement for non-zero first basis parameter.")
+    end
+end
 
 ###
 # Transforms
@@ -314,4 +447,4 @@ function \(A::Zernike{T}, wB::Weighted{V,Zernike{V}}) where {T,V}
     c = Int(B.a - A.a + B.b - A.b)
     @assert iszero(B.a)
     ModalInterlace{TV}((Normalized.(Jacobi{TV}.(A.b, A.a:∞)) .\ HalfWeighted{:a}.(Normalized.(Jacobi{TV}.(B.b, B.a:∞)))) .* convert(TV, 2)^(-c/2), (ℵ₀,ℵ₀), (2Int(B.b), 2Int(A.a+A.b)))
-end
+end

test/test_disk.jl

@@ -1,4 +1,4 @@
-using MultivariateOrthogonalPolynomials, ClassicalOrthogonalPolynomials, StaticArrays, BlockArrays, BandedMatrices, FastTransforms, LinearAlgebra, RecipesBase, Test, SpecialFunctions, LazyArrays
+using MultivariateOrthogonalPolynomials, ClassicalOrthogonalPolynomials, StaticArrays, BlockArrays, BandedMatrices, FastTransforms, LinearAlgebra, RecipesBase, Test, SpecialFunctions, LazyArrays, InfiniteArrays
 import MultivariateOrthogonalPolynomials: DiskTrav, grid, ZernikeTransform, ZernikeITransform, *, ModalInterlace
 import ClassicalOrthogonalPolynomials: HalfWeighted
 import ForwardDiff: hessian
@@ -68,6 +68,45 @@ import ForwardDiff: hessian
         end
     end
 
+    @testset "Jacobi matrices" begin
+        α = 10 * rand()
+        Z = Zernike(α)
+        xy = axes(Z, 1)
+        x, y = first.(xy), last.(xy)
+        n = 150
+            # X tests
+        JX = zeros(n,n)
+        for j = 1:n
+            JX[1:n,j] = (Z \ (x .* Z[:,j]))[1:n]
+        end 
+        X = jacobimatrix(Val(1),Z)
+        # The Zernike Jacobi matrices are symmetric for this normalization
+        @test issymmetric(X)
+        # Consistency with expansion
+        @test X[1:150,1:150] ≈ JX
+        # Multiplication by x
+        f = Z \ (sin.(x.*y) .+ x.^2 .- y)
+        xf = Z \ (x.*sin.(x.*y) .+ x.^3 .- x.*y)
+        @test X[Block.(1:20),Block.(1:20)]*f[Block.(1:20)] ≈ xf[Block.(1:20)]
+            # Y tests
+        JY = zeros(n,n)
+        for j = 1:n
+        JY[1:n,j] = (Z \ (y .* Z[:,j]))[1:n]
+        end 
+        Y = jacobimatrix(Val(2),Z)
+        # The Zernike Jacobi matrices are symmetric for this normalization
+        @test issymmetric(Y)
+        # Consistency with expansion
+        @test Y[1:150,1:150] ≈ JY
+        # Multiplication by x
+        f = Z \ (sin.(x.*y) .+ x.^2 .- y)
+        yf = Z \ (y.*sin.(x.*y) .+ y .* x.^2 .- y.^2)
+        @test Y[Block.(1:20),Block.(1:20)]*f[Block.(1:20)] ≈ yf[Block.(1:20)]
+        # data size tests
+        @test size((X.data).data) == (9, ℵ₀)
+        @test size((Y.data).data) == (15, ℵ₀)
+    end
+
     @testset "Transform" begin
         for (a,b) in ((0,0), (0.1, 0.2), (0,1))
             Zn = Zernike(a,b)[:,Block.(Base.OneTo(3))]
@@ -503,4 +542,4 @@ end
             end
         end
     end
-end
+end