Jp/zernike partials

src/disk.jl

@@ -342,4 +342,71 @@ function Base._sum(P::Zernike{T}, dims) where T
     @assert dims == 1
     @assert P.a == P.b == 0
     Hcat(sqrt(convert(T, π)), Zeros{T}(1,∞))
+end
+
+
+###
+# Partial derivatives
+###
+
+normal_jacobi(a::T,b::T,n::Int) where T = sqrt(2^(a+b+1) / (2n + a + b + 1) * gamma(n+a+1)*gamma(n+b+1) / (gamma(n+a+b+1) * factorial(n)))
+
+@simplify function *(∂ʸ::PartialDerivative{2}, WZ::Weighted{<:Any,<:Zernike})
+    @assert WZ.P.a == 0 && WZ.P.b == 1
+    T = eltype(eltype(WZ))
+
+    k = mortar(Base.OneTo.(oneto(∞)))     # k counts the the angular mode (+1)
+    n = mortar(Fill.(oneto(∞),oneto(∞))) .- 1  # n counts the block number which corresponds to the order
+    m = k .- isodd.(k).*iseven.(n) .- iseven.(k).*isodd.(n)
+
+    c=-(normal_jacobi.(T(0), T(1), (n .- m) .÷ 2) ./ normal_jacobi.(T(1), T(0), (n .- m) .÷ 2 .+ m))
+    c1 = c .* ((n  .- m) .÷ 2 .+ 1) .* (-1).^(iseven.(k.-n)) .* (iszero.(m) * sqrt(T(2)) + (1 .- iszero.(m)))
+    c2 = c1 .* iseven.(n .- k)
+    c3 = c1 - c2
+
+    d1 = c .* ((n  .- m) .÷ 2 .+ 1) .* (-1).^(isodd.(k.-n)) .* (isone.(m) * sqrt(T(2)) + (1 .- isone.(m)))
+    d2 = d1 .* iseven.(n .- k)
+    d3 = d1 - d2
+
+    A = BlockBandedMatrices._BandedBlockBandedMatrix(BlockBroadcastArray(hcat, c3, Zeros((axes(n,1),)), c2)', axes(n,1), (1,-1), (2,0))
+    B = BlockBandedMatrices._BandedBlockBandedMatrix(BlockBroadcastArray(hcat, d3, Zeros((axes(n,1),)), d2)', axes(n,1), (1,-1), (0,2))
+    Zernike{T}(0) * (A + B)
+end
+
+@simplify function *(∂ʸ::PartialDerivative{2}, Z::Zernike)
+    @assert Z.a == 0
+    T = eltype(eltype(Z))
+    b = convert(T, Z.b)
+
+    k = mortar(Base.OneTo.(oneto(∞)))     # k counts the the angular mode (+1)
+    n = mortar(Fill.(oneto(∞),oneto(∞))) .- 1  # n counts the block number which corresponds to the order
+    m = k .- isodd.(k).*iseven.(n) .- iseven.(k).*isodd.(n)
+
+    x=axes(Jacobi(0,0),1)
+    D = BroadcastVector(P->Derivative(x) * P, HalfWeighted{:b}.(Normalized.(Jacobi.(b,1:∞))))
+    Ds = BroadcastVector{AbstractMatrix{Float64}}((P,D) -> P \ D, HalfWeighted{:b}.(Normalized.(Jacobi.(b+1,0:∞))) , D)
+    M = ModalInterlace(Ds, (ℵ₀,ℵ₀), (0,0))
+    db = ones(axes(n)) .* (view(view(M, 1:∞, 1:∞),band(0)))
+
+    D = BroadcastVector(P->Derivative(x) * P, Normalized.(Jacobi.(b,-1:∞)))
+    Ds = BroadcastVector{AbstractMatrix{Float64}}((P,D) -> P \ D, Normalized.(Jacobi.(b+1,0:∞)), D)
+    Dss = BroadcastVector{AbstractMatrix{Float64}}(P -> Diagonal(view(P, band(1))), Ds)
+    M = ModalInterlace(Dss, (ℵ₀,ℵ₀), (0,0))
+    # d = ones(axes(n)) .*  view(Vcat(0,view(view(M, 1:∞, 1:∞),band(0))),1:∞)
+    d = ones(axes(n)) .*  view(view(M, 1:∞, 1:∞),band(0))
+
+    c1 = d .* (-1).^(isodd.(k.-n)) .* (isone.(m) .* sqrt(T(2)) + (1 .- isone.(m)))
+    c2 = c1 .* iseven.(n .- k)
+    c3 = c1 - c2
+
+
+    d1 = db .* (-1).^(iseven.(k.-n)) .* (iszero.(m) * sqrt(T(2)) + (1 .- iszero.(m)))
+    d2 = d1 .* iseven.(n .- k)
+    d3 = d1 - d2
+
+
+    A = BlockBandedMatrices._BandedBlockBandedMatrix(BlockBroadcastArray(hcat, c3, Zeros((axes(n,1),)), c2)', axes(n,1), (1,-1), (0,2))
+    B = BlockBandedMatrices._BandedBlockBandedMatrix(BlockBroadcastArray(hcat, d3, Zeros((axes(n,1),)), d2)', axes(n,1), (1,-1), (2,0))
+
+    Zernike{T}(Z.b+1) * (A+B)'
 end

test/test_disk.jl

@@ -319,6 +319,38 @@ import ForwardDiff: hessian
         g = MultivariateOrthogonalPolynomials.plotgrid(W[:,1:3])
         @test all(rep[1].args .≈ (first.(g),last.(g),u[g]))
     end
+
+    @testset "partial derivatives" begin
+        W = Weighted(Zernike(1))
+        Z⁰ = Zernike(0)
+        Z¹ = Zernike(1)
+        Z² = Zernike(2)
+
+        𝐱 = axes(W,1)
+        # ∂ˣ = PartialDerivative{1}(𝐱)
+        ∂ʸ = PartialDerivative{2}(𝐱)
+
+        ∂Y⁰ = Z⁰ \ (∂ʸ * W)
+        ∂Y¹ = Z¹ \ (∂ʸ * Z⁰)
+        ∂Y² = Z² \ (∂ʸ * Z¹)
+
+        B = Block.(1:10); xy = SVector(0.2,0.3)
+
+        for (i,n,m) in zip((1,2,3,4,5,6,14,17), (0,1,1,2,2,2,4,5), (0,-1,1,0,-2,2,-4,1))
+            g⁰ =  𝐱 -> ForwardDiff.gradient(𝐱 -> (1-norm(𝐱)^2)*zernikez(n, m, 1, 𝐱), 𝐱)[2]
+            c⁰ = ModalTrav(transform(Z⁰, g⁰)[B])[B]
+            @test Z⁰[xy,B]'*∂Y⁰[B,i] ≈ Z⁰[xy, B]' * c⁰
+
+            g¹ =  𝐱 -> ForwardDiff.gradient(𝐱 -> zernikez(n, m, 0, 𝐱), 𝐱)[2]
+            c¹ = ModalTrav(transform(Z¹, g¹)[B])[B]
+            @test Z¹[xy,B]'*∂Y¹[B,i] ≈ Z¹[xy, B]' * c¹
+
+            g² =  𝐱 -> ForwardDiff.gradient(𝐱 -> zernikez(n, m, 1, 𝐱), 𝐱)[2]
+            c² = ModalTrav(transform(Z², g²)[B])[B]
+            @test Z²[xy,B]'*∂Y²[B,i] ≈ Z²[xy, B]' * c²
+        end
+
+    end
 end
 
 @testset "Fractional Laplacian on Unit Disk" begin