Generalized Zernike (#65)

* Start Zernike weight

* fix m

* test other parameters

* Update disk.jl

* Update disk.jl

* Work on Laplacian

* Start xdisk conversion

* Update Project.toml

* Zernike conversion works

* Add Zernike \

* work on Heat equation

* tests pass

* increase coverage

Author: dlfivefifty

Project.toml

@@ -1,6 +1,6 @@
 name = "MultivariateOrthogonalPolynomials"
 uuid = "4f6956fd-4f93-5457-9149-7bfc4b2ce06d"
-version = "0.0.7"
+version = "0.0.8"
 
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
@@ -12,6 +12,7 @@ ContinuumArrays = "7ae1f121-cc2c-504b-ac30-9b923412ae5c"
 DomainSets = "5b8099bc-c8ec-5219-889f-1d9e522a28bf"
 FastTransforms = "057dd010-8810-581a-b7be-e3fc3b93f78c"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 HarmonicOrthogonalPolynomials = "e416a80e-9640-42f3-8df8-80a93ca01ea5"
 InfiniteArrays = "4858937d-0d70-526a-a4dd-2d5cb5dd786c"
 InfiniteLinearAlgebra = "cde9dba0-b1de-11e9-2c62-0bab9446c55c"
@@ -28,12 +29,12 @@ ArrayLayouts = "0.6"
 BandedMatrices = "0.16"
 BlockArrays = "0.14.1, 0.15"
 BlockBandedMatrices = "0.10.1"
-ClassicalOrthogonalPolynomials = "0.3"
-ContinuumArrays = "0.5, 0.6"
+ClassicalOrthogonalPolynomials = "0.3.1"
+ContinuumArrays = "0.6"
 DomainSets = "0.4"
 FastTransforms = "0.11, 0.12"
 FillArrays = "0.11"
-HarmonicOrthogonalPolynomials = "0.0.3"
+HarmonicOrthogonalPolynomials = "0.0.4"
 InfiniteArrays = "0.10"
 InfiniteLinearAlgebra = "0.5"
 LazyArrays = "0.21"

examples/diskheat.jl

@@ -0,0 +1,13 @@
+using MultivariateOrthogonalPolynomials, DifferentialEquations
+
+N = 20
+WZ = Weighted(Zernike(1))[:,Block.(Base.OneTo(N))]
+Δ = Laplacian(axes(WZ,1))
+(Δ*WZ).args[3].data |> typeof
+using LazyArrays: arguments, ApplyLayout
+arguments(ApplyLayout{typeof(*)}(), WZ)[2].data
+@ent arguments(ApplyLayout{typeof(*)}(), WZ)
+
+function heat!(du, u, (R,Δ), t)
+
+end

src/MultivariateOrthogonalPolynomials.jl

@@ -11,16 +11,17 @@ import QuasiArrays: LazyQuasiMatrix, LazyQuasiArrayStyle
 import ContinuumArrays: @simplify, Weight, grid, TransformFactorization, Expansion
 
 import BlockArrays: block, blockindex, BlockSlice, viewblock
-import BlockBandedMatrices: _BandedBlockBandedMatrix
+import BlockBandedMatrices: _BandedBlockBandedMatrix, AbstractBandedBlockBandedMatrix, _BandedMatrix, blockbandwidths, subblockbandwidths
 import LinearAlgebra: factorize
 import LazyArrays: arguments, paddeddata
 
-import ClassicalOrthogonalPolynomials: jacobimatrix
+import ClassicalOrthogonalPolynomials: jacobimatrix, Weighted, orthogonalityweight
 import HarmonicOrthogonalPolynomials: BivariateOrthogonalPolynomial, MultivariateOrthogonalPolynomial, Plan,
                                           PartialDerivative, BlockOneTo, BlockRange1, interlace
 
 export UnitTriangle, UnitDisk, JacobiTriangle, TriangleWeight, WeightedTriangle, PartialDerivative, Laplacian, 
-      MultivariateOrthogonalPolynomial, BivariateOrthogonalPolynomial, Zernike, RadialCoordinate
+      MultivariateOrthogonalPolynomial, BivariateOrthogonalPolynomial, Zernike, RadialCoordinate,
+      zerniker, zernikez, Weighted, Block, ZernikeWeight
 
 if VERSION < v"1.6-"
       oneto(n) = Base.OneTo(n)

src/disk.jl

@@ -43,25 +43,84 @@ getindex(A::DiskTrav, k::Int) = A[findblockindex(axes(A,1), k)]
 
 ClassicalOrthogonalPolynomials.checkpoints(d::UnitDisk{T}) where T = [SVector{2,T}(0.1,0.2), SVector{2,T}(0.2,0.3)]
 
+"""
+    ZernikeWeight(a, b)
+
+is a quasi-vector representing `r^(2a) * (1-r^2)^b`
+"""
+struct ZernikeWeight{T} <: Weight{T}
+    a::T
+    b::T
+end
+
+"""
+    ZernikeWeight(b)
 
-struct Zernike{T} <: BivariateOrthogonalPolynomial{T} end
+is a quasi-vector representing `(1-r^2)^b`
+"""
+
+ZernikeWeight(b) = ZernikeWeight(zero(b), b)
+
+axes(::ZernikeWeight{T}) where T = (Inclusion(UnitDisk{T}()),)
+
+==(w::ZernikeWeight, v::ZernikeWeight) = w.a == v.a && w.b == v.b
+
+function getindex(w::ZernikeWeight, xy::StaticVector{2})
+    r = norm(xy)
+    r^(2w.a) * (1-r^2)^w.b
+end
 
+
+"""
+    Zernike(a, b)
+
+is a quasi-matrix orthogonal `r^(2a) * (1-r^2)^b`
+"""
+struct Zernike{T} <: BivariateOrthogonalPolynomial{T}
+    a::T
+    b::T
+    Zernike{T}(a::T, b::T) where T = new{T}(a, b)
+end
+Zernike{T}(a, b) where T = Zernike{T}(convert(T,a), convert(T,b))
+Zernike(a::T, b::V) where {T,V} = Zernike{float(promote_type(T,V))}(a, b)
+Zernike{T}(b) where T = Zernike{T}(zero(b), b)
+Zernike{T}() where T = Zernike{T}(zero(T))
+
+"""
+    Zernike(b)
+
+is a quasi-matrix orthogonal `(1-r^2)^b`
+"""
+Zernike(b) = Zernike(zero(b), b)
 Zernike() = Zernike{Float64}()
 
 axes(P::Zernike{T}) where T = (Inclusion(UnitDisk{T}()),blockedrange(oneto(∞)))
 
+==(w::Zernike, v::Zernike) = w.a == v.a && w.b == v.b
+
 copy(A::Zernike) = A
 
+orthogonalityweight(Z::Zernike) = ZernikeWeight(Z.a, Z.b)
+
+zerniker(ℓ, m, a, b, r::T) where T = sqrt(convert(T,2)^(m+a+b+2-iszero(m))/π) * r^m * normalizedjacobip((ℓ-m) ÷ 2, b, m+a, 2r^2-1)
+zerniker(ℓ, m, b, r) = zerniker(ℓ, m, zero(b), b, r)
+zerniker(ℓ, m, r) = zerniker(ℓ, m, zero(r), r)
+
+function zernikez(ℓ, ms, a, b, rθ::RadialCoordinate{T}) where T
+    r,θ = rθ.r,rθ.θ
+    m = abs(ms)
+    zerniker(ℓ, m, a, b, r) * (signbit(ms) ? sin(m*θ) : cos(m*θ))
+end
+
+zernikez(ℓ, ms, a, b, xy::StaticVector{2}) = zernikez(ℓ, ms, a, b, RadialCoordinate(xy))
+zernikez(ℓ, ms, b, xy::StaticVector{2}) = zernikez(ℓ, ms, zero(b), b, xy)
+zernikez(ℓ, ms, xy::StaticVector{2,T}) where T = zernikez(ℓ, ms, zero(T), xy)
+
 function getindex(Z::Zernike{T}, rθ::RadialCoordinate, B::BlockIndex{1}) where T
-    r,θ = rθ.r, rθ.θ
     ℓ = Int(block(B))-1
     k = blockindex(B)
     m = iseven(ℓ) ? k-isodd(k) : k-iseven(k)
-    if iszero(m)
-        sqrt(convert(T,2)/π) * Normalized(Legendre{T}())[2r^2-1, ℓ ÷ 2 + 1]
-    else
-        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
+    zernikez(ℓ, (isodd(k+ℓ) ? 1 : -1) * m, Z.a, Z.b, rθ)
 end
 
 
@@ -70,6 +129,8 @@ getindex(Z::Zernike, xy::StaticVector{2}, B::Block{1}) = [Z[xy, B[j]] for j=1:In
 getindex(Z::Zernike, xy::StaticVector{2}, JR::BlockOneTo) = mortar([Z[xy,Block(J)] for J = 1:Int(JR[end])])
 
 
+
+
 ###
 # Transforms
 ###
@@ -93,10 +154,82 @@ struct ZernikeTransform{T} <: Plan{T}
     analysis::FastTransforms.FTPlan{T,2,FastTransforms.DISKANALYSIS}
 end
 
-function ZernikeTransform{T}(N::Int) where T<:Real
+function ZernikeTransform{T}(N::Int, a::Number, b::Number) where T<:Real
     Ñ = N ÷ 2 + 1
-    ZernikeTransform{T}(N, plan_disk2cxf(T, Ñ, 0, 0), plan_disk_analysis(T, Ñ, 4Ñ-3))
+    ZernikeTransform{T}(N, plan_disk2cxf(T, Ñ, a, b), 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)))
+factorize(S::FiniteZernike{T}) where T = TransformFactorization(grid(S), ZernikeTransform{T}(blocksize(S,2), parent(S).a, parent(S).b))
+
+# gives the entries for the Laplacian times (1-r^2) * Zernike(1)
+struct WeightedZernikeLaplacianDiag{T} <: AbstractBlockVector{T} end
+
+axes(::WeightedZernikeLaplacianDiag) = (blockedrange(oneto(∞)),)
+
+function Base.view(W::WeightedZernikeLaplacianDiag{T}, K::Block{1}) where T
+    K̃ = Int(K)
+    d = K̃÷2 + 1
+    if isodd(K̃)
+        v = (d:K̃) .* (d:(-1):1)
+        convert(AbstractVector{T}, -4*interlace(v, v[2:end]))
+    else
+        v = (d:K̃) .* (d-1:(-1):1)
+        convert(AbstractVector{T}, -4 * interlace(v, v))
+    end
+end
+
+getindex(W::WeightedZernikeLaplacianDiag, k::Integer) = W[findblockindex(axes(W,1),k)]
+
+@simplify function *(Δ::Laplacian, WZ::Weighted{<:Any,<:Zernike})
+    @assert WZ.P.a == 0 && WZ.P.b == 1
+    WZ.P * Diagonal(WeightedZernikeLaplacianDiag{eltype(eltype(WZ))}())
+end
+
+struct ZernikeConversion{T} <: AbstractBandedBlockBandedMatrix{T} end
+
+axes(Z::ZernikeConversion) = (blockedrange(oneto(∞)), blockedrange(oneto(∞)))
+
+blockbandwidths(::ZernikeConversion) = (0,2)
+subblockbandwidths(::ZernikeConversion) = (0,0)
+
+
+function Base.view(W::ZernikeConversion{T}, KJ::Block{2}) where T
+    K,J = KJ.n
+    dat = Matrix{T}(undef,1,J)
+    if J == K
+        if isodd(K)
+            R0 = Normalized(Jacobi(1,0)) \ Normalized(Jacobi(0,0))
+            dat[1,1] = R0[K÷2+1,K÷2+1]
+        end
+        for m in range(2-iseven(K); step=2, length=J÷2)
+            Rm = Normalized(Jacobi(1,m)) \ Normalized(Jacobi(0,m))
+            j = K÷2-m÷2+isodd(K)
+            dat[1,m] = dat[1,m+1] = Rm[j,j]
+        end
+    elseif J == K + 2
+        if isodd(K)
+            R0 = Normalized(Jacobi(1,0)) \ Normalized(Jacobi(0,0))
+            j = K÷2+1
+            dat[1,1] = R0[j,j+1]
+        end
+        for m in range(2-iseven(K); step=2, length=K÷2)
+            Rm = Normalized(Jacobi(1,m)) \ Normalized(Jacobi(0,m))
+            j = K÷2-m÷2+isodd(K)
+            dat[1,m] = dat[1,m+1] = Rm[j,j+1]
+        end
+    else
+        fill!(dat, zero(T))
+    end
+    dat ./= sqrt(2one(T))
+    _BandedMatrix(dat, K, 0, 0)
+end
+
+getindex(R::ZernikeConversion, k::Integer, j::Integer) = R[findblockindex.(axes(R),(k,j))...]
+
+function \(A::Zernike{T}, B::Zernike{V}) where {T,V}
+    A.a == B.a && A.b == B.b && return Eye{promote_type(T,V)}(∞)
+    @assert A.a == 0 && A.b == 1
+    @assert B.a == 0 && B.b == 0
+    ZernikeConversion{promote_type(T,V)}()
+end

src/triangle.jl

@@ -25,6 +25,12 @@ copy(A::JacobiTriangle) = A
 
 Base.summary(io::IO, P::JacobiTriangle) = print(io, "JacobiTriangle($(P.a), $(P.b), $(P.c))")
 
+
+"""
+    TriangleWeight(a, b, c)
+
+is a quasi-vector representing `x^a * y^b * (1-x-y)^c`.
+"""
 struct TriangleWeight{T,V} <: Weight{T}
     a::V
     b::V

test/runtests.jl

@@ -3,105 +3,3 @@ using MultivariateOrthogonalPolynomials, Test
 include("test_disk.jl")
 include("test_triangle.jl")
 # include("test_dirichlettriangle.jl")
-
-
-##  bessel
-
-# for k=0:10
-#     @test Fun(r->besselj(k,r),JacobiSquare(k))(0.9) ≈ besselj(k,0.9)
-# end
-#
-#
-# f=Fun([1.],Segment()^2)
-#     f(0.1,0.2)
-#
-#
-# f=Fun([1.],Disk())
-# @test f(0.1,0.2) ≈ 1.
-#
-# f=ProductFun(Fun([0.,1.]+0.0im,Disk()))
-# x,y=0.1,0.2
-# r,θ=sqrt(x^2+y^2),atan2(y,x)
-# @test exp(-im*θ)*r ≈ f(x,y)
-#
-#
-# ## Disk
-#
-#
-#
-#
-# f=(x,y)->exp(x.*sin(y))
-# u=ProductFun(f,Disk(),50,51)
-# @test u(.1,.1) ≈ f(.1,.1)
-#
-#
-#
-#
-# # write your own tests here
-# # Laplace
-# d=Disk()
-# u=[dirichlet(d),lap(d)]\Fun(z->real(exp(z)),Circle())
-# @test u(.1,.2) ≈ real(exp(.1+.2im))
-#
-#
-#
-#
-# # remaining numbers determined numerically, may be
-# # inaccurate
-#
-# # Poisson
-# f=Fun((x,y)->exp(-10(x+.2).^2-20(y-.1).^2),d)
-# u=[dirichlet(d),lap(d)]\[0.,f]
-# @test u(.1,.2) ≈ -0.039860694987858845
-#
-# #Helmholtz
-# u=[dirichlet(d),lap(d)+100I]\1.0
-# @test_approx_eq_eps u(.1,.2) -0.3675973169667076 1E-11
-# u=[neumann(d),lap(d)+100I]\1.0
-# @test_approx_eq_eps u(.1,.2) -0.20795862954551195 1E-11
-#
-# # Screened Poisson
-# u=[neumann(d),lap(d)-100.0I]\1.0
-# @test_approx_eq_eps u(.1,.9) 0.04313812031635443 1E-11
-#
-# # Lap^2
-# u=[dirichlet(d),neumann(d),lap(d)^2]\Fun(z->real(exp(z)),Circle())
-# @test_approx_eq_eps u(.1,.2) 1.1137317420521624 1E-11
-#
-#
-# # Speed Test
-#
-#
-# d = Disk()
-# f = Fun((x,y)->exp(-10(x+.2)^2-20(y-.1)^2),d)
-# S = discretize([dirichlet(d);lap(d)],100);
-# @time S = discretize([dirichlet(d);lap(d)],100);
-# u=S\Any[0.;f];
-# @time u=S\Any[0.;f];
-#
-# println("Disk Poisson: should be ~0.16,0.016")
-#
-#
-#
-# ## README Test
-#
-# d = Disk()
-# f = Fun((x,y)->exp(-10(x+.2)^2-20(y-.1)^2),d)
-# u = [dirichlet(d);lap(d)]\Any[0.,f]
-#
-#
-#
-# d = Disk()
-# u0 = Fun((x,y)->exp(-50x^2-40(y-.1)^2)+.5exp(-30(x+.5)^2-40(y+.2)^2),d)
-# B= [dirichlet(d);neumann(d)]
-# L = -lap(d)^2
-# h = 0.001
-#
-# L=Laplacian(Space(d),1)
-# @show real(ApproxFun.diagop(L,1)[1:10,1:10])
-#
-# d = Disk()
-# u0 = Fun((x,y)->exp(-50x^2-40(y-.1)^2)+.5exp(-30(x+.5)^2-40(y+.2)^2),d)
-# B= [dirichlet(d);neumann(d)]
-# L = -lap(d)^2
-# h = 0.001