Add Zernike transform (#64)

* Add Zernike transform

* Update test_disk.jl

Author: dlfivefifty

src/MultivariateOrthogonalPolynomials.jl

@@ -16,8 +16,8 @@ import LinearAlgebra: factorize
 import LazyArrays: arguments, paddeddata
 
 import ClassicalOrthogonalPolynomials: jacobimatrix
-import HarmonicOrthogonalPolynomials: BivariateOrthogonalPolynomial, MultivariateOrthogonalPolynomial, 
-                                          PartialDerivative, BlockOneTo, interlace
+import HarmonicOrthogonalPolynomials: BivariateOrthogonalPolynomial, MultivariateOrthogonalPolynomial, Plan,
+                                          PartialDerivative, BlockOneTo, BlockRange1, interlace
 
 export UnitTriangle, UnitDisk, JacobiTriangle, TriangleWeight, WeightedTriangle, PartialDerivative, Laplacian, 
       MultivariateOrthogonalPolynomial, BivariateOrthogonalPolynomial, Zernike, RadialCoordinate

src/disk.jl

@@ -60,11 +60,43 @@ function getindex(Z::Zernike{T}, rθ::RadialCoordinate, B::BlockIndex{1}) where
     if iszero(m)
         sqrt(convert(T,2)/π) * Normalized(Legendre{T}())[2r^2-1, ℓ ÷ 2 + 1]
     else
-        convert(T,2^((m+2)/2))/π * r^m * Normalized(Jacobi{T}(0, m))[2r^2-1,(ℓ-m) ÷ 2 + 1] * (isodd(k+ℓ) ? cos(m*θ) : sin(m*θ))
+        sqrt(convert(T,2)^(m+2)/π) * r^m * Normalized(Jacobi{T}(0, m))[2r^2-1,(ℓ-m) ÷ 2 + 1] * (isodd(k+ℓ) ? cos(m*θ) : sin(m*θ))
     end
 end
 
 
 getindex(Z::Zernike, xy::StaticVector{2}, B::BlockIndex{1}) = Z[RadialCoordinate(xy), B]
 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])])
+getindex(Z::Zernike, xy::StaticVector{2}, JR::BlockOneTo) = mortar([Z[xy,Block(J)] for J = 1:Int(JR[end])])
+
+
+###
+# Transforms
+###
+
+const FiniteZernike{T} = SubQuasiArray{T,2,Zernike{T},<:Tuple{<:Inclusion,<:BlockSlice{BlockRange1{OneTo{Int}}}}}
+
+function grid(S::FiniteZernike{T}) where T
+    N = blocksize(S,2) ÷ 2 + 1 # polynomial degree
+    M = 4N-3
+
+    r = sinpi.((N .-(0:N-1) .- one(T)/2) ./ (2N))
+
+    # The angular grid:
+    θ = (0:M-1)*convert(T,2)/M
+    RadialCoordinate.(r, π*θ')
+end
+
+struct ZernikeTransform{T} <: Plan{T}
+    N::Int
+    disk2cxf::FastTransforms.FTPlan{T,2,FastTransforms.DISK}
+    analysis::FastTransforms.FTPlan{T,2,FastTransforms.DISKANALYSIS}
+end
+
+function ZernikeTransform{T}(N::Int) where T<:Real
+    Ñ = N ÷ 2 + 1
+    ZernikeTransform{T}(N, plan_disk2cxf(T, Ñ, 0, 0), plan_disk_analysis(T, Ñ, 4Ñ-3))
+end
+*(P::ZernikeTransform{T}, f::Matrix{T}) where T = DiskTrav(P.disk2cxf \ (P.analysis * f))[Block.(1:P.N)]
+
+factorize(S::FiniteZernike{T}) where T = TransformFactorization(grid(S), ZernikeTransform{T}(blocksize(S,2)))

test/test_disk.jl

@@ -1,5 +1,5 @@
-using MultivariateOrthogonalPolynomials, StaticArrays, BlockArrays, Test
-import MultivariateOrthogonalPolynomials: DiskTrav
+using MultivariateOrthogonalPolynomials, StaticArrays, BlockArrays, FastTransforms, LinearAlgebra, Test
+import MultivariateOrthogonalPolynomials: DiskTrav, grid
 
 function chebydiskeval(c::AbstractMatrix{T}, r, θ) where T
     ret = zero(T)
@@ -21,8 +21,8 @@ end
         xy = SVector(rθ)
         @test Zernike()[rθ,1] ≈ Zernike()[xy,1] ≈ inv(sqrt(π))
         @test Zernike()[rθ,Block(1)] ≈ Zernike()[xy,Block(1)] ≈ [inv(sqrt(π))]
-        @test Zernike()[rθ,Block(2)] ≈ [2r/π*sin(θ), 2r/π*cos(θ)]
-        @test Zernike()[rθ,Block(3)] ≈ [sqrt(3/π)*(2r^2-1),sqrt(6)/π*r^2*sin(2θ),sqrt(6)/π*r^2*cos(2θ)]
+        @test Zernike()[rθ,Block(2)] ≈ [2r/sqrt(π)*sin(θ), 2r/sqrt(π)*cos(θ)]
+        @test Zernike()[rθ,Block(3)] ≈ [sqrt(3/π)*(2r^2-1),sqrt(6/π)*r^2*sin(2θ),sqrt(6/π)*r^2*cos(2θ)]
     end
 
     @testset "DiskTrav" begin
@@ -40,4 +40,19 @@ end
         @test_throws ArgumentError DiskTrav([1 2 3; 4 5 6])
         @test_throws ArgumentError DiskTrav([1 2 3 4; 5 6 7 8])
     end
-end
+
+    @testset "Transform" begin
+        Zn = Zernike()[:,Block.(Base.OneTo(3))]
+        for k = 1:6
+            @test factorize(Zn) \ Zernike()[:,k] ≈ [zeros(k-1); 1; zeros(6-k)]
+        end
+
+        Z = Zernike();
+        xy = axes(Z,1); x,y = first.(xy),last.(xy);
+        u = Z * (Z \ exp.(x .* cos.(y)))
+        @test u[SVector(0.1,0.2)] ≈ exp(0.1cos(0.2))
+    end
+end
+
+
+