* and \ for FTPlan applied to an Array

Author: ioannisPApapadopoulos

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,48 @@
+struct ArrayPlan{T, F<:FTPlan{<:T}, Szs<:Tuple, Dims<:Tuple{<:Int}} <: Plan{T}
+    F::F
+    szs::Szs
+    dims::Dims
+end
+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
+
+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), prod(size(fp)[2:end]))
+
+    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), prod(size(fp)[2:end]))
+
+    permutedims(reshape(F\fr, size(fp)...), inv_perm(perm))
+end

test/arraystests.jl

@@ -0,0 +1,23 @@
+using FastTransforms, Test
+import FastTransforms: ArrayPlan
+
+@testset "Array transform"  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])
+    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/runtests.jl

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