Cleanup fft backend tests

Author: gaurav-arya

src/TestUtils.jl

@@ -10,22 +10,12 @@ using AbstractFFTs: Plan
 using LinearAlgebra
 using Test
 
-import Unitful
-
-function test_fft_backend()
-    @testset "rfft sizes" begin
-        A = rand(11, 10)
-        @test @inferred(AbstractFFTs.rfft_output_size(A, 1)) == (6, 10)
-        @test @inferred(AbstractFFTs.rfft_output_size(A, 2)) == (11, 6)
-        A1 = rand(6, 10); A2 = rand(11, 6)
-        @test @inferred(AbstractFFTs.brfft_output_size(A1, 11, 1)) == (11, 10)
-        @test @inferred(AbstractFFTs.brfft_output_size(A2, 10, 2)) == (11, 10)
-        @test_throws AssertionError AbstractFFTs.brfft_output_size(A1, 10, 2)
-    end
-
+"""
+"""
+function test_fft_backend(array_constructor)
     @testset "fft correctness" begin
         # DFT along last dimension, results computed using FFTW
-        for (x, fftw_fft) in (
+        for (_x, _fftw_fft) in (
             (collect(1:7),
              [28.0 + 0.0im,
               -3.5 + 7.267824888003178im,
@@ -51,6 +41,9 @@ function test_fft_backend()
               15.0+0.0im  -4.5+2.598076211353316im  -4.5-2.598076211353316im;
               18.0+0.0im  -4.5+2.598076211353316im  -4.5-2.598076211353316im]),
         )
+            x = array_constructor(_x)
+            fftw_fft = array_constructor(_fftw_fft)
+
             # FFT
             dims = ndims(x)
             y = AbstractFFTs.fft(x, dims)
@@ -59,39 +52,37 @@ function test_fft_backend()
             # functionally identical plans
             for P in [plan_fft(x, dims), inv(plan_ifft(x, dims)), 
                       AbstractFFTs.plan_inv(plan_ifft(x, dims))]
-                @test eltype(P) === ComplexF64
+                @test eltype(P) <: Complex
                 @test P * x ≈ fftw_fft
                 @test P \ (P * x) ≈ x
                 @test fftdims(P) == dims
             end
 
+            # BFFT
             fftw_bfft = complex.(size(x, dims) .* x)
             @test AbstractFFTs.bfft(y, dims) ≈ fftw_bfft
             P = plan_bfft(x, dims)
             @test P * y ≈ fftw_bfft
             @test P \ (P * y) ≈ y
             @test fftdims(P) == dims
 
+            # IFFT
             fftw_ifft = complex.(x)
             @test AbstractFFTs.ifft(y, dims) ≈ fftw_ifft
-            # test plan_ifft and also inv and plan_inv of plan_fft, which should all give 
-            # functionally identical plans
             for P in [plan_ifft(x, dims), inv(plan_fft(x, dims)), 
                       AbstractFFTs.plan_inv(plan_fft(x, dims))]
                 @test P * y ≈ fftw_ifft
                 @test P \ (P * y) ≈ y
                 @test fftdims(P) == dims
             end
 
-            # real FFT
+            # RFFT
             fftw_rfft = fftw_fft[
                 (Colon() for _ in 1:(ndims(fftw_fft) - 1))...,
                 1:(size(fftw_fft, ndims(fftw_fft)) ÷ 2 + 1)
             ]
             ry = AbstractFFTs.rfft(x, dims)
             @test ry ≈ fftw_rfft
-            # test plan_rfft and also inv and plan_inv of plan_irfft, which should all give 
-            # functionally identical plans
             for P in [plan_rfft(x, dims), inv(plan_irfft(ry, size(x, dims), dims)), 
                       AbstractFFTs.plan_inv(plan_irfft(ry, size(x, dims), dims))]
                 @test eltype(P) <: Real
@@ -100,17 +91,17 @@ function test_fft_backend()
                 @test fftdims(P) == dims
             end
 
+            # BRFFT
             fftw_brfft = complex.(size(x, dims) .* x)
             @test AbstractFFTs.brfft(ry, size(x, dims), dims) ≈ fftw_brfft
             P = plan_brfft(ry, size(x, dims), dims)
             @test P * ry ≈ fftw_brfft
             @test P \ (P * ry) ≈ ry
             @test fftdims(P) == dims
 
+            # IRFFT
             fftw_irfft = complex.(x)
             @test AbstractFFTs.irfft(ry, size(x, dims), dims) ≈ fftw_irfft
-            # test plan_rfft and also inv and plan_inv of plan_irfft, which should all give 
-            # functionally identical plans
             for P in [plan_irfft(ry, size(x, dims), dims), inv(plan_rfft(x, dims)), 
                       AbstractFFTs.plan_inv(plan_rfft(x, dims))]
                 @test P * ry ≈ fftw_irfft

test/runtests.jl

@@ -3,13 +3,24 @@ using Test
 using AbstractFFTs
 using AbstractFFTs.TestUtils
 using ChainRulesTestUtils
+import Unitful
 
 Random.seed!(1234)
 
 include("TestPlans.jl")
 
 using .TestPlans
-test_fft_backend() # Tests for FFT plans (and operations in AbstractFFTs that derive from them)
+test_fft_backend(Array) # Tests for FFT plans (and operations in AbstractFFTs that derive from them)
+
+@testset "rfft sizes" begin
+    A = rand(11, 10)
+    @test @inferred(AbstractFFTs.rfft_output_size(A, 1)) == (6, 10)
+    @test @inferred(AbstractFFTs.rfft_output_size(A, 2)) == (11, 6)
+    A1 = rand(6, 10); A2 = rand(11, 6)
+    @test @inferred(AbstractFFTs.brfft_output_size(A1, 11, 1)) == (11, 10)
+    @test @inferred(AbstractFFTs.brfft_output_size(A2, 10, 2)) == (11, 10)
+    @test_throws AssertionError AbstractFFTs.brfft_output_size(A1, 10, 2)
+end
 
 @testset "Shift functions" begin
     @test @inferred(AbstractFFTs.fftshift([1 2 3])) == [3 1 2]
@@ -99,7 +110,7 @@ end
     # normalization should be inferable even if region is only inferred as ::Any,
     # need to wrap in another function to test this (note that p.region::Any for
     # p::TestPlan)
-    f9(p::Plan{T}, sz) where {T} = AbstractFFTs.normalization(real(T), sz, fftdims(p))
+    f9(p::AbstractFFTs.Plan{T}, sz) where {T} = AbstractFFTs.normalization(real(T), sz, fftdims(p))
     @test @inferred(f9(plan_fft(zeros(10), 1), 10)) == 1/10
 end