Extend FTPlans to Arrays and higher dimensions

src/FastTransforms.jl

@@ -128,6 +128,7 @@ for f in (:jac2jac,
     @eval $f(x::AbstractArray, y...; z...) = $lib_f(x, y...; z...)
 end
 
+include("arrays.jl")
 # following use Toeplitz-Hankel to avoid expensive plans
 # for f in (:leg2cheb, :cheb2leg, :ultra2ultra)
 #     th_f = Symbol("th_", f)

src/arrays.jl

@@ -0,0 +1,99 @@
+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
+    szs = size(c)
+    @assert F.n == szs[dims[1]]
+    ArrayPlan(F, size(c), dims)
+end
+
+function inv_perm(d::Vector{<:Int})
+    inv_d = Vector{Int}(undef, length(d))
+    for (i, val) in enumerate(d)
+        inv_d[val] = i
+    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
+    @assert length(dims) == 1
+    @assert szs == size(f)
+    d = first(dims)
+
+    perm = [d; setdiff(1:ndims(f), d)]
+    fp = permutedims(f, perm)
+
+    fr = reshape(fp, size(fp,1), :)
+
+    permutedims(reshape(F*fr, size(fp)...), inv_perm(perm))
+end
+
+function \(P::ArrayPlan, f::AbstractArray)
+    F, dims, szs = P.F, P.dims, P.szs
+    @assert length(dims) == 1
+    @assert szs == size(f)
+    d = first(dims)
+
+    perm = [d; setdiff(1:ndims(f), d)]
+    fp = permutedims(f, perm)
+
+    fr = reshape(fp, size(fp,1), :)
+
+    permutedims(reshape(F\fr, size(fp)...), inv_perm(perm))
+end
+
+struct NDimsPlan{T, FF<:ArrayPlan{<:T}, Szs<:Tuple, Dims<:Tuple} <: Plan{T}
+    F::FF
+    szs::Szs
+    dims::Dims
+    function NDimsPlan(F, szs, dims)
+        if length(Set(szs[[dims...]])) > 1
+            error("Different size in dims axes not yet implemented in N-dimensional transform.")
+        end
+        new{eltype(F), typeof(F), typeof(szs), typeof(dims)}(F, szs, dims)
+    end
+end
+
+size(P::NDimsPlan) = P.szs
+size(P::NDimsPlan, k::Int) = P.szs[k]
+size(P::NDimsPlan, k...) = P.szs[[k...]]
+
+function NDimsPlan(F::FTPlan, szs::Tuple, dims::Tuple)
+    NDimsPlan(ArrayPlan(F, szs, (first(dims),)), szs, dims)
+end
+
+function *(P::NDimsPlan, f::AbstractArray)
+    F, dims = P.F, P.dims
+    @assert size(P) == size(f)
+    g = copy(f)
+    t = 1:ndims(g)
+    d1 = dims[1]
+    for d in dims
+        perm = ntuple(k -> k == d1 ? t[d] : k == d ? t[d1] : t[k], ndims(g))
+        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(P) == size(f)
+    g = copy(f)
+    t = 1:ndims(g)
+    d1 = dims[1]
+    for d in dims
+        perm = ntuple(k -> k == d1 ? t[d] : k == d ? t[d1] : t[k], ndims(g))
+        gp = permutedims(g, perm)
+        g = permutedims(F\gp, inv_perm(perm))
+    end
+    return g
+end

test/arraystests.jl

@@ -0,0 +1,68 @@
+using FastTransforms, Test
+import FastTransforms: ArrayPlan, NDimsPlan
+
+@testset "Array transform"  begin
+    @testset "ArrayPlan" begin
+        c = randn(5,20,10)
+        F = plan_cheb2leg(c)
+        FT = ArrayPlan(F, c)
+
+        @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
+
+    @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), size(c), (1,2))        
+
+        c = randn(5,20)
+        F = plan_cheb2leg(c)
+        FT = ArrayPlan(F, c)
+        P = NDimsPlan(F, size(c), (1,))
+        @test F*c ≈ FT*c ≈ P*c
+
+        c = randn(20,20,5);
+        F = plan_cheb2leg(c)
+        FT = ArrayPlan(F, c)
+        P = NDimsPlan(FT, size(c), (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(5,10,10,60)
+        F = plan_cheb2leg(randn(10))
+        P = NDimsPlan(F, size(c), (2,3))
+        f = similar(c)
+        for i in axes(f,1), j in axes(f,4)
+            f[i,:,:,j] = (F*(F*c[i,:,:,j])')'
+        end
+        @test f ≈ P*c
+        @test c ≈ P\f
+    end
+end
+
+

test/runtests.jl

@@ -12,3 +12,4 @@ include("clenshawtests.jl")
 include("toeplitzplanstests.jl")
 include("toeplitzhankeltests.jl")
 include("symmetrictoeplitzplushankeltests.jl")
+include("arraystests.jl")