reshape and vec

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"

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, vec_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
+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

@@ -190,10 +190,16 @@ broadcastbasis(::typeof(+), A::KronPolynomial, B::KronPolynomial) =  KronPolynom
 
 
 ####
-# vec
+# reshape/vec
 ####
 
-function vec_layout(::ExpansionLayout{KronOPLayout{2}}, f)
+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 ClassicalOrthogonalPolynomials: expand, coefficients, recurrencecoefficients
 using MultivariateOrthogonalPolynomials: weaklaplacian, ClenshawKron
-using ContinuumArrays: plotgridvalues
+using ContinuumArrays: plotgridvalues, ExpansionLayout
 using Base: oneto
 
 @testset "RectPolynomial" begin
@@ -206,9 +206,31 @@ using Base: oneto
         𝐜 = expand(RectPolynomial(Legendre(),Jacobi(1,0)),splat((x,y) -> cos(x*sin(y))))
     end
 
-    @testset "vec" begin
-        P = RectPolynomial(Legendre(),Chebyshev())        
+    @testset "reshape/vec" begin
+        P = RectPolynomial(Legendre(),Chebyshev())
         f = expand(P, splat((x,y) -> cos((x-0.1)*exp(y))))
-        @test vec(f)[0.1,0.2] ≈ f[SVector(0.1,0.2)]
+        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