Add DualPlan

Author: dlfivefifty

ext/AbstractFFTsForwardDiffExt.jl

@@ -1,18 +1,11 @@
 module AbstractFFTsForwardDiffExt
 
 using AbstractFFTs
+using AbstractFFTs.LinearAlgebra
 import ForwardDiff
 import ForwardDiff: Dual
-import AbstractFFTs: Plan
+import AbstractFFTs: Plan, mul!
 
-for P in (:Plan, :ScaledPlan)  # need ScaledPlan to avoid ambiguities
-    @eval begin
-        AbstractFFTs.plan_mul(p::AbstractFFTs.$P, x::AbstractArray{DT}) where DT<:Dual = array2dual(DT, p * dual2array(x))
-        AbstractFFTs.plan_mul(p::AbstractFFTs.$P, x::AbstractArray{Complex{DT}}) where DT<:Dual = array2dual(DT, p * dual2array(x))
-    end
-end
-
-mul!(y::AbstractArray{<:Union{Dual,Complex{<:Dual}}}, p::Plan, x::AbstractArray{<:Union{Dual,Complex{<:Dual}}}) = copyto!(y, p*x)
 
 AbstractFFTs.complexfloat(x::AbstractArray{<:Dual}) = AbstractFFTs.complexfloat.(x)
 AbstractFFTs.complexfloat(d::Dual{T,V,N}) where {T,V,N} = convert(Dual{T,float(V),N}, d) + 0im
@@ -26,10 +19,37 @@ array2dual(DT::Type{<:Dual}, x::Array{T}) where T = reinterpret(reshape, DT, rea
 array2dual(DT::Type{<:Dual}, x::Array{<:Complex{T}}) where T = complex.(array2dual(DT, real(x)), array2dual(DT, imag(x)))
 
 
+########
+# DualPlan
+# represents a plan acting on dual numbers. We wrap a plan acting on a higher dimensional tensor
+# as an array of duals can be reinterpreted as a higher dimensional array.
+# This allows standard FFTW plans to act on arrays of duals.
+#####
+struct DualPlan{T,P} <: Plan{T}
+    p::P
+    DualPlan{T,P}(p) where {T,P} = new(p)
+end
+
+DualPlan(::Type{Dual{Tag,V,N}}, p::Plan{T}) where {Tag,T<:Real,V,N} = DualPlan{Dual{Tag,T,N},typeof(p)}(p)
+DualPlan(::Type{Dual{Tag,V,N}}, p::Plan{Complex{T}}) where {Tag,T<:Real,V,N} = DualPlan{Complex{Dual{Tag,T,N}},typeof(p)}(p)
+Base.size(p::DualPlan) = Base.tail(size(p.p))
+Base.:*(p::DualPlan{DT}, x::AbstractArray{DT}) where DT<:Dual = array2dual(DT, p.p * dual2array(x))
+Base.:*(p::DualPlan{Complex{DT}}, x::AbstractArray{Complex{DT}}) where DT<:Dual = array2dual(DT, p.p * dual2array(x))
+
+function LinearAlgebra.mul!(y::AbstractArray{<:Dual}, p::DualPlan, x::AbstractArray{<:Dual})
+    LinearAlgebra.mul!(dual2array(y), p.p, dual2array(x)) # even though `Dual` are immutable, when in an `Array` they can be modified.
+    y
+end
+
+function LinearAlgebra.mul!(y::AbstractArray{<:Complex{<:Dual}}, p::DualPlan, x::AbstractArray{<:Union{Dual,Complex{<:Dual}}})
+    copyto!(y, p*x) # Complex duals cannot be reinterpret in-place
+end
+
+
 for plan in (:plan_fft, :plan_ifft, :plan_bfft, :plan_rfft)
     @eval begin
-        AbstractFFTs.$plan(x::AbstractArray{<:Dual}, dims=1:ndims(x)) = AbstractFFTs.$plan(dual2array(x), 1 .+ dims)
-        AbstractFFTs.$plan(x::AbstractArray{<:Complex{<:Dual}}, dims=1:ndims(x)) = AbstractFFTs.$plan(dual2array(x), 1 .+ dims)
+        AbstractFFTs.$plan(x::AbstractArray{D}, dims=1:ndims(x)) where D<:Dual = DualPlan(D, AbstractFFTs.$plan(dual2array(x), 1 .+ dims))
+        AbstractFFTs.$plan(x::AbstractArray{<:Complex{D}}, dims=1:ndims(x)) where D<:Dual = DualPlan(D, AbstractFFTs.$plan(dual2array(x), 1 .+ dims))
     end
 end
 

test/abstractfftsforwarddiff.jl

@@ -3,6 +3,9 @@ using ForwardDiff
 using Test
 using ForwardDiff: Dual, partials, value
 
+# Needed until https://github.com/JuliaDiff/ForwardDiff.jl/pull/732 is merged
+complexpartials(x, k) = partials(real(x), k) + im*partials(imag(x), k)
+
 @testset "ForwardDiff extension tests" begin
     x1 = Dual.(1:4.0, 2:5, 3:6)
 
@@ -13,42 +16,42 @@ using ForwardDiff: Dual, partials, value
 
     @testset "$f" for f in (fft, ifft, rfft, bfft)
         @test value.(f(x1)) == f(value.(x1))
-        @test partials.(real(f(x1)), 1) + im*partials.(imag(f(x1)), 1) == f(partials.(x1, 1))
-        @test partials.(real(f(x1)), 2) + im*partials.(imag(f(x1)), 2) == f(partials.(x1, 2))
+        @test complexpartials.(f(x1), 1) == f(partials.(x1, 1))
+        @test complexpartials.(f(x1), 2) == f(partials.(x1, 2))
     end
 
     @test ifft(fft(x1)) ≈ x1
     @test irfft(rfft(x1), length(x1)) ≈ x1
     @test brfft(rfft(x1), length(x1)) ≈ 4x1
 
     f = x -> real(fft([x; 0; 0])[1])
-    @test derivative(f,0.1) ≈ 1
+    @test ForwardDiff.derivative(f,0.1) ≈ 1
 
     r = x -> real(rfft([x; 0; 0])[1])
-    @test derivative(r,0.1) ≈ 1
+    @test ForwardDiff.derivative(r,0.1) ≈ 1
 
 
     n = 100
     θ = range(0,2π; length=n+1)[1:end-1]
     # emperical from Mathematical
-    @test derivative(ω -> fft(exp.(ω .* cos.(θ)))[1]/n, 1) ≈ 0.565159103992485
+    @test ForwardDiff.derivative(ω -> fft(exp.(ω .* cos.(θ)))[1]/n, 1) ≈ 0.565159103992485
 
     # c = x -> dct([x; 0; 0])[1]
     # @test derivative(c,0.1) ≈ 1
 
     @testset "matrix" begin
         A = x1 * (1:10)'
         @test value.(fft(A)) == fft(value.(A))
-        @test partials.(fft(A), 1) == fft(partials.(A, 1))
-        @test partials.(fft(A), 2) == fft(partials.(A, 2))
+        @test complexpartials.(fft(A), 1) == fft(partials.(A, 1))
+        @test complexpartials.(fft(A), 2) == fft(partials.(A, 2))
 
         @test value.(fft(A, 1)) == fft(value.(A), 1)
-        @test partials.(fft(A, 1), 1) == fft(partials.(A, 1), 1)
-        @test partials.(fft(A, 1), 2) == fft(partials.(A, 2), 1)
+        @test complexpartials.(fft(A, 1), 1) == fft(partials.(A, 1), 1)
+        @test complexpartials.(fft(A, 1), 2) == fft(partials.(A, 2), 1)
 
         @test value.(fft(A, 2)) == fft(value.(A), 2)
-        @test partials.(fft(A, 2), 1) == fft(partials.(A, 1), 2)
-        @test partials.(fft(A, 2), 2) == fft(partials.(A, 2), 2)
+        @test complexpartials.(fft(A, 2), 1) == fft(partials.(A, 1), 2)
+        @test complexpartials.(fft(A, 2), 2) == fft(partials.(A, 2), 2)
     end
 
     c1 = complex.(x1)