diff --git a/src/MultivariateOrthogonalPolynomials.jl b/src/MultivariateOrthogonalPolynomials.jl index f3679fe..b012858 100644 --- a/src/MultivariateOrthogonalPolynomials.jl +++ b/src/MultivariateOrthogonalPolynomials.jl @@ -21,7 +21,7 @@ import LazyArrays: arguments, paddeddata, LazyArrayStyle, LazyLayout, PaddedLayo import LazyBandedMatrices: LazyBandedBlockBandedLayout, AbstractBandedBlockBandedLayout, AbstractLazyBandedBlockBandedLayout, _krontrav_axes, DiagTravLayout import InfiniteArrays: InfiniteCardinal, OneToInf -import ClassicalOrthogonalPolynomials: jacobimatrix, Weighted, orthogonalityweight, HalfWeighted, WeightedBasis, pad, recurrencecoefficients, clenshaw +import ClassicalOrthogonalPolynomials: jacobimatrix, Weighted, orthogonalityweight, HalfWeighted, WeightedBasis, pad, recurrencecoefficients, clenshaw, massmatrix import HarmonicOrthogonalPolynomials: BivariateOrthogonalPolynomial, MultivariateOrthogonalPolynomial, Plan, PartialDerivative, AngularMomentum, BlockOneTo, BlockRange1, interlace, MultivariateOPLayout, MAX_PLOT_BLOCKS diff --git a/src/disk.jl b/src/disk.jl index 4595766..3218405 100644 --- a/src/disk.jl +++ b/src/disk.jl @@ -342,4 +342,76 @@ 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 +### + +@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) # Fourier mode number + + # Raising Fourier mode coefficients, split into sin and cos parts + m=0 change + N = sqrt.(massmatrix(Jacobi{T}(0,1)).diag) + c=-BroadcastVector(n->N[n+1], (n .- m) .÷ 2) ./ BroadcastVector(n->N[n+1], (n .- m) .÷ 2 .+ m) + c1 = c .* ((n .- m) .÷ 2 .+ 1) .* (-1).^(iseven.(k.-n)) .* (iszero.(m) * sqrt(2*one(T)) + (1 .- iszero.(m))) + c2 = c1 .* iseven.(n .- k) + c3 = c1 .* isodd.(n .- k) + + # Lowering Fourier mode coefficients, split into sin and cos parts + m=1 change + d1 = c .* ((n .- m) .÷ 2 .+ 1) .* (-1).^(isodd.(k.-n)) .* (isone.(m) * sqrt(2*one(T)) + (1 .- isone.(m))) + d2 = d1 .* iseven.(n .- k) + d3 = d1 .* isodd.(n .- k) + + # Put coefficients together + 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) # Fourier mode number + + # Computes the entries for the component that lowers the Fourier mode + D = BroadcastVector(P->diff(P; dims=1), HalfWeighted{:b}.(Normalized.(Jacobi.(b,1:∞)))) + Ds = BroadcastVector{AbstractMatrix{T}}((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))) + + # Computes the entries for the component that lowers the Fourier mode + # the -1 in the parameters is a trick to make ModalInterlace think that the + # 0-parameter is the second Fourier mode (we use the -1 as a dummy matrix in the ModalInterlace) + D = BroadcastVector(P->diff(P; dims=1), Normalized.(Jacobi.(b,-1:∞))) + Ds = BroadcastVector{AbstractMatrix{T}}((P,D) -> P \ D, Normalized.(Jacobi.(b+1,0:∞)), D) + Dss = BroadcastVector{AbstractMatrix{T}}(P -> Diagonal(view(P, band(1))), Ds) + M = ModalInterlace(Dss, (ℵ₀,ℵ₀), (0,0)) + d = ones(axes(n)) .* view(view(M, 1:∞, 1:∞),band(0)) + + # Lowering Fourier mode coefficients, split into sin and cos parts + m=1 change + c1 = d .* (-1).^(isodd.(k.-n)) .* (isone.(m) .* sqrt(2*one(T)) + (1 .- isone.(m))) + c2 = c1 .* iseven.(n .- k) + c3 = c1 .* isodd.(n .- k) + + # Raising Fourier mode coefficients, split into sin and cos parts + m=0 change + d1 = db .* (-1).^(iseven.(k.-n)) .* (iszero.(m) * sqrt(2*one(T)) + (1 .- iszero.(m))) + d2 = d1 .* iseven.(n .- k) + d3 = d1 .* isodd.(n .- k) + + # Put coefficients together + A = BlockBandedMatrices._BandedBlockBandedMatrix(BlockBroadcastArray(hcat, c3, Zeros{T}((axes(n,1),)), c2)', axes(n,1), (1,-1), (0,2)) + B = BlockBandedMatrices._BandedBlockBandedMatrix(BlockBroadcastArray(hcat, d3, Zeros{T}((axes(n,1),)), d2)', axes(n,1), (1,-1), (2,0)) + + Zernike{T}(Z.b+1) * (A+B)' end \ No newline at end of file diff --git a/test/test_disk.jl b/test/test_disk.jl index 2cc358f..f884047 100644 --- a/test/test_disk.jl +++ b/test/test_disk.jl @@ -1,7 +1,7 @@ using MultivariateOrthogonalPolynomials, ClassicalOrthogonalPolynomials, StaticArrays, BlockArrays, BandedMatrices, FastTransforms, LinearAlgebra, RecipesBase, Test, SpecialFunctions, LazyArrays, InfiniteArrays import MultivariateOrthogonalPolynomials: ModalTrav, grid, ZernikeTransform, ZernikeITransform, *, ModalInterlace import ClassicalOrthogonalPolynomials: HalfWeighted, expand -import ForwardDiff: hessian +import ForwardDiff: hessian, gradient @testset "Disk" begin @testset "Transform" begin @@ -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⁰ = 𝐱 -> 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¹ = 𝐱 -> 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² = 𝐱 -> 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