Merge pull request #90 from jusack/2drfft-for-nd

2D rfft for N-d arrays

Author: andreasnoack

src/plan.jl

@@ -65,9 +65,6 @@ function AbstractFFTs.plan_rfft(x::AbstractArray{T,N}, region; kwargs...)::FFTAP
         pinv = FFTAInvPlan{Complex{T},FFTN}()
         return FFTAPlan_re{Complex{T},FFTN}(tuple(g), region, FFT_FORWARD, pinv, size(x,region[]))
     elseif FFTN == 2
-        if N !== 2
-            throw(ArgumentError("2D real FFT only supported for 2D arrays"))
-        end
         sort!(region)
         g1 = CallGraph{Complex{T}}(size(x,region[1]))
         g2 = CallGraph{Complex{T}}(size(x,region[2]))
@@ -85,9 +82,6 @@ function AbstractFFTs.plan_brfft(x::AbstractArray{T,N}, len, region; kwargs...):
         pinv = FFTAInvPlan{T,FFTN}()
         return FFTAPlan_re{T,FFTN}((g,), region, FFT_BACKWARD, pinv, len)
     elseif FFTN == 2
-        if N !== 2
-            throw(ArgumentError("2D real FFT only supported for 2D arrays"))
-        end
         sort!(region)
         g1 = CallGraph{T}(len)
         g2 = CallGraph{T}(size(x,region[2]))
@@ -120,8 +114,8 @@ function LinearAlgebra.mul!(y::AbstractArray{U,N}, p::FFTAPlan_cx{T,1}, x::Abstr
     if size(p, 1) != size(x, p.region[])
         throw(DimensionMismatch("plan has size $(size(p, 1)), but input array has size $(size(x, p.region[])) along region $(p.region[])"))
     end
-    Rpre = CartesianIndices(size(x)[1:p.region-1])
-    Rpost = CartesianIndices(size(x)[p.region+1:end])
+    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)
@@ -262,37 +256,43 @@ function Base.:*(p::FFTAPlan_re{T,1}, x::AbstractArray{T,N}) where {T<:Complex,
     throw(ArgumentError("only FFT_BACKWARD supported for complex arrays"))
 end
 
-#### 2D plan 2D array
+#### 2D plan ND array
 ##### Forward
-function Base.:*(p::FFTAPlan_re{Complex{T},2}, x::AbstractArray{T,2}) where {T<:Real}
+function Base.:*(p::FFTAPlan_re{Complex{T},2}, x::AbstractArray{T,N}) where {T<:Real, N}
     Base.require_one_based_indexing(x)
     if p.dir === FFT_FORWARD
         half_1 = 1:(p.flen ÷ 2 + 1)
         x_c = similar(x, Complex{T})
         copy!(x_c, x)
         y = similar(x_c)
         LinearAlgebra.mul!(y, complex(p), x_c)
-        return y[half_1, :]
+        return copy(selectdim(y, p.region[1], half_1))
     end
     throw(ArgumentError("only FFT_FORWARD supported for real arrays"))
 end
 
 ##### Backward
-function Base.:*(p::FFTAPlan_re{T,2}, x::AbstractArray{T,2}) where {T<:Complex}
+function Base.:*(p::FFTAPlan_re{T,2}, x::AbstractArray{T,N}) where {T<:Complex, N}
     Base.require_one_based_indexing(x)
-    if size(p, 1) ÷ 2 + 1 != size(x, 1)
-        throw(DimensionMismatch("real 2D plan has size $(size(p)). First dimension of input array should have size ($(size(p, 1) ÷ 2 + 1)), but has size $(size(x, 1))"))
+    if size(p, 1) ÷ 2 + 1 != size(x, p.region[1])
+        throw(DimensionMismatch("real 2D plan has size $(size(p)). First transform dimension of input array should have size ($(size(p, 1) ÷ 2 + 1)), but has size $(size(x, p.region[1]))"))
     end
     if p.dir === FFT_BACKWARD
+        res_size = ntuple(i->ifelse(i==p.region[1], p.flen, size(x,i)), ndims(x))
         # for the inverse transformation we have to reconstruct the full array
-        m, n = size(x)
         half_1 = 1:(p.flen ÷ 2 + 1)
         half_2 = half_1[end]+1:p.flen
-        x_full = similar(x, p.flen, n)
-        x_full[1:m, :] = x
-        start_reverse = m - iseven(p.flen)
-        map!(conj, view(x_full, (m + 1):p.flen, 1), view(x, start_reverse:-1:2, 1))
-        map!(conj, view(x_full, half_2, 2:n), view(x, start_reverse:-1:2, n:-1:2))
+        x_full = similar(x, res_size)
+        # use first half as is
+        copy!(selectdim(x_full, p.region[1], half_1), x)
+
+        # the second half in the first transform dimension is reversed and conjugated
+        x_half_2 = selectdim(x_full, p.region[1], half_2) # view to the second half of x
+        start_reverse = size(x, p.region[1]) - iseven(p.flen)
+        
+        map!(conj, x_half_2, selectdim(x, p.region[1], start_reverse:-1:2))
+        # for the 2D transform we have to reverse index 2:end of the same block in the second transform dimension as well
+        reverse!(selectdim(x_half_2, p.region[2], 2:size(x, p.region[2])), dims=p.region[2])
 
         y = similar(x_full)
         LinearAlgebra.mul!(y, complex(p), x_full)

test/argument_checking.jl

@@ -82,10 +82,3 @@ end
         end
     end
 end
-
-@testset "2D real FFT only supported for 2D arrays" begin
-    xr = zeros(2, 2, 2)
-    xc = complex(xr)
-    @test_throws ArgumentError("2D real FFT only supported for 2D arrays") plan_rfft(xr, 2:3)
-    @test_throws ArgumentError("2D real FFT only supported for 2D arrays") plan_brfft(xc, 2, 2:3)
-end

test/onedim/real_backward.jl

@@ -29,6 +29,10 @@ end
 @testset "1D plan, ND array. Size: $n" for n in 1:64
     x = randn(n, n + 1, n + 2)
 
+    @testset "round tripping with irfft, r=$r" for r in 1:3
+        @test irfft(rfft(x, r), size(x,r), r) ≈ x
+    end
+
     @testset "against 1D array with mapslices, r=$r" for r in 1:3
         y = rfft(x, r)
         @test brfft(y, size(x, r), r) == mapslices(t -> brfft(t, size(x, r)), y; dims = r)

test/twodim/real_backward.jl

@@ -27,16 +27,19 @@ end
 @testset "2D plan, ND array. Size: $n" for n in 1:64
     x = randn(n, n + 1, n + 2)
 
-    @testset "against 1D array with mapslices, r=$r" for r in [[1,2], [1,3], [2,3]]
-        # y = rfft(x, r)
-        y = fft(x, r) # to produce y while tests are broken
-        @test_broken brfft(y, size(x, r), r) == mapslices(t -> brfft(t, size(x, r)), y; dims = r)
+    @testset "round trip with irfft, r=$r" for r in [[1,2], [1,3], [2,3]]
+        @test x ≈ irfft(rfft(x,r), size(x,r[1]), r)
+    end
+
+    @testset "against 2D array with mapslices, r=$r" for r in [[1,2], [1,3], [2,3]]
+        y = rfft(x, r)
+        @test brfft(y, size(x, r[1]), r) == mapslices(t -> brfft(t, size(x, r[1])), y; dims = r)
     end
 end
 
 @testset "allocations" begin
     X = randn(256, 256)
     Y = rfft(X)
     brfft(Y, 256) # compile
-    @test (@test_allocations brfft(Y, 256)) <= 68
+    @test (@test_allocations brfft(Y, 256)) <= 82
 end

test/twodim/real_forward.jl

@@ -7,9 +7,9 @@ using FFTA, Test
     y_ref[1] = length(x)
     @test y ≈ y_ref
     x = randn(N,N)
-    @test_broken rfft(x) ≈ rfft(reshape(x,1,N,N), [2,3])[1,:,:]
-    @test_broken rfft(x) ≈ rfft(reshape(x,1,N,N,1), [2,3])[1,:,:,1]
-    @test_broken rfft(x) ≈ rfft(reshape(x,1,1,N,N,1), [3,4])[1,1,:,:,1]
+    @test rfft(x) ≈ rfft(reshape(x,1,N,N), [2,3])[1,:,:]
+    @test rfft(x) ≈ rfft(reshape(x,1,N,N,1), [2,3])[1,:,:,1]
+    @test rfft(x) ≈ rfft(reshape(x,1,1,N,N,1), [3,4])[1,1,:,:,1]
     @test size(rfft(x)) == (N÷2+1, N)
 end
 
@@ -31,12 +31,12 @@ end
     x = randn(n, n + 1, n + 2)
 
     @testset "against 1D array with mapslices, r=$r" for r in [[1,2], [1,3], [2,3]]
-        @test_broken rfft(x, r) == mapslices(rfft, x; dims = r)
+        @test rfft(x, r) == mapslices(rfft, x; dims = r)
     end
 end
 
 @testset "allocations" begin
     X = randn(256, 256)
     rfft(X) # compile
-    @test (@test_allocations rfft(X)) <= 62
+    @test (@test_allocations rfft(X)) <= 63
 end