Improve array construction in tests

Author: gaurav-arya

src/TestUtils.jl

@@ -11,34 +11,33 @@ using LinearAlgebra
 using Test
 
 """
-    test_fft_backend(array_constructor; test_real=true, test_inplace=true) 
+    test_fft_backend(ArrayType=Array; test_real=true, test_inplace=true) 
 
 Run tests to verify correctness of FFT functions using a particular 
 backend plan implementation. The backend implementation is assumed to be loaded 
 prior to calling this function.
 
 # Arguments
 
-- `array_constructor`: determines the `AbstractArray` implementation for
-  which the correctness tests are run. It is assumed to be a callable object that
-  takes in input arrays of type `Array` and return arrays of the desired type for
-  testing. For example, this can be a constructor such as `Array` or `CUDA.CuArray`.
+- `ArrayType`: determines the `AbstractArray` implementation for
+  which the correctness tests are run. Arrays are constructed via
+  `convert(ArrayType, ...)`.
 - `test_real=true`: whether to test real-to-complex and complex-to-real FFTs.
 - `test_inplace=true`: whether to test in-place plans. 
 """
-function test_fft_backend(array_constructor; test_real=true, test_inplace=true)
+function test_fft_backend(ArrayType=Array; test_real=true, test_inplace=true)
     @testset "fft correctness" begin
         # DFT along last dimension, results computed using FFTW
-        for (_x, dims, real_input, _fftw_fft) in (
-            (collect(1:7), 1, true,
+        for (_x, dims, _fftw_fft) in (
+            (collect(1:7), 1,
              [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), 1, true,
+            (collect(1:8), 1,
              [36.0 + 0.0im,
               -4.0 + 9.65685424949238im,
               -4.0 + 4.0im,
@@ -47,49 +46,50 @@ 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)), 2, true,
+            (collect(reshape(1:8, 2, 4)), 2,
              [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)), 2, true,
+            (collect(reshape(1:9, 3, 3)), 2,
              [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,
+            (collect(reshape(1:8, 2, 2, 2)), 1:2,
              cat([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],
                  dims=3)),
-            (collect(1:7) + im * collect(8:14), 1, false,
+            (collect(1:7) + im * collect(8:14), 1,
              [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,
+            (collect(reshape(1:8, 2, 2, 2)) + im * reshape(9:16, 2, 2, 2), 1:2,
              cat([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],
                  dims=3)),
         )
-            x = array_constructor(_x) # dummy array that will be passed to plans
-            x_complex = complex.(float.(x)) # for testing complex FFTs
-            fftw_fft = array_constructor(_fftw_fft)
+            x = convert(ArrayType, _x) # dummy array that will be passed to plans
+            x_complex = convert(ArrayType, complex.(x)) # for testing complex FFTs
+            x_complexfloat = convert(ArrayType, complex.(float.(x))) # for in-place operations 
+            fftw_fft = convert(ArrayType, _fftw_fft)
 
             # FFT
             y = AbstractFFTs.fft(x_complex, dims)
             @test y ≈ fftw_fft
             if test_inplace
-                @test AbstractFFTs.fft!(copy(x_complex), dims) ≈ fftw_fft
+                @test AbstractFFTs.fft!(copy(x_complexfloat), dims) ≈ fftw_fft
             end
             # test plan_fft and also inv and plan_inv of plan_ifft, which should all give 
             # functionally identical plans
             plans_to_test = [plan_fft(x, dims), inv(plan_ifft(x, dims)), 
                              AbstractFFTs.plan_inv(plan_ifft(x, dims))]
             for P in plans_to_test
-                @test mul!(similar(y), P, copy(x_complex)) ≈ fftw_fft
+                @test mul!(similar(y), P, copy(x_complexfloat)) ≈ fftw_fft
             end
             if test_inplace
-                push!(plans_to_test, plan_fft!(similar(x_complex), dims))
+                push!(plans_to_test, plan_fft!(similar(x_complexfloat), dims))
             end
             for P in plans_to_test 
                 @test eltype(P) <: Complex
@@ -105,7 +105,7 @@ function test_fft_backend(array_constructor; test_real=true, test_inplace=true)
                 @test AbstractFFTs.bfft!(copy(y), dims) ≈ fftw_bfft
             end
             P = plan_bfft(similar(y), dims)
-            @test mul!(similar(x_complex), P, copy(y)) ≈ fftw_bfft
+            @test mul!(similar(x_complexfloat), P, copy(y)) ≈ fftw_bfft
             plans_to_test = if test_inplace
                 [P, plan_bfft!(similar(y), dims)]
             else
@@ -127,10 +127,10 @@ function test_fft_backend(array_constructor; test_real=true, test_inplace=true)
             plans_to_test = [plan_ifft(x, dims), inv(plan_fft(x, dims)), 
                              AbstractFFTs.plan_inv(plan_fft(x, dims))]
             for P in plans_to_test
-                @test mul!(similar(x_complex), P, copy(y)) ≈ fftw_ifft
+                @test mul!(similar(x_complexfloat), P, copy(y)) ≈ fftw_ifft
             end
             if test_inplace
-                push!(plans_to_test, plan_ifft!(similar(x_complex), dims))
+                push!(plans_to_test, plan_ifft!(similar(x_complexfloat), dims))
             end
             for P in plans_to_test
                 @test eltype(P) <: Complex
@@ -139,7 +139,7 @@ function test_fft_backend(array_constructor; test_real=true, test_inplace=true)
                 @test fftdims(P) == dims
             end
 
-            if test_real && real_input
+            if test_real && (x isa Real) 
                 x_real = float.(x) # for testing real FFTs
                 # RFFT
                 fftw_rfft = selectdim(fftw_fft, first(dims), 1:(size(fftw_fft, first(dims)) ÷ 2 + 1))