Support vec with RectPolynomial and fix sum ordering

Author: dlfivefifty

Project.toml

@@ -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,7 +10,7 @@ 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 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 ArrayLayouts: MemoryLayout, sublayout, sub_materialize

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)
@@ -164,9 +166,9 @@ end
 
 ## sum
 
-function Base._sum(P::RectPolynomial, dims)
+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,15 @@ 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)...)
+
+
+
+####
+# vec
+####
+
+function vec_layout(::ExpansionLayout{KronOPLayout{2}}, f)
+    A,B = basis(f).args
+    A*invdiagtrav(coefficients(f))*B'
+end

test/test_rect.jl

@@ -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 "KronTrav bug" begin
@@ -204,4 +205,10 @@ using Base: oneto
 
         𝐜 = expand(RectPolynomial(Legendre(),Jacobi(1,0)),splat((x,y) -> cos(x*sin(y))))
     end
+
+    @testset "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)]
+    end
 end