Extend FTPlans to Arrays and higher dimensions

src/arrays.jl

@@ -1,8 +1,10 @@
-struct ArrayPlan{T, F<:FTPlan{<:T}, Szs<:Tuple, Dims<:Tuple{<:Int}} <: Plan{T}
-    F::F
+struct ArrayPlan{T, FF<:FTPlan{<:T}, Szs<:Tuple, Dims<:Tuple{<:Int}} <: Plan{T}
+    F::FF
     szs::Szs
     dims::Dims
 end
+size(P::ArrayPlan) = P.szs
+size(P::ArrayPlan, k::Int) = P.szs[k]
 size(P::ArrayPlan, k...) = P.szs[[k...]]
 
 function ArrayPlan(F::FTPlan{<:T}, c::AbstractArray{T}, dims::Tuple{<:Int}=(1,)) where T
@@ -18,6 +20,7 @@ function inv_perm(d::Vector{<:Int})
     end
     return inv_d
 end
+inv_perm(d::Tuple) = inv_perm([d...])
 
 function *(P::ArrayPlan, f::AbstractArray)
     F, dims, szs = P.F, P.dims, P.szs
@@ -45,4 +48,45 @@ function \(P::ArrayPlan, f::AbstractArray)
     fr = reshape(fp, size(fp,1), prod(size(fp)[2:end]))
 
     permutedims(reshape(F\fr, size(fp)...), inv_perm(perm))
+end
+
+struct NDimsPlan{T, FF<:ArrayPlan{<:T}, Dims<:Tuple} <: Plan{T}
+    F::FF
+    dims::Dims
+    function NDimsPlan(F, dims)
+        if length(Set(size(F, dims...))) > 1
+            error("Different size in dims axes not yet implemented in N-dimensional transform.")
+        end
+        new{eltype(F), typeof(F), typeof(dims)}(F, dims)
+    end
+end
+
+size(P::NDimsPlan) = size(P.F)
+size(P::NDimsPlan, k::Int) = size(P.F, k)
+size(P::NDimsPlan, k...) = size(P.F, k...)
+
+function *(P::NDimsPlan, f::AbstractArray)
+    F, dims = P.F, P.dims
+    @assert size(F) == size(f)
+    g = copy(f)
+    t = 1:ndims(g)
+    for d in dims
+        perm = ntuple(k -> k == 1 ? t[d] : k == d ? t[1] : t[k], length(t))
+        gp = permutedims(g, perm)
+        g = permutedims(F*gp, inv_perm(perm))
+    end
+    return g
+end
+
+function \(P::NDimsPlan, f::AbstractArray)
+    F, dims = P.F, P.dims
+    @assert size(F) == size(f)
+    g = copy(f)
+    t = 1:ndims(g)
+    for d in dims
+        perm = ntuple(k -> k == 1 ? t[d] : k == d ? t[1] : t[k], length(t))
+        gp = permutedims(g, perm)
+        g = permutedims(F\gp, inv_perm(perm))
+    end
+    return g
 end

test/arraystests.jl

@@ -1,23 +1,63 @@
 using FastTransforms, Test
-import FastTransforms: ArrayPlan
+import FastTransforms: ArrayPlan, NDimsPlan
 
 @testset "Array transform"  begin
-    c = randn(5,20,10)
-    F = plan_cheb2leg(c)
-    FT = ArrayPlan(F, c)
+    @testset "ArrayPlan" begin
+        c = randn(5,20,10)
+        F = plan_cheb2leg(c)
+        FT = ArrayPlan(F, c)
 
-    f = similar(c);
-    for k in axes(c,3)
-        f[:,:,k] = (F*c[:,:,k])
+        @test size(FT) == size(c)
+        @test size(FT,1) == size(c,1)
+        @test size(FT,1,2) == (size(c,1), size(c,2))
+
+        f = similar(c);
+        for k in axes(c,3)
+            f[:,:,k] = (F*c[:,:,k])
+        end
+        @test f ≈ FT*c
+        @test c ≈ FT\f
+
+        F = plan_cheb2leg(Vector{Float64}(axes(c,2)))
+        FT = ArrayPlan(F, c, (2,))
+        for k in axes(c,3)
+            f[:,:,k] = (F*c[:,:,k]')'
+        end
+        @test f ≈ FT*c
+        @test c ≈ FT\f
     end
-    @test f ≈ FT*c
-    @test c ≈ FT\f
 
-    F = plan_cheb2leg(Vector{Float64}(axes(c,2)))
-    FT = ArrayPlan(F, c, (2,))
-    for k in axes(c,3)
-        f[:,:,k] = (F*c[:,:,k]')'
+    @testset "NDimsPlan" begin
+        c = randn(20,10,20)
+        @test_throws ErrorException("Different size in dims axes not yet implemented in N-dimensional transform.") NDimsPlan(ArrayPlan(plan_cheb2leg(c), c), (1,2))        
+
+        c = randn(20,20,5);
+        F = plan_cheb2leg(c)
+        FT = ArrayPlan(F, c)
+        P = NDimsPlan(FT, (1,2))
+
+        @test size(P) == size(c)
+        @test size(P,1) == size(c,1)
+        @test size(P,1,2) == (size(c,1), size(c,2))
+
+
+        f = similar(c);
+        for k in axes(f,3)
+            f[:,:,k] = (F*(F*c[:,:,k])')'
+        end
+        @test f ≈ P*c
+        @test c ≈ P\f
+
+        c = randn(10,5,10,60)
+        F = plan_cheb2leg(c)
+        P = NDimsPlan(ArrayPlan(F, c), (1,3))
+        f = similar(c)
+        for i in axes(f,2), j in axes(f,4)
+            f[:,i,:,j] = (F*(F*c[:,i,:,j])')'
+        end
+        @test f ≈ P*c
+        @test c ≈ P\f
     end
-    @test f ≈ FT*c
-    @test c ≈ FT\f
-end
+end
+
+