Add angularmomentum

Author: dlfivefifty

src/HarmonicOrthogonalPolynomials.jl

@@ -5,8 +5,10 @@ import Base: OneTo, oneto, axes, getindex, convert, to_indices, tail, eltype, *,
 import BlockArrays: block, blockindex, unblock, BlockSlice
 import DomainSets: indomain, Sphere
 import LinearAlgebra: norm, factorize
-import QuasiArrays: to_quasi_index, SubQuasiArray, *
-import ContinuumArrays: TransformFactorization, @simplify, ProjectionFactorization, plan_grid_transform, plan_transform, grid, grid_layout, plotgrid_layout, AbstractBasisLayout, MemoryLayout, abslaplacian, laplacian
+import QuasiArrays: to_quasi_index, SubQuasiArray, *, AbstractQuasiVecOrMat
+import ContinuumArrays: TransformFactorization, @simplify, ProjectionFactorization, plan_grid_transform, plan_transform, grid, grid_layout, plotgrid_layout,
+                        AbstractBasisLayout, MemoryLayout, abslaplacian, laplacian, AbstractDifferentialQuasiMatrix, operatorcall, similaroperator, SubBasisLayout,
+                        ApplyLayout, arguments, ExpansionLayout
 import ClassicalOrthogonalPolynomials: checkpoints, _sum, cardinality, increasingtruncations
 import BlockBandedMatrices: BlockRange1, _BandedBlockBandedMatrix
 import FastTransforms: Plan, interlace

src/angularmomentum.jl

@@ -5,21 +5,50 @@ Applies the partial derivative with respect to the last angular variable in the
 For example, in polar coordinates (r, θ) in ℝ² or cylindrical coordinates (r, θ, z) in ℝ³, we apply ∂ / ∂θ = (x ∂ / ∂y - y ∂ / ∂x).
 In spherical coordinates (ρ, θ, φ) in ℝ³, we apply ∂ / ∂φ = (x ∂ / ∂y - y ∂ / ∂x).
 """
-struct AngularMomentum{T,Ax<:Inclusion} <: LazyQuasiMatrix{T}
+struct AngularMomentum{T,Ax<:Inclusion,Order} <: AbstractDifferentialQuasiMatrix{T}
     axis::Ax
+    order::Order
 end
 
-AngularMomentum{T}(axis::Inclusion) where T = AngularMomentum{T,typeof(axis)}(axis)
-AngularMomentum{T}(domain) where T = AngularMomentum{T}(Inclusion(domain))
-AngularMomentum(axis) = AngularMomentum{eltype(eltype(axis))}(axis)
+AngularMomentum{T, D}(axis::D, order) where {T,D<:Inclusion} = AngularMomentum{T,D,typeof(order)}(axis, order)
+AngularMomentum{T, D}(axis::D) where {T,D<:Inclusion} = AngularMomentum{T,D,Nothing}(axis, nothing)
+AngularMomentum{T}(axis::Inclusion, order...) where T = AngularMomentum{T,typeof(axis)}(axis, order...)
+AngularMomentum{T}(domain, order...) where T = AngularMomentum{T}(Inclusion(domain), order...)
+AngularMomentum(domain, order...) = AngularMomentum(Inclusion(domain), order...)
+AngularMomentum(ax::AbstractQuasiVector{T}, order...) where T = AngularMomentum{eltype(eltype(ax))}(ax, order...)
+AngularMomentum(L::AbstractQuasiMatrix, order...) = AngularMomentum(axes(L,1), order...)
 
-axes(A::AngularMomentum) = (A.axis, A.axis)
-==(a::AngularMomentum, b::AngularMomentum) = a.axis == b.axis
-copy(A::AngularMomentum) = AngularMomentum(copy(A.axis))
 
-^(A::AngularMomentum, k::Integer) = ApplyQuasiArray(^, A, k)
+operatorcall(::AngularMomentum) = angularmomentum
+similaroperator(D::AngularMomentum, ax, ord) = AngularMomentum{eltype(D)}(ax, ord)
 
-@simplify function *(A::AngularMomentum, P::SphericalHarmonic)
+simplifiable(::typeof(*), A::AngularMomentum, B::AngularMomentum) = Val(true)
+*(D1::AngularMomentum, D2::AngularMomentum) = similaroperator(convert(AbstractQuasiMatrix{promote_type(eltype(D1),eltype(D2))}, D1), D1.axis, operatororder(D1)+operatororder(D2))
+
+angularmomentum(A, order...; dims...) = angularmomentum_layout(MemoryLayout(A), A, order...; dims...)
+
+angularmomentum_layout(::AbstractBasisLayout, Vm, order...; dims...) = error("Overload angularmomentum(::$(typeof(Vm)))")
+function angularmomentum_layout(::AbstractBasisLayout, a, order::Int; dims...)
+    order < 0 && throw(ArgumentError("order must be non-negative"))
+    order == 0 && return a
+    isone(order) ? angularmomentum(a) : angularmomentum(angularmomentum(a), order-1)
+end
+
+angularmomentum_layout(::ExpansionLayout, A, order...; dims...) = angularmomentum_layout(ApplyLayout{typeof(*)}(), A, order...; dims...)
+
+
+function angularmomentum_layout(::SubBasisLayout, Vm, order...; dims::Integer=1)
+    dims == 1 || error("not implemented")
+    angularmomentum(parent(Vm), order...)[:,parentindices(Vm)[2]]
+end
+
+function angularmomentum_layout(LAY::ApplyLayout{typeof(*)}, V::AbstractQuasiVecOrMat, order...; dims=1)
+    a = arguments(LAY, V)
+    dims == 1 || throw(ArgumentError("cannot take angularmomentum a vector along dimension $dims"))
+    *(angularmomentum(a[1], order...), tail(a)...)
+end
+
+function angularmomentum(P::SphericalHarmonic)
     # Spherical harmonics are the eigenfunctions of the angular momentum operator on the unit sphere
     T = real(eltype(P))
     k = mortar(Base.OneTo.(1:2:∞))

test/runtests.jl

@@ -417,9 +417,9 @@ end
 @testset "Angular momentum" begin
     S = SphericalHarmonic()
     R = RealSphericalHarmonic()
-    ∂θ = AngularMomentum(axes(S, 1))
+    ∂θ = AngularMomentum(S)
     @test axes(∂θ) == (axes(S, 1), axes(S, 1))
-    @test ∂θ == AngularMomentum(axes(R, 1)) == AngularMomentum(axes(S, 1).domain)
+    @test ∂θ == AngularMomentum(R) == AngularMomentum(axes(S, 1).domain)
     @test copy(∂θ) ≡ ∂θ
     A = S \ (∂θ * S)
     A2 = S \ (∂θ^2 * S)