Works for vectors

Author: dannys4

src/plan.jl

@@ -1,8 +1,11 @@
+import Base: *
+import LinearAlgebra: mul!
+
 struct FFTAInvPlan{T} <: Plan{T} end
 
 @computed struct FFTAPlan{T<:Union{Real, Complex},N} <: Plan{T}
     callgraph::NTuple{N, CallGraph{(T<:Real) ? Complex{T} : T}}
-    region::NTuple{N, Int}
+    region
     dir::Direction
     pinv::FFTAInvPlan{T}
 end
@@ -12,12 +15,12 @@ function AbstractFFTs.plan_fft(x::AbstractArray{T}, region; kwargs...)::FFTAPlan
     @assert N <= 2 "Only supports vectors and matrices"
     if N == 1
         g = CallGraph{T}(size(x,region[]))
-        pinv = FFTAInvPlan()
+        pinv = FFTAInvPlan{T}()
         return FFTAPlan{T,N}((g,), region, FFT_FORWARD, pinv)
     else
         g1 = CallGraph{T}(size(x,region[1]))
         g2 = CallGraph{T}(size(x,region[2]))
-        pinv = FFTAInvPlan()
+        pinv = FFTAInvPlan{T}()
         return FFTAPlan{T,N}((g1,g2), region, FFT_FORWARD, pinv)
     end
 end
@@ -26,11 +29,21 @@ function AbstractFFTs.plan_bfft(p::FFTAPlan{T,N}) where {T,N}
     return FFTAPlan{T,N}(p.callgraph, p.region, -p.dir, p.pinv)
 end
 
-function LinearAlgebra.mul!(y, p::FFTAPlan, x)
-    fft!(y, x, 1, 1, p.dir, p.callgraph[1].type, p.callgraph, 1)
+function LinearAlgebra.mul!(y::AbstractVector{T}, p::FFTAPlan{T,1}, x::AbstractVector{T}) where T
+    fft!(y, x, 1, 1, p.dir, p.callgraph[1][1].type, p.callgraph[1], 1)
+end
+
+function LinearAlgebra.mul!(y::AbstractArray{T,N}, p::FFTAPlan{T,1}, x::AbstractArray{T,N}) where {T,N}
+    Rpre = CartesianIndices(size(x)[1:p.region-1])
+    Rpost = CartesianIndices(size(x)[p.region+1:end])
+    for Ipre in Rpre
+        for Ipost in Rpost
+            @views fft!(y[Ipre,:,Ipost], x[Ipre,:,Ipost], 1, 1, p.dir, p.callgraph[1][1].type, p.callgraph[1], 1)
+        end
+    end
 end
 
-function *(p::FFTAPlan, x)
+function *(p::FFTAPlan{T,1}, x::AbstractVector{T}) where {T<:Union{Real, Complex}}
     y = similar(x)
     LinearAlgebra.mul!(y, p, x)
     y