Support 2D Clenshaw on rectangles (#196)

* Support 2D Clenshaw on rectangles

* Update test_rect.jl

* ClenshawKron implemented

* Update Project.toml

* Update Project.toml

* overload broadcast for RectPoly

* v0.8

Author: dlfivefifty

Project.toml

@@ -1,6 +1,6 @@
 name = "MultivariateOrthogonalPolynomials"
 uuid = "4f6956fd-4f93-5457-9149-7bfc4b2ce06d"
-version = "0.7.2"
+version = "0.8"
 
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
@@ -19,6 +19,7 @@ LazyArrays = "5078a376-72f3-5289-bfd5-ec5146d43c02"
 LazyBandedMatrices = "d7e5e226-e90b-4449-9968-0f923699bf6f"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 QuasiArrays = "c4ea9172-b204-11e9-377d-29865faadc5c"
+RecurrenceRelationships = "807425ed-42ea-44d6-a357-6771516d7b2c"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
@@ -28,17 +29,18 @@ ArrayLayouts = "1.0.9"
 BandedMatrices = "1"
 BlockArrays = "1.0"
 BlockBandedMatrices = "0.13"
-ClassicalOrthogonalPolynomials = "0.14"
+ClassicalOrthogonalPolynomials = "0.14.1"
 ContinuumArrays = "0.18"
-DomainSets = "0.6, 0.7"
-FastTransforms = "0.16"
+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"
+LazyBandedMatrices = "0.11.1"
 QuasiArrays = "0.11"
+RecurrenceRelationships = "0.2"
 SpecialFunctions = "1, 2"
 StaticArrays = "1"
 julia = "1.10"

src/MultivariateOrthogonalPolynomials.jl

@@ -4,27 +4,27 @@ using QuasiArrays: AbstractVector
 using ClassicalOrthogonalPolynomials, FastTransforms, BlockBandedMatrices, BlockArrays, DomainSets,
       QuasiArrays, StaticArrays, ContinuumArrays, InfiniteArrays, InfiniteLinearAlgebra,
       LazyArrays, SpecialFunctions, LinearAlgebra, BandedMatrices, LazyBandedMatrices, ArrayLayouts,
-      HarmonicOrthogonalPolynomials
+      HarmonicOrthogonalPolynomials, RecurrenceRelationships
 
 import Base: axes, in, ==, +, -, /, *, ^, \, copy, copyto!, OneTo, getindex, size, oneto, all, resize!, BroadcastStyle, similar, fill!, setindex!, convert, show, summary
 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
+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,
@@ -36,6 +36,7 @@ export MultivariateOrthogonalPolynomial, BivariateOrthogonalPolynomial,
        grammatrix, oneto
 
 include("ModalInterlace.jl")
+include("clenshawkron.jl")
 include("rect.jl")
 include("disk.jl")
 include("rectdisk.jl")

src/clenshawkron.jl

@@ -0,0 +1,56 @@
+
+"""
+    ClenshawKron(coefficientmatrix, (A₁,A₂), (B₁,B₂), (C₁,C₂), (X₁,X₂)))
+
+represents the multiplication operator of a tensor product of two orthogonal polynomials.  That is,
+if `f(x,y) = P(x)*F*Q(y)'` and `a(x,y) = R(x)*A*S(y)'` then
+
+    M = ClenshawKron(A, tuple.(recurrencecoefficients(R), recurrencecoefficients(S)), (jacobimatrix(P), jacobimatrix(Q))
+    M * KronTrav(F)
+
+Gives the coefficients of `a(x,y)*f(x,y)`.
+"""
+struct ClenshawKron{T, Coefs<:AbstractMatrix{T}, Rec<:NTuple{2,NTuple{3,AbstractVector}}, Jac<:NTuple{2,AbstractMatrix}} <: AbstractBandedBlockBandedMatrix{T}
+    c::Coefs
+    recurrencecoeffients::Rec
+    X::Jac
+end
+
+
+copy(M::ClenshawKron) = M
+size(M::ClenshawKron) = (ℵ₀, ℵ₀)
+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
+    # a(X,Y) = 𝐞_0' * inv(L^1_X') * c * Q(Y)'
+    # So we first apply 1D clenshaw to each column
+
+    (A₁,B₁,C₁), (A₂,B₂,C₂) = M.recurrencecoeffients
+    X,Y = M.X
+    m,n = size(M.c)
+    @assert m == n
+    # Apply Clenshaw to each column
+    g = (a,b,N) -> LazyBandedMatrices.krontrav(a[1:N,1:N], b[1:N,1:N])
+    cols = [Clenshaw(M.c[1:m-j+1,j], A₁, B₁, C₁, X) for j=1:n]
+
+    M = m-2+Int(N) # over sample
+    Q_Y = forwardrecurrence(n, A₂, B₂, C₂, Y[1:M,1:M])
+    +(broadcast((a,b,N) -> LazyBandedMatrices.krontrav(a[1:N,1:N], b[1:N,1:N]), Q_Y, cols, Int(N))...)
+end
+
+getindex(M::ClenshawKron, KR::BlockRange{1}, JR::BlockRange{1}) = square_getindex(M, max(maximum(KR), maximum(JR)))[KR,JR]
+getindex(M::ClenshawKron, K::Block{1}, J::Block{1}) = square_getindex(M, max(K, J))[K,J]
+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) 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(clenshaw(paddeddata(c.array), recurrencecoefficients(A)..., x), recurrencecoefficients(B)..., 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

@@ -1,7 +1,8 @@
 using MultivariateOrthogonalPolynomials, ClassicalOrthogonalPolynomials, StaticArrays, LinearAlgebra, BlockArrays, FillArrays, Base64, LazyBandedMatrices, Test
-using ClassicalOrthogonalPolynomials: expand
-using MultivariateOrthogonalPolynomials: weaklaplacian
+using ClassicalOrthogonalPolynomials: expand, coefficients, recurrencecoefficients
+using MultivariateOrthogonalPolynomials: weaklaplacian, ClenshawKron
 using ContinuumArrays: plotgridvalues
+using Base: oneto
 
 @testset "RectPolynomial" begin
     @testset "Evaluation" begin
@@ -154,4 +155,47 @@ using ContinuumArrays: plotgridvalues
         Dₓ = KronTrav(D²,M)
         @test Dₓ[Block.(1:1),Block.(1:1)] == Dₓ[Block(1,1)]
     end
+
+    @testset "variable coefficients" begin
+        T,U = ChebyshevT(), ChebyshevU()
+        P = RectPolynomial(T, U)
+        𝐱 = axes(P,1)
+        x,y = first.(𝐱), last.(𝐱)
+        X = P\(x .* P)
+        Y = P\(y .* P)
+
+        @test X isa KronTrav
+        @test Y isa KronTrav
+
+        a =  (x,y) -> I + x + 2y + 3x^2 +4x*y + 5y^2
+        𝐚 = expand(P,splat(a))
+
+        @testset "ClenshawKron" begin
+            C = LazyBandedMatrices.paddeddata(LazyBandedMatrices.invdiagtrav(coefficients(𝐚)))
+
+            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 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