Test plan adjoints and AD rules

Author: gaurav-arya

test/runtests.jl

@@ -3,6 +3,7 @@
 using AbstractFFTs
 using AbstractFFTs: Plan
 using ChainRulesTestUtils
+using ChainRulesCore: NoTangent
 
 using LinearAlgebra
 using Random
@@ -197,6 +198,39 @@ end
     @test @inferred(f9(plan_fft(zeros(10), 1), 10)) == 1/10
 end
 
+@testset "adjoint" begin
+    @testset "complex fft adjoint" begin
+        for x in (randn(3), randn(3, 4), randn(3, 4, 5))
+            N = ndims(x)
+            y = randn(size(x))
+            for dims in unique((1, 1:N, N))
+                P = plan_fft(x, dims)
+                @test dot(y, P * x) ≈ dot(P' * y, x)
+                @test_broken dot(y, P \ x) ≈ dot(P' \ y, x)
+                Pinv = plan_ifft(x)
+                @test dot(x, Pinv * y) ≈ dot(Pinv' * x, y)
+                @test_broken dot(x, Pinv \ y) ≈ dot(Pinv' \ x, y)
+            end
+        end
+    end
+    @testset "real fft adjoint" begin
+        for x in (randn(3), randn(4), randn(3, 4), randn(3, 4, 5)) # test odd and even lengths
+            N = ndims(x)
+            for dims in unique((1, 1:N, N))
+                P = plan_rfft(similar(x), dims)        
+                y_real = randn(size(P * x))
+                y_imag = randn(size(P * x))
+                y = y_real .+ y_imag .* im 
+                @test dot(y_real, real.(P * x)) + dot(y_imag, imag.(P * x)) ≈ dot(P' * y, x)
+                @test_broken dot(y_real, real.(P \ x)) + dot(y_imag, imag.(P \ x)) ≈ dot(P' * y, x)
+                Pinv = plan_irfft(similar(y), size(x)[first(dims)], dims)
+                @test dot(x, Pinv * y) ≈ dot(y_real, real.(Pinv' * x)) + dot(y_imag, imag.(Pinv' * x))
+                @test_broken dot(x, Pinv \ y) ≈ dot(y_real, real.(Pinv' \ x)) + dot(y_imag, imag.(Pinv' \ x))
+            end
+        end
+    end
+end
+
 @testset "ChainRules" begin
     @testset "shift functions" begin
         for x in (randn(3), randn(3, 4), randn(3, 4, 5))
@@ -218,7 +252,7 @@ end
     end
 
     @testset "fft" begin
-        for x in (randn(3), randn(3, 4), randn(3, 4, 5))
+        for x in (randn(2), randn(2, 3), randn(3, 4, 5))
             N = ndims(x)
             complex_x = complex.(x)
             for dims in unique((1, 1:N, N))
@@ -229,8 +263,12 @@ end
                     test_rrule(f, complex_x, dims)
                 end
 
-                test_frule(rfft, x, dims)
-                test_rrule(rfft, x, dims)
+                for pf in (plan_fft, plan_ifft, plan_bfft) 
+                    test_frule(*, pf(x, dims) ⊢ NoTangent(), x)
+                    test_rrule(*, pf(x, dims) ⊢ NoTangent(), x)
+                    test_frule(*, pf(complex_x, dims) ⊢ NoTangent(), complex_x)
+                    test_rrule(*, pf(complex_x, dims) ⊢ NoTangent(), complex_x)
+                end
 
                 for f in (irfft, brfft)
                     for d in (2 * size(x, first(dims)) - 1, 2 * size(x, first(dims)) - 2)
@@ -240,6 +278,14 @@ end
                         test_rrule(f, complex_x, d, dims)
                     end
                 end
+
+                for pf in (plan_irfft, plan_brfft) 
+                    for d in (2 * size(x, first(dims)) - 1, 2 * size(x, first(dims)) - 2)
+                        test_frule(*, pf(complex_x, d, dims) ⊢ NoTangent(), complex_x)
+                        test_rrule(*, pf(complex_x, d, dims) ⊢ NoTangent(), complex_x)
+                    end
+                end
+
             end
         end
     end

test/testplans.jl

@@ -1,18 +1,18 @@
-mutable struct TestPlan{T,N} <: Plan{T}
-    region
+mutable struct TestPlan{T,N,G} <: Plan{T}
+    region::G
     sz::NTuple{N,Int}
     pinv::Plan{T}
-    function TestPlan{T}(region, sz::NTuple{N,Int}) where {T,N}
-        return new{T,N}(region, sz)
+    function TestPlan{T}(region::G, sz::NTuple{N,Int}) where {T,N,G}
+        return new{T,N,G}(region, sz)
     end
 end
 
-mutable struct InverseTestPlan{T,N} <: Plan{T}
-    region
+mutable struct InverseTestPlan{T,N,G} <: Plan{T}
+    region::G
     sz::NTuple{N,Int}
     pinv::Plan{T}
-    function InverseTestPlan{T}(region, sz::NTuple{N,Int}) where {T,N}
-        return new{T,N}(region, sz)
+    function InverseTestPlan{T}(region::G, sz::NTuple{N,Int}) where {T,N,G}
+        return new{T,N,G}(region, sz)
     end
 end
 
@@ -21,6 +21,9 @@ Base.ndims(::TestPlan{T,N}) where {T,N} = N
 Base.size(p::InverseTestPlan) = p.sz
 Base.ndims(::InverseTestPlan{T,N}) where {T,N} = N
 
+AbstractFFTs.projection_style(::TestPlan) = :none
+AbstractFFTs.projection_style(::InverseTestPlan) = :none
+
 function AbstractFFTs.plan_fft(x::AbstractArray{T}, region; kwargs...) where {T}
     return TestPlan{T}(region, size(x))
 end
@@ -90,24 +93,28 @@ end
 Base.:*(p::TestPlan, x::AbstractArray) = mul!(similar(x, complex(float(eltype(x)))), p, x)
 Base.:*(p::InverseTestPlan, x::AbstractArray) = mul!(similar(x, complex(float(eltype(x)))), p, x)
 
-mutable struct TestRPlan{T,N} <: Plan{T}
-    region
+mutable struct TestRPlan{T,N,G} <: Plan{T}
+    region::G
     sz::NTuple{N,Int}
     pinv::Plan{T}
-    TestRPlan{T}(region, sz::NTuple{N,Int}) where {T,N} = new{T,N}(region, sz)
+    TestRPlan{T}(region::G, sz::NTuple{N,Int}) where {T,N,G} = new{T,N,G}(region, sz)
 end
 
-mutable struct InverseTestRPlan{T,N} <: Plan{T}
+mutable struct InverseTestRPlan{T,N,G} <: Plan{T}
     d::Int
-    region
+    region::G
     sz::NTuple{N,Int}
     pinv::Plan{T}
-    function InverseTestRPlan{T}(d::Int, region, sz::NTuple{N,Int}) where {T,N}
+    function InverseTestRPlan{T}(d::Int, region::G, sz::NTuple{N,Int}) where {T,N,G}
         sz[first(region)::Int] == d ÷ 2 + 1 || error("incompatible dimensions")
-        return new{T,N}(d, region, sz)
+        return new{T,N,G}(d, region, sz)
     end
 end
 
+AbstractFFTs.projection_style(::TestRPlan) = :real
+AbstractFFTs.projection_style(::InverseTestRPlan) = :real_inv
+AbstractFFTs.irfft_dim(p::InverseTestRPlan) = p.d
+
 function AbstractFFTs.plan_rfft(x::AbstractArray{T}, region; kwargs...) where {T}
     return TestRPlan{T}(region, size(x))
 end