Simplify Zernike Jacobi matrices + bugfixes (#123)

* Add tests for addition and multiplication

* Update disk.jl

* Update Project.toml

* Update test_disk.jl

* Update test_disk.jl

* add tests for copy

* Overload copy() so that Z \ (x .* Z) works

* fix minor inconsistency

* fix typo

* remove custom type for Zernike bands

* adjust tests

* up LazyBandedMatrices

Author: TSGut

Project.toml

@@ -1,6 +1,6 @@
 name = "MultivariateOrthogonalPolynomials"
 uuid = "4f6956fd-4f93-5457-9149-7bfc4b2ce06d"
-version = "0.2.5"
+version = "0.2.6"
 
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
@@ -24,7 +24,7 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]
-ArrayLayouts = "0.7, 0.8"
+ArrayLayouts = "0.8.6"
 BandedMatrices = "0.16, 0.17"
 BlockArrays = "0.16.14"
 BlockBandedMatrices = "0.11.5"
@@ -35,9 +35,9 @@ FastTransforms = "0.13"
 FillArrays = "0.12, 0.13"
 HarmonicOrthogonalPolynomials = "0.2.4"
 InfiniteArrays = "0.12"
-InfiniteLinearAlgebra = "0.6"
+InfiniteLinearAlgebra = "0.6.6"
 LazyArrays = "0.22.10"
-LazyBandedMatrices = "0.7.12"
+LazyBandedMatrices = "0.7.15"
 QuasiArrays = "0.9"
 SpecialFunctions = "1, 2"
 StaticArrays = "1"

src/disk.jl

@@ -160,30 +160,13 @@ 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
+function jacobimatrix(::Val{1}, Z::Zernike{T}) where T
+    if iszero(Z.a)
+        α = 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)
@@ -224,65 +207,48 @@ function zernikejacobibandsX(Z::Zernike)
     du = sqrt.( (dufirst .+ dueven .+ duodd ) ./ quotient)
     dl = sqrt.( (dleven .+ dlodd .+ dlspecial)  ./ quotient)
 
-    return BlockBroadcastArray(hcat, du, Zeros((axes(n,1),)), dl)
-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]
-
-    # generate and return bands
-    return dat = BlockBroadcastArray(hcat, dl, Zeros((axes(n,1),)), d, Zeros((axes(n,1),)), du)
-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)))
+    return Symmetric(BlockBandedMatrices._BandedBlockBandedMatrix(BlockBroadcastArray(hcat, du, Zeros((axes(n,1),)), dl)', axes(n,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)))
+        α = 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) .* view(d,2:∞)
+    
+        # generate and return bands
+        return Symmetric(BlockBandedMatrices._BandedBlockBandedMatrix(BlockBroadcastArray(hcat, dl, Zeros((axes(n,1),)), d, Zeros((axes(n,1),)), du)', axes(n,1), (-1,1), (2,2)))
     else
         error("Implement for non-zero first basis parameter.")
     end

test/test_disk.jl

@@ -68,18 +68,20 @@ import ForwardDiff: hessian
         end
     end
 
-    @testset "Jacobi matrices" begin
+       @testset "Jacobi matrices" begin
+        # Setup
         α = 10 * rand()
         Z = Zernike(α)
         xy = axes(Z, 1)
         x, y = first.(xy), last.(xy)
         n = 150
-            # X tests
+
+        # 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)
+        X = Z \ (x .* Z)
         # The Zernike Jacobi matrices are symmetric for this normalization
         @test issymmetric(X)
         # Consistency with expansion
@@ -88,23 +90,29 @@ import ForwardDiff: hessian
         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
+
+        # 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)
+        Y = Z \ (y .* 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
+        # Multiplication by y
         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) == (3, ℵ₀)
-        @test size((Y.data).data) == (5, ℵ₀)
+            
+        # Multiplication of Jacobi matrices
+        @test (X*X)[Block.(1:6),Block.(1:6)] ≈ (X[Block.(1:10),Block.(1:10)]*X[Block.(1:10),Block.(1:10)])[Block.(1:6),Block.(1:6)]
+        @test (X*Y)[Block.(1:6),Block.(1:6)] ≈ (X[Block.(1:10),Block.(1:10)]*Y[Block.(1:10),Block.(1:10)])[Block.(1:6),Block.(1:6)]
+
+        # Addition of Jacobi matrices
+        @test (X+Y)[Block.(1:6),Block.(1:6)] ≈ (X[Block.(1:6),Block.(1:6)]+Y[Block.(1:6),Block.(1:6)])
+        @test (Y+Y)[Block.(1:6),Block.(1:6)] ≈ (Y[Block.(1:6),Block.(1:6)]+Y[Block.(1:6),Block.(1:6)])
 
         # for now, reject non-zero first parameter options
         @test_throws ErrorException("Implement for non-zero first basis parameter.") jacobimatrix(Val(1),Zernike(1,1))