overload broadcast for RectPoly

Author: dlfivefifty

Project.toml

@@ -19,7 +19,6 @@ LazyArrays = "5078a376-72f3-5289-bfd5-ec5146d43c02"
 LazyBandedMatrices = "d7e5e226-e90b-4449-9968-0f923699bf6f"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 QuasiArrays = "c4ea9172-b204-11e9-377d-29865faadc5c"
-RecurrenceRelationshipArrays = "b889d2dc-af3c-4820-88a8-238fa91d3518"
 RecurrenceRelationships = "807425ed-42ea-44d6-a357-6771516d7b2c"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
@@ -35,12 +34,13 @@ ContinuumArrays = "0.18"
 DomainSets = "0.7"
 FastTransforms = "0.17"
 FillArrays = "1.0"
-HarmonicOrthogonalPolynomials = "0.6"
+HarmonicOrthogonalPolynomials = "0.6.3"
 InfiniteArrays = "0.15"
 InfiniteLinearAlgebra = "0.9"
 LazyArrays = "2.3.1"
 LazyBandedMatrices = "0.11.1"
 QuasiArrays = "0.11"
+RecurrenceRelationships = "0.2"
 SpecialFunctions = "1, 2"
 StaticArrays = "1"
 julia = "1.10"

src/MultivariateOrthogonalPolynomials.jl

@@ -11,20 +11,20 @@ import Base.Broadcast: Broadcasted, broadcasted, DefaultArrayStyle
 import DomainSets: boundary, EuclideanDomain
 
 import QuasiArrays: LazyQuasiMatrix, LazyQuasiArrayStyle, domain
-import ContinuumArrays: @simplify, Weight, weight, grid, plotgrid, TransformFactorization, ExpansionLayout, plotvalues, unweighted, plan_transform, checkpoints, transform_ldiv, AbstractBasisLayout, basis_axes, Inclusion, grammatrix, weaklaplacian
+import ContinuumArrays: @simplify, Weight, weight, grid, plotgrid, TransformFactorization, ExpansionLayout, plotvalues, unweighted, plan_transform, checkpoints, transform_ldiv, AbstractBasisLayout, basis_axes, Inclusion, grammatrix, weaklaplacian, layout_broadcasted
 
 import ArrayLayouts: MemoryLayout, sublayout, sub_materialize
 import BlockArrays: block, blockindex, BlockSlice, viewblock, blockcolsupport, AbstractBlockStyle, BlockStyle
 import BlockBandedMatrices: _BandedBlockBandedMatrix, AbstractBandedBlockBandedMatrix, _BandedMatrix, blockbandwidths, subblockbandwidths
 import LinearAlgebra: factorize
 import LazyArrays: arguments, paddeddata, LazyArrayStyle, LazyLayout, PaddedLayout, applylayout, LazyMatrix, ApplyMatrix
-import LazyBandedMatrices: LazyBandedBlockBandedLayout, AbstractBandedBlockBandedLayout, AbstractLazyBandedBlockBandedLayout, _krontrav_axes, DiagTravLayout, invdiagtrav, ApplyBandedBlockBandedLayout
+import LazyBandedMatrices: LazyBandedBlockBandedLayout, AbstractBandedBlockBandedLayout, AbstractLazyBandedBlockBandedLayout, _krontrav_axes, DiagTravLayout, invdiagtrav, ApplyBandedBlockBandedLayout, krontrav
 import InfiniteArrays: InfiniteCardinal, OneToInf
 
 import ClassicalOrthogonalPolynomials: jacobimatrix, Weighted, orthogonalityweight, HalfWeighted, WeightedBasis, pad, recurrencecoefficients, clenshaw, weightedgrammatrix, Clenshaw
 import HarmonicOrthogonalPolynomials: BivariateOrthogonalPolynomial, MultivariateOrthogonalPolynomial, Plan,
                                           PartialDerivative, AngularMomentum, BlockOneTo, BlockRange1, interlace,
-                                          MultivariateOPLayout, MAX_PLOT_BLOCKS
+                                          MultivariateOPLayout, AbstractMultivariateOPLayout, MAX_PLOT_BLOCKS
 
 export MultivariateOrthogonalPolynomial, BivariateOrthogonalPolynomial,
        UnitTriangle, UnitDisk,

src/clenshawkron.jl

@@ -23,6 +23,9 @@ axes(M::ClenshawKron) = (blockedrange(oneto(∞)),blockedrange(oneto(∞)))
 blockbandwidths(M::ClenshawKron) = (size(M.c,1)-1,size(M.c,1)-1)
 subblockbandwidths(M::ClenshawKron) = (size(M.c,2)-1,size(M.c,2)-1)
 
+struct ClenshawKronLayout <: AbstractLazyBandedBlockBandedLayout end
+MemoryLayout(::Type{<:ClenshawKron}) = ClenshawKronLayout()
+
 function square_getindex(M::ClenshawKron, N::Block{1})
     # Consider P(x) = L^1_x \ 𝐞_0 
     # So that if a(x,y) = P(x)*c*Q(y)' then we have
@@ -47,5 +50,7 @@ getindex(M::ClenshawKron, K::Block{1}, J::Block{1}) = square_getindex(M, max(K,
 getindex(M::ClenshawKron, Kk::BlockIndex{1}, Jj::BlockIndex{1}) = M[block(Kk), block(Jj)][blockindex(Kk), blockindex(Jj)]
 getindex(M::ClenshawKron, k::Int, j::Int) = M[findblockindex(axes(M,1),k), findblockindex(axes(M,2),j)]
 
-Base.array_summary(io::IO, C::ClenshawKron{T}, inds::Tuple{Vararg{OneToInf{Int}}}) where T =
+Base.array_summary(io::IO, C::ClenshawKron{T}, inds) where T =
     print(io, Base.dims2string(length.(inds)), " ClenshawKron{$T} with $(size(C.c)) polynomial")
+
+Base.summary(io::IO, C::ClenshawKron) = Base.array_summary(io, C, axes(C))

src/rect.jl

@@ -18,6 +18,10 @@ end
 
 ==(A::KronPolynomial, B::KronPolynomial) = length(A.args) == length(B.args) && all(map(==, A.args, B.args))
 
+
+struct KronOPLayout{d} <: AbstractMultivariateOPLayout{d} end
+MemoryLayout(::Type{<:KronPolynomial{d}}) where d = KronOPLayout{d}()
+
 const RectPolynomial{T, PP} = KronPolynomial{2, T, PP}
 
 
@@ -71,7 +75,7 @@ end
 function \(P::RectPolynomial, Q::RectPolynomial)
     PA,PB = P.args
     QA,QB = Q.args
-    KronTrav(PA\QA, PB\QB)
+    krontrav(PA\QA, PB\QB)
 end
 
 @simplify function *(Ac::QuasiAdjoint{<:Any,<:RectPolynomial}, B::RectPolynomial)
@@ -144,14 +148,14 @@ pad(C::DiagTrav, ::BlockedOneTo{Int,RangeCumsum{Int,OneToInf{Int}}}) = DiagTrav(
 
 QuasiArrays.mul(A::BivariateOrthogonalPolynomial, b::DiagTrav) = ApplyQuasiArray(*, A, b)
 
-function Base.unsafe_getindex(f::Mul{MultivariateOPLayout{2},<:DiagTravLayout{<:PaddedLayout}}, 𝐱::SVector)
+function Base.unsafe_getindex(f::Mul{KronOPLayout{2},<:DiagTravLayout{<:PaddedLayout}}, 𝐱::SVector)
     P,c = f.A, f.B
     A,B = P.args
     x,y = 𝐱
     clenshaw(vec(clenshaw(paddeddata(c.array), recurrencecoefficients(A)..., x; dims=1)), recurrencecoefficients(B)..., y)
 end
 
-Base.@propagate_inbounds function getindex(f::Mul{MultivariateOPLayout{2},<:DiagTravLayout{<:PaddedLayout}}, x::SVector, j...)
+Base.@propagate_inbounds function getindex(f::Mul{KronOPLayout{2},<:DiagTravLayout{<:PaddedLayout}}, x::SVector, j...)
     @inbounds checkbounds(ApplyQuasiArray(*,f.A,f.B), x, j...)
     Base.unsafe_getindex(f, x, j...)
 end
@@ -171,3 +175,16 @@ function Base._sum(P::RectPolynomial, dims)
     @assert dims == 1
     KronTrav(sum.(P.args; dims=1)...)
 end
+
+## multiplication
+
+function layout_broadcasted(::Tuple{ExpansionLayout{KronOPLayout{2}},KronOPLayout{2}}, ::typeof(*), a, P)
+    axes(a,1) == axes(P,1) || throw(DimensionMismatch())
+
+    A,B = basis(a).args
+    T,U = P.args
+
+    C = paddeddata(invdiagtrav(coefficients(a)))
+        
+    P * ClenshawKron(C, (recurrencecoefficients(A), recurrencecoefficients(B)), (jacobimatrix(T), jacobimatrix(U)))
+end

test/test_rect.jl

@@ -169,19 +169,33 @@ using Base: oneto
 
         a =  (x,y) -> I + x + 2y + 3x^2 +4x*y + 5y^2
         𝐚 = expand(P,splat(a))
-        
-        C = LazyBandedMatrices.paddeddata(LazyBandedMatrices.invdiagtrav(coefficients(𝐚)))
-        
-        A = ClenshawKron(C, (recurrencecoefficients(T), recurrencecoefficients(U)), (jacobimatrix(T), jacobimatrix(U)))
 
+        @testset "ClenshawKron" begin
+            C = LazyBandedMatrices.paddeddata(LazyBandedMatrices.invdiagtrav(coefficients(𝐚)))
 
-        Ã = a(X,Y)
-        for (k,j) in ((Block.(oneto(5)),Block.(oneto(5))), Block.(oneto(5)),Block.(oneto(6)), (Block(2), Block(3)), (4,5),
-                    (Block(2)[2], Block(3)[3]), (Block(2)[2], Block(3)))
-            @test A[k,j] ≈ Ã[k,j]
+            A = ClenshawKron(C, (recurrencecoefficients(T), recurrencecoefficients(U)), (jacobimatrix(T), jacobimatrix(U)))
+
+            @test copy(A) ≡ A
+            @test size(A) == size(X)
+            @test summary(A) == "ℵ₀×ℵ₀ ClenshawKron{Float64} with (3, 3) polynomial"
+
+            Ã = a(X,Y)
+            for (k,j) in ((Block.(oneto(5)),Block.(oneto(5))), Block.(oneto(5)),Block.(oneto(6)), (Block(2), Block(3)), (4,5),
+                        (Block(2)[2], Block(3)[3]), (Block(2)[2], Block(3)))
+                @test A[k,j] ≈ Ã[k,j]
+            end
+
+            @test A[Block(1,2)] ≈ Ã[Block(1,2)]
+            @test A[Block(1,2)][1,2] ≈ Ã[Block(1,2)[1,2]]
         end
 
-        @test A[Block(1,2)] ≈ Ã[Block(1,2)]
-        @test A[Block(1,2)][1,2] ≈ Ã[Block(1,2)[1,2]]
+        @test P \ (𝐚 .* P) isa ClenshawKron
+
+        @test (𝐚 .* 𝐚)[SVector(0.1,0.2)] ≈ 𝐚[SVector(0.1,0.2)]^2
+
+        𝐛 = expand(RectPolynomial(Legendre(),Ultraspherical(3/2)),splat((x,y) -> cos(x*sin(y))))
+        @test (𝐛 .* 𝐚)[SVector(0.1,0.2)] ≈ 𝐚[SVector(0.1,0.2)]𝐛[SVector(0.1,0.2)]
+
+        𝐜 = expand(RectPolynomial(Legendre(),Jacobi(1,0)),splat((x,y) -> cos(x*sin(y))))
     end
 end