Finish implementations of vectorized 1D & 2D plans (#80)

* implement _inv_perm_blockvec for vectorized 1D plan

* vectorized 2D transform

* add return for CodeCov

* missed 2

* fix axes

---------

Co-authored-by: ioannisPApapadopoulos <john.papadopoulos0@gmail.com>

Author: ioannisPApapadopoulos

src/arrowhead.jl

@@ -206,7 +206,7 @@ function materialize!(M::MatMulVecAdd{<:ArrowheadLayouts,<:AbstractStridedLayout
     for k = 1:d
         mul!(view(y, m+k:d:size(y,1)), A.D[k], view(x, n+k:d:size(x,1)), α, one(α))
     end
-    y_in
+    return y_in
 end
 
 function materialize!(M::MatMulMatAdd{<:ArrowheadLayouts,<:AbstractColumnMajor,<:AbstractColumnMajor})
@@ -229,7 +229,7 @@ function materialize!(M::MatMulMatAdd{<:ArrowheadLayouts,<:AbstractColumnMajor,<
     for k = 1:d
         mul!(view(Y, m+k:d:size(Y,1), :), A.D[k], view(X, n+k:d:size(Y,1), :), α, one(α))
     end
-    Y_in
+    return Y_in
 end
 
 function materialize!(M::MatMulMatAdd{<:AbstractColumnMajor,<:ArrowheadLayouts,<:AbstractColumnMajor})
@@ -252,7 +252,7 @@ function materialize!(M::MatMulMatAdd{<:AbstractColumnMajor,<:ArrowheadLayouts,<
     for k = 1:d
         mul!(view(Y, :, n+k:d:size(Y,2)), view(X, :, m+k:d:size(Y,2)), A.D[k], α, one(α))
     end
-    Y_in
+    return Y_in
 end
 
 
@@ -450,7 +450,7 @@ for (UNIT, Tri) in (('U',UnitUpperTriangular), ('N', UpperTriangular))
 
         ArrayLayouts.ldiv!($Tri(A), b̃_1)
 
-        dest
+        return dest
     end
 end
 for (UNIT, Tri) in (('U',UnitLowerTriangular), ('N', LowerTriangular))
@@ -479,7 +479,7 @@ for (UNIT, Tri) in (('U',UnitLowerTriangular), ('N', LowerTriangular))
             ArrayLayouts.ldiv!($Tri(D[k]), view(dest, n+k:m:length(dest)))
         end
 
-        dest
+        return dest
     end
 end
 

src/continuouspolynomial.jl

@@ -123,7 +123,7 @@ function *(Pl::ContinuousPolynomialTransform{T,<:Any,<:Any,Int}, X::AbstractMatr
         end
     end
     cfs[Block.(2:M-1)] .= vec(dat[3:end,:]')
-    cfs
+    return cfs
 end
 
 function \(Pl::ContinuousPolynomialTransform{T,<:Any,<:Any,Int}, cfs::AbstractVector{T}) where T
@@ -165,7 +165,7 @@ function *(Pl::ContinuousPolynomialTransform{T}, X::AbstractArray{T,4}) where T
     for k = 1:M, j = 1:N, l = 1:O, m = 1:P
         cfs[_contpolyinds2blocks(k,j), _contpolyinds2blocks(l,m)] = dat[k,j,l,m]
     end
-    cfs
+    return cfs
 end
 
 function \(Pl::ContinuousPolynomialTransform{T}, cfs::AbstractMatrix{T}) where T

src/dirichlet.jl

@@ -91,7 +91,7 @@ function *(Pl::DirichletPolynomialTransform{T,<:Any,<:Any,Int}, X::AbstractMatri
         end
     end
     cfs[Block.(2:M-1)] .= vec(dat[3:end,:]')
-    cfs
+    return cfs
 end
 
 function \(Pl::DirichletPolynomialTransform{T,<:Any,<:Any,Int}, cfs::AbstractVector{T}) where T
@@ -137,7 +137,7 @@ function *(Pl::DirichletPolynomialTransform{T}, X::AbstractArray{T,4}) where T
         l == 2 && m == P && continue
         cfs[_dirichletpolyinds2blocks(k,j), _dirichletpolyinds2blocks(l,m)] = dat[k,j,l,m]
     end
-    cfs
+    return cfs
 end
 
 function \(Pl::DirichletPolynomialTransform{T}, cfs::AbstractMatrix{T}) where T

src/piecewisepolynomial.jl

@@ -22,7 +22,7 @@ function repeatgrid(ax, g, pts)
     for j in axes(ret, 2)
         ret[:, j] = affine(ax, pts[j] .. pts[j+1])[g]
     end
-    ret
+    return ret
 end
 
 
@@ -83,12 +83,10 @@ function _perm_blockvec(X::AbstractArray{T,3}, dims=1) where T
     for k = 2:size(X,3)
         ret[:,k] = _perm_blockvec(X[:,:,k])
     end
-    ret
+    return ret
 end
 
 
-
-
 function _perm_blockvec(X::AbstractArray{T,4}, dims=(1,2)) where T
     @assert dims == 1:2 || dims == (1,2)
     X1 = _perm_blockvec(X[:,:,1,1])
@@ -97,18 +95,48 @@ function _perm_blockvec(X::AbstractArray{T,4}, dims=(1,2)) where T
     for k = axes(X,1), j = axes(X,2), l = axes(X,3), m = axes(X,4)
         ret[Block(k)[j], Block(l)[m]] = X[k,j,l,m]
     end
-    ret
+    return ret
 end
 
 function _inv_perm_blockvec(X::AbstractMatrix{T}, dims=(1,2)) where T
-    @assert dims == 1:2 || dims == (1,2)
+    @assert dims == 1:2 || dims == (1,2) || dims == 1
     M,N = blocksize(X)
     m,n = size(X)
-    ret = Array{T}(undef, M, m ÷ M, N, n ÷ N)
-    for k = axes(ret,1), j = axes(ret,2), l = axes(ret,3), m = axes(ret,4)
-        ret[k,j,l,m] = X[Block(k)[j], Block(l)[m]]
+    if dims == 1:2 || dims == (1,2)
+        ret = Array{T}(undef, M, m ÷ M, N, n ÷ N)
+        for k = axes(ret,1), j = axes(ret,2), l = axes(ret,3), m = axes(ret,4)
+            ret[k,j,l,m] = X[Block(k)[j], Block(l)[m]]
+        end
+    elseif dims == 1
+        ret = Array{T}(undef, M, m ÷ M, n ÷ N)
+        for k = axes(ret,1), j = axes(ret,2), l = axes(ret,3)
+            ret[k,j,l] = X[Block(k)[j], Block(1)[l]]
+        end
+    end
+    return ret
+end
+
+function _perm_blockvec(X::AbstractArray{T,5}, dims=(1,2)) where T
+    @assert dims == 1:2 || dims == (1,2)
+    X1 = _perm_blockvec(X[:,:,:,:,1])
+    ret = BlockedArray{T}(undef, (axes(X1,1), axes(X1,2), axes(X,5)))
+    ret[:, :, 1] = X1
+    for k = 2:lastindex(ret,3)
+        ret[:, :, k] = _perm_blockvec(X[:,:,:,:,k])
     end
-    ret
+    return ret
+end
+
+function _inv_perm_blockvec(X::AbstractArray{T,3}, dims=(1,2)) where T
+    @assert dims == 1:2 || dims == (1,2)
+    M,N,L = blocksize(X)
+    m,n,ℓ = size(X)
+
+    ret = Array{T}(undef, M, m ÷ M, N, n ÷ N, ℓ÷L)
+    for k = axes(ret,5)
+        ret[:,:,:,:,k] = _inv_perm_blockvec(X[:,:,k])
+    end
+    return ret
 end
 
 \(F::ApplyPlan{<:Any,typeof(_perm_blockvec)}, X::AbstractArray) = F.F \ _inv_perm_blockvec(X, F.args...)
@@ -137,6 +165,13 @@ function plan_transform(P::PiecewisePolynomial, (M,n)::Tuple{Block{1},Int}, dims
     ApplyPlan(_perm_blockvec, F, (dims,))
 end
 
+function plan_transform(P::PiecewisePolynomial, (N,M,n)::Tuple{Block{1},Block{1},Int}, dims=ntuple(identity,Val(2)))
+    @assert dims == 1:2 || dims == ntuple(identity,Val(2))
+    Ns = (N,M)
+    F = plan_transform(P.basis, (_interlace_const(length(P.points)-1, Int.(Ns)...)..., n), _doubledims(dims...))
+    ApplyPlan(_perm_blockvec, F, (dims,))
+end
+
 function factorize(V::SubQuasiArray{<:Any,2,<:PiecewisePolynomial,<:Tuple{Inclusion,BlockSlice}}, dims...)
     P = parent(V)
     _,JR = parentindices(V)

test/test_piecewisepolynomial.jl

@@ -31,6 +31,7 @@ using LazyBandedMatrices: BlockVec
                 x,F = plan_grid_transform(P, (Block(10),2), 1)
                 KR = Block.(1:10)
                 @test F * [exp.(x) ;;; cos.(x)] ≈ [transform(P,exp)[KR] transform(P,cos)[KR]]
+                @test F \ (F * [exp.(x) ;;; cos.(x)]) ≈ [exp.(x) ;;; cos.(x)]
 
                 t = axes(P,1)
                 @test (P \ [cos.(t) sin.(t)])[KR,:] ≈ [(P\cos.(t))[KR,:] (P\sin.(t))[KR,:]]
@@ -64,4 +65,22 @@ using LazyBandedMatrices: BlockVec
             @test pl\(pl*X) ≈ X
         end
     end
+
+    @testset "tensor transform" begin
+        for P in (PiecewisePolynomial(Chebyshev(), range(0,1; length=3)), PiecewisePolynomial(Legendre(), range(0,1; length=3)))
+            xy, pl = plan_grid_transform(P, (Block(10),Block(11), 3), (1,2))
+            @test size(pl) == (10, 2, 11, 2, 3)
+            x,y=first(xy),last(xy)
+            y = reshape(y,1,1,size(y)...)
+
+            X = [exp.(x .+ cos.(y)) ;;;;; exp.(y .+ sin.(x)) ;;;;; cos.(x .+ sin.(y))]
+            C = pl*X
+
+            @test P[0.1, Block.(1:10)]' * C[:,:,1] * P[0.2, Block.(1:11)] ≈ exp(0.1 + cos(0.2))
+            @test P[0.1, Block.(1:10)]' * C[:,:,2] * P[0.2, Block.(1:11)] ≈ exp(0.2 + sin(0.1))
+            @test P[0.1, Block.(1:10)]' * C[:,:,3] * P[0.2, Block.(1:11)] ≈ cos(0.1 + sin(0.2))
+
+            @test pl\C ≈ X
+        end
+    end
 end