Add tests for multidimensional FFTs and complex input arrays

Author: gaurav-arya

src/TestUtils.jl

@@ -29,16 +29,16 @@ testing: this would most commonly be a constructor such as `Array` or `CuArray`.
 function test_fft_backend(array_constructor; test_real=true, test_inplace=true)
     @testset "fft correctness" begin
         # DFT along last dimension, results computed using FFTW
-        for (_x, _fftw_fft) in (
-            (collect(1:7),
+        for (_x, dims, real_input, _fftw_fft) in (
+            (collect(1:7), 1, true,
              [28.0 + 0.0im,
               -3.5 + 7.267824888003178im,
               -3.5 + 2.7911568610884143im,
               -3.5 + 0.7988521603655248im,
               -3.5 - 0.7988521603655248im,
               -3.5 - 2.7911568610884143im,
               -3.5 - 7.267824888003178im]),
-            (collect(1:8),
+            (collect(1:8), 1, true,
              [36.0 + 0.0im,
               -4.0 + 9.65685424949238im,
               -4.0 + 4.0im,
@@ -47,21 +47,32 @@ function test_fft_backend(array_constructor; test_real=true, test_inplace=true)
               -4.0 - 1.6568542494923806im,
               -4.0 - 4.0im,
               -4.0 - 9.65685424949238im]),
-            (collect(reshape(1:8, 2, 4)),
+            (collect(reshape(1:8, 2, 4)), 2, true,
              [16.0+0.0im  -4.0+4.0im  -4.0+0.0im  -4.0-4.0im;
               20.0+0.0im  -4.0+4.0im  -4.0+0.0im  -4.0-4.0im]),
-            (collect(reshape(1:9, 3, 3)),
+            (collect(reshape(1:9, 3, 3)), 2, true,
              [12.0+0.0im  -4.5+2.598076211353316im  -4.5-2.598076211353316im;
               15.0+0.0im  -4.5+2.598076211353316im  -4.5-2.598076211353316im;
               18.0+0.0im  -4.5+2.598076211353316im  -4.5-2.598076211353316im]),
+            (collect(reshape(1:8, 2, 2, 2)), 1:2, true,
+             [10.0 + 0.0im -4.0 + 0.0im; -2.0 + 0.0im 0.0 + 0.0im;;;
+              26.0 + 0.0im -4.0 + 0.0im; -2.0 + 0.0im 0.0 + 0.0im]),
+            (collect(1:7) + im * collect(8:14), 1, false,
+             [28.0 + 77.0im,
+              -10.76782488800318 + 3.767824888003175im,
+              -6.291156861088416 - 0.7088431389115883im,
+              -4.298852160365525 - 2.7011478396344746im,
+              -2.7011478396344764 - 4.298852160365524im,
+              -0.7088431389115866 - 6.291156861088417im,
+              3.767824888003177 - 10.76782488800318im]),
+            (collect(reshape(1:8, 2, 2, 2)) + im * reshape(9:16, 2, 2, 2), 1:2, false,
+             [10.0 + 42.0im -4.0 - 4.0im; -2.0 - 2.0im 0.0 + 0.0im;;;
+              26.0 + 58.0im -4.0 - 4.0im; -2.0 - 2.0im 0.0 + 0.0im]),
         )
             x = array_constructor(_x) # dummy array that will be passed to plans
-            x_real = float.(x) # for testing real FFTs
-            x_complex = complex.(x_real) # for testing complex FFTs
+            x_complex = complex.(float.(x)) # for testing complex FFTs
             fftw_fft = array_constructor(_fftw_fft)
 
-            dims = ndims(x) # TODO: this is a single dimension, should check multidimensional FFTs too
-
             # FFT
             y = AbstractFFTs.fft(x_complex, dims)
             @test y ≈ fftw_fft
@@ -82,7 +93,7 @@ function test_fft_backend(array_constructor; test_real=true, test_inplace=true)
             end
 
             # BFFT
-            fftw_bfft = size(x_complex, dims) .* x_complex
+            fftw_bfft = prod(size(x_complex, d) for d in dims) .* x_complex
             @test AbstractFFTs.bfft(y, dims) ≈ fftw_bfft
             test_inplace && (@test AbstractFFTs.bfft!(copy(y), dims) ≈ fftw_bfft)
             plans_to_test = [plan_bfft(similar(y), dims)]
@@ -114,16 +125,18 @@ function test_fft_backend(array_constructor; test_real=true, test_inplace=true)
                 @test fftdims(P) == dims
             end
 
-            if test_real
+            if test_real && real_input
+                x_real = float.(x) # for testing real FFTs
                 # RFFT
                 fftw_rfft = fftw_fft[
-                    (Colon() for _ in 1:(ndims(fftw_fft) - 1))...,
-                    1:(size(fftw_fft, ndims(fftw_fft)) ÷ 2 + 1)
+                    (Colon() for _ in 1:(first(dims) - 1))...,
+                    1:(size(fftw_fft, first(dims)) ÷ 2 + 1),
+                    (Colon() for _ in (first(dims) + 1):ndims(fftw_fft))...
                 ]
                 ry = AbstractFFTs.rfft(x_real, dims)
                 @test ry ≈ fftw_rfft
-                for P in [plan_rfft(x_real, dims), inv(plan_irfft(ry, size(x, dims), dims)), 
-                          AbstractFFTs.plan_inv(plan_irfft(ry, size(x, dims), dims))]
+                for P in [plan_rfft(x_real, dims), inv(plan_irfft(ry, size(x, first(dims)), dims)), 
+                          AbstractFFTs.plan_inv(plan_irfft(ry, size(x, first(dims)), dims))]
                     @test eltype(P) <: Real
                     @test P * x_real ≈ fftw_rfft
                     @test mul!(similar(ry), P, x_real) ≈ fftw_rfft
@@ -132,18 +145,18 @@ function test_fft_backend(array_constructor; test_real=true, test_inplace=true)
                 end
 
                 # BRFFT
-                fftw_brfft = complex.(size(x, dims) .* x_real)
-                @test AbstractFFTs.brfft(ry, size(x_real, dims), dims) ≈ fftw_brfft
-                P = plan_brfft(ry, size(x_real, dims), dims)
+                fftw_brfft = prod(size(x_real, d) for d in dims) .* x_real
+                @test AbstractFFTs.brfft(ry, size(x_real, first(dims)), dims) ≈ fftw_brfft
+                P = plan_brfft(ry, size(x_real, first(dims)), dims)
                 @test P * ry ≈ fftw_brfft
                 @test mul!(similar(x_real), P, ry) ≈ fftw_brfft
                 @test P \ (P * ry) ≈ ry
                 @test fftdims(P) == dims
 
                 # IRFFT
                 fftw_irfft = x_complex
-                @test AbstractFFTs.irfft(ry, size(x, dims), dims) ≈ fftw_irfft
-                for P in [plan_irfft(ry, size(x, dims), dims), inv(plan_rfft(x_real, dims)), 
+                @test AbstractFFTs.irfft(ry, size(x, first(dims)), dims) ≈ fftw_irfft
+                for P in [plan_irfft(ry, size(x, first(dims)), dims), inv(plan_rfft(x_real, dims)), 
                           AbstractFFTs.plan_inv(plan_rfft(x_real, dims))]
                     @test P * ry ≈ fftw_irfft
                     @test mul!(similar(x_real), P, ry) ≈ fftw_irfft