diff --git a/src/HarmonicOrthogonalPolynomials.jl b/src/HarmonicOrthogonalPolynomials.jl index 7af6175..eca726e 100644 --- a/src/HarmonicOrthogonalPolynomials.jl +++ b/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))) diff --git a/src/multivariateops.jl b/src/multivariateops.jl index 3e69b4d..656cdc4 100644 --- a/src/multivariateops.jl +++ b/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) diff --git a/test/runtests.jl b/test/runtests.jl index 0cd999a..74fdf1f 100644 --- a/test/runtests.jl +++ b/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]] +