Support vec with RectPolynomial and fix sum ordering (#208)

* Support vec with RectPolynomial and fix sum ordering

* reshape and vec

* Update test_rect.jl

Author: dlfivefifty

Project.toml

@@ -30,7 +30,7 @@ BandedMatrices = "1"
 BlockArrays = "1.0"
 BlockBandedMatrices = "0.13"
 ClassicalOrthogonalPolynomials = "0.15"
-ContinuumArrays = "0.19"
+ContinuumArrays = "0.19.5"
 DomainSets = "0.7"
 FastTransforms = "0.17"
 FillArrays = "1.0"
@@ -39,7 +39,7 @@ InfiniteArrays = "0.15"
 InfiniteLinearAlgebra = "0.9, 0.10"
 LazyArrays = "2.3.1"
 LazyBandedMatrices = "0.11.3"
-QuasiArrays = "0.12"
+QuasiArrays = "0.12.2"
 RecurrenceRelationships = "0.2"
 SpecialFunctions = "1, 2"
 StaticArrays = "1"

src/MultivariateOrthogonalPolynomials.jl

@@ -10,8 +10,8 @@ import Base: axes, in, ==, +, -, /, *, ^, \, copy, copyto!, OneTo, getindex, siz
 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, layout_broadcasted, laplacian, abslaplacian, laplacian_axis, abslaplacian_axis, diff_layout, operatororder, broadcastbasis
+import QuasiArrays: LazyQuasiMatrix, LazyQuasiArrayStyle, domain, vec_layout, reshape_layout
+import ContinuumArrays: @simplify, Weight, weight, grid, plotgrid, TransformFactorization, ExpansionLayout, plotvalues, unweighted, plan_transform, checkpoints, transform_ldiv, AbstractBasisLayout, basis_axes, Inclusion, grammatrix, weaklaplacian, layout_broadcasted, laplacian, abslaplacian, laplacian_axis, abslaplacian_axis, diff_layout, operatororder, broadcastbasis, KronExpansionLayout
 
 import ArrayLayouts: MemoryLayout, sublayout, sub_materialize
 import BlockArrays: block, blockindex, BlockSlice, viewblock, blockcolsupport, AbstractBlockStyle, BlockStyle
@@ -21,7 +21,7 @@ import LazyArrays: arguments, paddeddata, LazyArrayStyle, LazyLayout, PaddedLayo
 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 ClassicalOrthogonalPolynomials: jacobimatrix, Weighted, orthogonalityweight, HalfWeighted, WeightedBasis, pad, recurrencecoefficients, clenshaw, weightedgrammatrix, Clenshaw, OPLayout
 import HarmonicOrthogonalPolynomials: BivariateOrthogonalPolynomial, MultivariateOrthogonalPolynomial, Plan,
                                           AngularMomentum, angularmomentum, BlockOneTo, BlockRange1, interlace,
                                           MultivariateOPLayout, AbstractMultivariateOPLayout, MAX_PLOT_BLOCKS

src/rect.jl

@@ -27,6 +27,8 @@ const RectPolynomial{T, PP} = KronPolynomial{2, T, PP}
 
 
 axes(P::KronPolynomial) = (Inclusion(×(map(domain, axes.(P.args, 1))...)), _krontrav_axes(axes.(P.args, 2)...))
+
+
 function getindex(P::RectPolynomial{T}, xy::StaticVector{2}, Jj::BlockIndex{1})::T where T
     a,b = P.args
     J,j = Int(block(Jj)),blockindex(Jj)
@@ -166,7 +168,7 @@ end
 
 function Base._sum(P::RectPolynomial, dims::Int)
     @assert dims == 1
-    KronTrav(sum.(P.args; dims=1)...)
+    KronTrav(reverse(sum.(P.args; dims=1))...)
 end
 
 ## multiplication
@@ -183,4 +185,21 @@ function layout_broadcasted(::Tuple{ExpansionLayout{KronOPLayout{2}},KronOPLayou
 end
 
 
-broadcastbasis(::typeof(+), A::KronPolynomial, B::KronPolynomial) =  KronPolynomial(broadcastbasis.(+, A.args, B.args)...)
+broadcastbasis(::typeof(+), A::KronPolynomial, B::KronPolynomial) =  KronPolynomial(broadcastbasis.(+, A.args, B.args)...)
+
+
+
+####
+# reshape/vec
+####
+
+function reshape_layout(::ExpansionLayout{KronOPLayout{2}}, f)
+    A,B = basis(f).args
+    A*invdiagtrav(coefficients(f))*B'
+end
+
+# TODO: generalise
+function vec_layout(::KronExpansionLayout{OPLayout, OPLayout}, f)
+    A,C,Bc = arguments(f)
+    RectPolynomial(A, Bc') * DiagTrav(C)
+end

test/test_rect.jl

@@ -1,7 +1,7 @@
-using MultivariateOrthogonalPolynomials, ClassicalOrthogonalPolynomials, StaticArrays, LinearAlgebra, BlockArrays, FillArrays, Base64, LazyBandedMatrices, Test
+using MultivariateOrthogonalPolynomials, ClassicalOrthogonalPolynomials, StaticArrays, LinearAlgebra, BlockArrays, FillArrays, Base64, LazyBandedMatrices, ArrayLayouts, Test
 using ClassicalOrthogonalPolynomials: expand, coefficients, recurrencecoefficients
 using MultivariateOrthogonalPolynomials: weaklaplacian, ClenshawKron
-using ContinuumArrays: plotgridvalues
+using ContinuumArrays: plotgridvalues, ExpansionLayout
 using Base: oneto
 
 @testset "RectPolynomial" begin
@@ -148,11 +148,12 @@ using Base: oneto
     end
 
     @testset "sum" begin
-        P = RectPolynomial(Legendre(),Legendre())
-        p₀ = expand(P, 𝐱 -> 1)
-        @test sum(p₀) ≈ 4.0
-        f = expand(P, splat((x,y) -> exp(cos(x^2*y))))
-        @test sum(f) ≈ 10.546408460894801 # empirical
+        for P in (RectPolynomial(Legendre(),Legendre()), RectPolynomial(Legendre(),Chebyshev()))
+            p₀ = expand(P, 𝐱 -> 1)
+            @test sum(p₀) ≈ 4.0
+            f = expand(P, splat((x,y) -> exp(cos(x^2*y))))
+            @test sum(f) ≈ 10.546408460894801 # empirical
+        end
     end
 
     @testset "diff" begin
@@ -219,4 +220,32 @@ using Base: oneto
 
         𝐜 = expand(RectPolynomial(Legendre(),Jacobi(1,0)),splat((x,y) -> cos(x*sin(y))))
     end
+
+    @testset "reshape/vec" begin
+        P = RectPolynomial(Legendre(),Chebyshev())
+        f = expand(P, splat((x,y) -> cos((x-0.1)*exp(y))))
+        F = reshape(f)
+        @test F[0.1,0.2] ≈ f[SVector(0.1,0.2)]
+        @test vec(F)[SVector(0.1,0.2)] ≈ f[SVector(0.1,0.2)]
+
+        g = F[:,0.2]
+        h = F[0.1,:]
+        @test MemoryLayout(g) isa ExpansionLayout
+        @test MemoryLayout(h) isa ExpansionLayout
+        @test g[0.1] ≈ f[SVector(0.1,0.2)]
+        @test h[0.2] ≈ f[SVector(0.1,0.2)]
+
+        @test sum(F; dims=1)[1,0.2] ≈ exp(-0.2)*(sin(0.9exp(0.2)) + sin(1.1exp(0.2)))
+        # TODO: should be matrix but isn't because of InfiniteArrays/src/reshapedarray.jl:77
+        @test_broken sum(F; dims=2)[0.1,1] ≈ 2
+        @test sum(F; dims=2)[0.1] ≈ 2
+    end
+
+    @testset "sample" begin
+        P = RectPolynomial(Legendre(),Legendre())
+        f = expand(P, splat((x,y) -> exp(x*cos(y-0.1))))
+        F = reshape(f)
+        @test sum(F; dims=1)[1,0.2] ≈ 2.346737615950585
+        @test sum(F; dims=2)[0.1,1] ≈ 2.1748993079723618
+    end
 end