DiagTrav coefficients should be treated as an ExpansionLayout (#158)

Author: dlfivefifty

src/MultivariateOrthogonalPolynomials.jl

@@ -11,13 +11,13 @@ import Base.Broadcast: Broadcasted, broadcasted, DefaultArrayStyle
 import DomainSets: boundary
 
 import QuasiArrays: LazyQuasiMatrix, LazyQuasiArrayStyle, domain
-import ContinuumArrays: @simplify, Weight, weight, grid, plotgrid, TransformFactorization, ExpansionLayout, plotvalues, unweighted, plan_grid_transform, checkpoints, transform_ldiv
+import ContinuumArrays: @simplify, Weight, weight, grid, plotgrid, TransformFactorization, ExpansionLayout, plotvalues, unweighted, plan_grid_transform, checkpoints, transform_ldiv, AbstractBasisLayout, basis_axes
 
 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
+import LazyArrays: arguments, paddeddata, LazyArrayStyle, LazyLayout, PaddedLayout, applylayout
 import LazyBandedMatrices: LazyBandedBlockBandedLayout, AbstractBandedBlockBandedLayout, AbstractLazyBandedBlockBandedLayout, _krontrav_axes, DiagTravLayout
 import InfiniteArrays: InfiniteCardinal, OneToInf
 

src/rect.jl

@@ -77,6 +77,9 @@ function plan_grid_transform(P::KronPolynomial{d,<:Any,<:Fill}, B::Tuple{Block{1
     SVector.(x̃, x̃'), ApplyPlan(DiagTrav, F)
 end
 
+applylayout(::Type{typeof(*)}, ::Lay, ::DiagTravLayout) where Lay <: AbstractBasisLayout = ExpansionLayout{Lay}()
+ContinuumArrays._mul_plotgrid(::Tuple{Any,DiagTravLayout{<:PaddedLayout}}, (P,c)) = plotgrid(P, maximum(blockcolsupport(c)))
+
 pad(C::DiagTrav, ::BlockedUnitRange{RangeCumsum{Int,OneToInf{Int}}}) = DiagTrav(pad(C.array, ∞, ∞))
 
 QuasiArrays.mul(A::BivariateOrthogonalPolynomial, b::DiagTrav) = ApplyQuasiArray(*, A, b)

src/triangle.jl

@@ -27,6 +27,8 @@ show(io::IO, P::JacobiTriangle) = summary(io, P)
 summary(io::IO, P::JacobiTriangle) = print(io, "JacobiTriangle($(P.a), $(P.b), $(P.c))")
 
 
+basis_axes(::Inclusion{<:Any,<:UnitTriangle}, v) = JacobiTriangle()
+
 """
     TriangleWeight(a, b, c)
 
@@ -639,6 +641,13 @@ function tridenormalize!(F̌,a,b,c)
     F̌
 end
 
+function trinormalize!(F̌,a,b,c)
+    for n = 0:size(F̌,1)-1, k = 0:n
+        F̌[n-k+1,k+1] /= _ft_trinorm(n,k,a,b,c)
+    end
+    F̌
+end
+
 function trigrid(N::Integer)
     M = N
     x = [sinpi((2N-2n-1)/(4N))^2 for n in 0:N-1]
@@ -666,4 +675,35 @@ function plan_grid_transform(P::JacobiTriangle, Bs::Tuple{Block{1}}, dims=1:1)
     T = eltype(P)
     N = Bs[1]
     grid(P, N), TriPlan{T}(N, P.a, P.b, P.c)
+end
+
+struct TriIPlan{T}
+    tri2cheb::FastTransforms.FTPlan{T,2,FastTransforms.TRIANGLE}
+    cheb2grid::FastTransforms.FTPlan{T,2,FastTransforms.TRIANGLESYNTHESIS}
+    a::T
+    b::T
+    c::T
+end
+
+TriIPlan{T}(F::AbstractMatrix{T}, a, b, c) where T =
+    TriIPlan{T}(plan_tri2cheb(F, a, b, c), plan_tri_synthesis(F), a, b, c)
+
+TriIPlan{T}(N::Block{1}, a, b, c) where T = TriIPlan{T}(Matrix{T}(undef, Int(N), Int(N)), a, b, c)
+
+*(T::TriIPlan, F::DiagTrav) = T.cheb2grid*(T.tri2cheb*trinormalize!(Matrix(F.array),T.a,T.b,T.c))
+
+
+function plotgrid(S::JacobiTriangle{T}, B::Block{1}) where T
+    N = min(2Int(B), MAX_PLOT_BLOCKS)
+    grid(S, Block(N)) # double sampling
+end
+
+function plotvalues(u::ApplyQuasiVector{T,typeof(*),<:Tuple{JacobiTriangle, AbstractVector}}, x) where T
+    P,c = u.args
+    B = findblock(axes(P,2), last(colsupport(c)))
+
+    N = min(2Int(B), MAX_PLOT_BLOCKS)
+    F = TriIPlan{T}(Block(N), P.a, P.b, P.c)
+    C = F * DiagTrav(c.array[1:N,1:N]) # transform to grid
+    C
 end

test/test_triangle.jl

@@ -1,6 +1,6 @@
 using MultivariateOrthogonalPolynomials, StaticArrays, BlockArrays, BlockBandedMatrices, ArrayLayouts,
         QuasiArrays, Test, ClassicalOrthogonalPolynomials, BandedMatrices, FastTransforms, LinearAlgebra
-import MultivariateOrthogonalPolynomials: tri_forwardrecurrence, grid, TriangleRecurrenceA, TriangleRecurrenceB, TriangleRecurrenceC, xy_muladd!
+import MultivariateOrthogonalPolynomials: tri_forwardrecurrence, grid, TriangleRecurrenceA, TriangleRecurrenceB, TriangleRecurrenceC, xy_muladd!, ExpansionLayout
 
 @testset "Triangle" begin
     @testset "basics" begin
@@ -52,6 +52,7 @@ import MultivariateOrthogonalPolynomials: tri_forwardrecurrence, grid, TriangleR
             𝐱 = SVector(0.1,0.2)
             c = PseudoBlockVector([1; Zeros(∞)], (axes(P,2),))
             f = P*c
+            @test MemoryLayout(f) isa ExpansionLayout
             @test @inferred(f[𝐱]) == 1.0
             c = PseudoBlockVector([1:3; Zeros(∞)], (axes(P,2),))
             f = P*c
@@ -115,13 +116,16 @@ import MultivariateOrthogonalPolynomials: tri_forwardrecurrence, grid, TriangleR
             N = 20
             P_N = P[:,Block.(Base.OneTo(N))]
             u = P_N * (P_N \ (exp.(x) .* cos.(y)))
+            @test MemoryLayout(u) isa ExpansionLayout
             @test u[SVector(0.1,0.2)] ≈ exp(0.1)*cos(0.2)
 
             P_n = P[:,1:200]
             u = P_n * (P_n \ (exp.(x) .* cos.(y)))
+            @test MemoryLayout(u) isa ExpansionLayout
             @test u[SVector(0.1,0.2)] ≈ exp(0.1)*cos(0.2)
 
             @time u = P * (P \ (exp.(x) .* cos.(y)))
+            @test MemoryLayout(u) isa ExpansionLayout
             @test u[SVector(0.1,0.2)] ≈ exp(0.1)*cos(0.2)
         end
     end