overload _getindex instead of getindex (#97)

* overload _getindex instead of getindex

* import _getindex

* reorg

* Update runtests.jl

Author: dlfivefifty

src/HarmonicOrthogonalPolynomials.jl

@@ -12,7 +12,7 @@ import ContinuumArrays: TransformFactorization, @simplify, ProjectionFactorizati
 import ClassicalOrthogonalPolynomials: checkpoints, _sum, cardinality, increasingtruncations
 import BlockBandedMatrices: BlockRange1, _BandedBlockBandedMatrix
 import FastTransforms: Plan, interlace
-import QuasiArrays: LazyQuasiMatrix, LazyQuasiArrayStyle
+import QuasiArrays: LazyQuasiMatrix, LazyQuasiArrayStyle, _getindex
 import InfiniteArrays: InfStepRange, RangeCumsum
 
 export SphericalHarmonic, UnitSphere, SphericalCoordinate, RadialCoordinate, Block, associatedlegendre, RealSphericalHarmonic, sphericalharmonicy, abs, -, ^, AngularMomentum, Laplacian, AbsLaplacian
@@ -86,7 +86,6 @@ end
 getindex(S::AbstractSphericalHarmonic, x::StaticVector{3}, K::BlockIndex{1}) = S[SphericalCoordinate(x), K]
 getindex(S::AbstractSphericalHarmonic, x::StaticVector{3}, K::Block{1}) = S[x, axes(S,2)[K]]
 getindex(S::AbstractSphericalHarmonic, x::StaticVector{3}, KR::BlockOneTo) = mortar([S[x, K] for K in KR])
-getindex(S::AbstractSphericalHarmonic, x::StaticVector{3}, k::Int) = S[x, findblockindex(axes(S,2), k)]
 getindex(S::AbstractSphericalHarmonic, x::StaticVector{3}, kr::AbstractUnitRange{Int}) = [S[x, k] for k in kr]
 
 # @simplify *(Ac::QuasiAdjoint{<:Any,<:SphericalHarmonic}, B::SphericalHarmonic) =
@@ -128,9 +127,6 @@ RealSphericalHarmonicTransform{T}(N::Int) where T<:Real = RealSphericalHarmonicT
 plan_transform(P::SphericalHarmonic{T}, (N,)::Tuple{Block{1}}, dims=1) where T = SphericalHarmonicTransform{T}(Int(N))
 plan_transform(P::RealSphericalHarmonic{T}, (N,)::Tuple{Block{1}}, dims=1) where T = RealSphericalHarmonicTransform{T}(Int(N))
 
-grid(P::MultivariateOrthogonalPolynomial, n::Int) = grid(P, findblock(axes(P,2),n))
-plan_transform(P::MultivariateOrthogonalPolynomial, Bs::NTuple{N,Int}, dims=ntuple(identity,Val(N))) where N = plan_transform(P, findblock.(Ref(axes(P,2)), Bs), dims)
-
 function _sum(A::AbstractSphericalHarmonic{T}, dims) where T
     @assert dims == 1
     BlockedArray(Hcat(sqrt(4convert(T, π)), Zeros{T}(1,∞)), (Base.OneTo(1),axes(A,2)))

src/multivariateops.jl

@@ -10,12 +10,12 @@ const BlockOneTo = BlockRange{1,Tuple{OneTo{Int}}}
 
 copy(P::MultivariateOrthogonalPolynomial) = P
 
-getindex(P::MultivariateOrthogonalPolynomial{D}, xy::StaticVector{D}, JR::BlockOneTo) where D = error("Overload")
-getindex(P::MultivariateOrthogonalPolynomial{D}, xy::StaticVector{D}, J::Block{1}) where D = P[xy, Block.(OneTo(Int(J)))][J]
-getindex(P::MultivariateOrthogonalPolynomial{D}, xy::StaticVector{D}, JR::BlockRange{1}) where D = P[xy, Block.(OneTo(Int(maximum(JR))))][JR]
-getindex(P::MultivariateOrthogonalPolynomial{D}, xy::StaticVector{D}, Jj::BlockIndex{1}) where D = P[xy, block(Jj)][blockindex(Jj)]
-getindex(P::MultivariateOrthogonalPolynomial{D}, xy::StaticVector{D}, j::Integer) where D = P[xy, findblockindex(axes(P,2), j)]
-getindex(P::MultivariateOrthogonalPolynomial{D}, xy::StaticVector{D}, jr::AbstractVector{<:Integer}) where D = P[xy, Block.(OneTo(Int(findblock(axes(P,2), maximum(jr)))))][jr]
+_getindex(::Type{Tuple{IND1,IND2}}, P::MultivariateOrthogonalPolynomial, (𝐱,JR)::Tuple{IND1,BlockOneTo}) where {IND1,IND2} = error("Overload")
+_getindex(::Type{Tuple{IND1,IND2}}, P::MultivariateOrthogonalPolynomial, (𝐱,J)::Tuple{IND1,Block{1}}) where {IND1,IND2} = P[𝐱, Block.(OneTo(Int(J)))][J]
+_getindex(::Type{Tuple{IND1,IND2}}, P::MultivariateOrthogonalPolynomial, (𝐱,JR)::Tuple{IND1,BlockRange{1}}) where {IND1,IND2} = P[𝐱, Block.(OneTo(Int(maximum(JR))))][JR]
+_getindex(::Type{Tuple{IND1,IND2}}, P::MultivariateOrthogonalPolynomial, (𝐱,Jj)::Tuple{IND1,BlockIndex{1}}) where {IND1,IND2} = P[𝐱, block(Jj)][blockindex(Jj)]
+_getindex(::Type{Tuple{IND1,IND2}}, P::MultivariateOrthogonalPolynomial, (𝐱,j)::Tuple{IND1,IND2}) where {IND1,IND2} = P[𝐱, findblockindex(axes(P,2), j)]
+_getindex(::Type{Tuple{IND1,IND2}}, P::MultivariateOrthogonalPolynomial, (𝐱,jr)::Tuple{IND1,AbstractArray{IND2}}) where {IND1,IND2} = P[𝐱, Block.(OneTo(Int(findblock(axes(P,2), maximum(jr)))))][jr]
 
 const FirstInclusion = BroadcastQuasiVector{<:Any, typeof(first), <:Tuple{Inclusion}}
 const LastInclusion = BroadcastQuasiVector{<:Any, typeof(last), <:Tuple{Inclusion}}
@@ -91,3 +91,7 @@ const MAX_PLOT_BLOCKS = 200
 grid_layout(::AbstractMultivariateOPLayout, S, n::Integer) = grid(S, findblock(axes(S,2), n))
 plotgrid_layout(::AbstractMultivariateOPLayout, S, n::Integer) = plotgrid(S, findblock(axes(S,2), n))
 plotgrid_layout(::AbstractMultivariateOPLayout, S, B::Block{1}) = grid(S, min(2B, Block(MAX_PLOT_BLOCKS)))
+
+
+grid(P::MultivariateOrthogonalPolynomial, n::Int) = grid(P, findblock(axes(P,2),n))
+plan_transform(P::MultivariateOrthogonalPolynomial, Bs::NTuple{N,Int}, dims=ntuple(identity,Val(N))) where N = plan_transform(P, findblock.(Ref(axes(P,2)), Bs), dims)

test/runtests.jl

@@ -1,5 +1,5 @@
 using HarmonicOrthogonalPolynomials, StaticArrays, Test, InfiniteArrays, LinearAlgebra, BlockArrays, ClassicalOrthogonalPolynomials, QuasiArrays
-import HarmonicOrthogonalPolynomials: ZSphericalCoordinate, associatedlegendre, grid, SphereTrav, RealSphereTrav, plotgrid
+import HarmonicOrthogonalPolynomials: ZSphericalCoordinate, associatedlegendre, grid, SphereTrav, RealSphereTrav, plotgrid, BivariateOrthogonalPolynomial
 
 # @testset "associated legendre" begin
 #     m = 2
@@ -74,7 +74,7 @@ end
                           0.25sqrt(105/2π)sin(θ)^2*cos(θ)*exp(2im*φ),
                           0.125sqrt(35/π)sin(θ)^3*exp(3im*φ)]
 
-    @test S[x,Block.(1:4)] == [S[x,Block(1)]; S[x,Block(2)]; S[x,Block(3)]; S[x,Block(4)]]
+    @test S[x,Block.(1:4)] == S[x,Block.(Base.OneTo(4))] == [S[x,Block(1)]; S[x,Block(2)]; S[x,Block(3)]; S[x,Block(4)]] 
 end
 
 @testset "Real Evaluation" begin
@@ -428,3 +428,13 @@ end
     @test isdiag(A[1:N, 1:N])
     @test A[1:N, 1:N]^2 ≈ A2[1:N, 1:N]
 end
+
+
+struct IncompleteMultivariateOP <: BivariateOrthogonalPolynomial{Float64} end
+Base.axes(::IncompleteMultivariateOP) = Inclusion((-1.0..1)^2), blockedrange(Base.oneto(∞))
+
+@test_throws "Overload" IncompleteMultivariateOP()[SVector(0.1,0.2),2]
+@test_throws "Overload" IncompleteMultivariateOP()[SVector(0.1,0.2),Block(2)]
+@test_throws "Overload" IncompleteMultivariateOP()[SVector(0.1,0.2),Block(2)[2]]
+@test_throws "Overload" IncompleteMultivariateOP()[SVector(0.1,0.2),[1,2]]
+