From 17adac8b812346040417adf2066224c722d6a89e Mon Sep 17 00:00:00 2001 From: Gaurav Arya Date: Wed, 19 Jul 2023 18:34:06 -0400 Subject: [PATCH] Clarify documentation, rename _output_size -> output_size --- docs/src/api.md | 23 +++++++++++---- docs/src/implementations.md | 9 +++--- src/definitions.jl | 57 +++++++++++++++++++++++++++---------- 3 files changed, 64 insertions(+), 25 deletions(-) diff --git a/docs/src/api.md b/docs/src/api.md index ef998f8..713e62d 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -1,5 +1,7 @@ # Public Interface +## FFT and FFT planning functions + ```@docs AbstractFFTs.fft AbstractFFTs.fft! @@ -20,15 +22,26 @@ AbstractFFTs.plan_rfft AbstractFFTs.plan_brfft AbstractFFTs.plan_irfft AbstractFFTs.fftdims -Base.adjoint -AbstractFFTs.FFTAdjointStyle -AbstractFFTs.RFFTAdjointStyle -AbstractFFTs.IRFFTAdjointStyle -AbstractFFTs.UnitaryAdjointStyle AbstractFFTs.fftshift AbstractFFTs.fftshift! AbstractFFTs.ifftshift AbstractFFTs.ifftshift! AbstractFFTs.fftfreq AbstractFFTs.rfftfreq +Base.size +``` + +## Adjoint functionality + +The following API is supported by plans that support adjoint functionality. +It is also relevant to implementers of FFT plans that wish to support adjoints. +```@docs +Base.adjoint +AbstractFFTs.AdjointStyle +AbstractFFTs.output_size +AbstractFFTs.adjoint_mul +AbstractFFTs.FFTAdjointStyle +AbstractFFTs.RFFTAdjointStyle +AbstractFFTs.IRFFTAdjointStyle +AbstractFFTs.UnitaryAdjointStyle ``` diff --git a/docs/src/implementations.md b/docs/src/implementations.md index 364d21b..efa3be4 100644 --- a/docs/src/implementations.md +++ b/docs/src/implementations.md @@ -32,11 +32,10 @@ To define a new FFT implementation in your own module, you should * You can also define similar methods of `plan_rfft` and `plan_brfft` for real-input FFTs. -* We offer an experimental `AdjointStyle` trait to enable automatic computation of adjoint plans via [`Base.adjoint`](@ref). - To support adjoints in a new plan, define the trait `AbstractFFTs.AdjointStyle(::MyPlan)`. This should return a subtype of `AS <: AbstractFFTs.AdjointStyle` supporting `AbstractFFTs.adjoint_mul(::Plan, ::AbstractArray, ::AS)` and - `AbstractFFTs._output_size(::Plan, ::AS)`. - - `AbstractFFTs` pre-implements the following adjoint styles: [`AbstractFFTs.FFTAdjointStyle`](@ref), [`AbstractFFTs.RFFTAdjointStyle`](@ref), [`AbstractFFTs.IRFFTAdjointStyle`](@ref), and [`AbstractFFTs.UnitaryAdjointStyle`](@ref). +* To support adjoints in a new plan, define the trait [`AbstractFFTs.AdjointStyle`](@ref). + `AbstractFFTs` implements the following adjoint styles: [`AbstractFFTs.FFTAdjointStyle`](@ref), [`AbstractFFTs.RFFTAdjointStyle`](@ref), [`AbstractFFTs.IRFFTAdjointStyle`](@ref), and [`AbstractFFTs.UnitaryAdjointStyle`](@ref). + To define a new adjoint style of type `AS <: AbstractFFTs.AdjointStyle`, define the methods + [`AbstractFFTs.adjoint_mul`](@ref) and [`AbstractFFTs.output_size`](@ref). The normalization convention for your FFT should be that it computes ``y_k = \sum_j x_j \exp(-2\pi i j k/n)`` for a transform of length ``n``, and the "backwards" (unnormalized inverse) transform computes the same thing but with ``\exp(+2\pi i jk/n)``. diff --git a/src/definitions.jl b/src/definitions.jl index 4c59b00..b642ac5 100644 --- a/src/definitions.jl +++ b/src/definitions.jl @@ -583,12 +583,19 @@ plan_brfft ############################################################################## +""" + AbstractFFTs.AdjointStyle(::Plan) + +Return the adjoint style of a plan, enabling automatic computation of adjoint plans via +[`Base.adjoint`](@ref). Instructions for supporting adjoint styles are provided in the +[implementation instructions](implementations.md#Defining-a-new-implementation). +""" abstract type AdjointStyle end """ FFTAdjointStyle() -Projection style for complex to complex discrete Fourier transforms that normalize +Adjoint style for complex to complex discrete Fourier transforms that normalize the output analogously to [`fft`](@ref). Since the Fourier transform is unitary up to a scaling, the adjoint simply applies @@ -599,8 +606,8 @@ struct FFTAdjointStyle <: AdjointStyle end """ RFFTAdjointStyle() -Projection style for real to complex discrete Fourier transforms that halve -one of the output's dimensions and normalize the output analogously to [`rfft`](@ref). +Adjoint style for real to complex discrete Fourier transforms that halve one of +the output's dimensions and normalize the output, analogously to [`rfft`](@ref). Since the Fourier transform is unitary up to a scaling, the adjoint applies the transform's inverse, but with additional logic to handle the fact that the output is projected @@ -611,8 +618,8 @@ struct RFFTAdjointStyle <: AdjointStyle end """ IRFFTAdjointStyle(d::Dim) -Projection style for complex to real discrete Fourier transforms that expect -an input with a halved dimension and normalize the output analogously to [`irfft`](@ref), +Adjoint style for complex to real discrete Fourier transforms that expect an input +with a halved dimension and normalize the output, analogously to [`irfft`](@ref), where `d` is the original length of the dimension. Since the Fourier transform is unitary up to a scaling, the adjoint applies the transform's @@ -626,15 +633,22 @@ end """ UnitaryAdjointStyle() -Projection style for unitary transforms, whose adjoint equals their inverse. +Adjoint style for unitary transforms, whose adjoint equals their inverse. """ struct UnitaryAdjointStyle <: AdjointStyle end -output_size(p::Plan) = _output_size(p, AdjointStyle(p)) -_output_size(p::Plan, ::FFTAdjointStyle) = size(p) -_output_size(p::Plan, ::RFFTAdjointStyle) = rfft_output_size(size(p), fftdims(p)) -_output_size(p::Plan, s::IRFFTAdjointStyle) = brfft_output_size(size(p), s.dim, fftdims(p)) -_output_size(p::Plan, ::UnitaryAdjointStyle) = size(p) +""" + output_size(p::Plan) + +Return the size of the output of a plan `p`. + +Implementations of a new adjoint style `AS <: AbstractFFTs.AdjointStyle` should define `output_size(::Plan, ::AS)`. +""" +output_size(p::Plan) = output_size(p, AdjointStyle(p)) +output_size(p::Plan, ::FFTAdjointStyle) = size(p) +output_size(p::Plan, ::RFFTAdjointStyle) = rfft_output_size(size(p), fftdims(p)) +output_size(p::Plan, s::IRFFTAdjointStyle) = brfft_output_size(size(p), s.dim, fftdims(p)) +output_size(p::Plan, ::UnitaryAdjointStyle) = size(p) struct AdjointPlan{T,P<:Plan} <: Plan{T} p::P @@ -645,9 +659,7 @@ end (p::Plan)' adjoint(p::Plan) -Form the adjoint operator of an FFT plan. Returns a plan that performs the adjoint operation of -the original plan. Note that this differs from the corresponding backwards plan in the case of real -FFTs due to the halving of one of the dimensions of the FFT output, as described in [`rfft`](@ref). +Return a plan that performs the adjoint operation of the original plan. !!! note Adjoint plans do not currently support `LinearAlgebra.mul!`. Further, as a new addition to `AbstractFFTs`, @@ -658,10 +670,25 @@ Base.adjoint(p::AdjointPlan) = p.p # always have AdjointPlan inside ScaledPlan. Base.adjoint(p::ScaledPlan) = ScaledPlan(p.p', p.scale) +""" + size(p::Plan) + +Return the size of the input of a plan `p`. +""" size(p::AdjointPlan) = output_size(p.p) output_size(p::AdjointPlan) = size(p.p) -Base.:*(p::AdjointPlan, x::AbstractArray) = adjoint_mul(p.p, x, AdjointStyle(p.p)) +Base.:*(p::AdjointPlan, x::AbstractArray) = adjoint_mul(p.p, x) + +""" + adjoint_mul(p::Plan, x::AbstractArray) + +Multiply an array `x` by the adjoint of a plan `p`. This is equivalent to `p' * x`. + +Implementations of a new adjoint style `AS <: AbstractFFTs.AdjointStyle` should define +`adjoint_mul(::Plan, ::AbstractArray, ::AS)`. +""" +adjoint_mul(p::Plan, x::AbstractArray) = adjoint_mul(p, x, AdjointStyle(p)) function adjoint_mul(p::Plan{T}, x::AbstractArray, ::FFTAdjointStyle) where {T} dims = fftdims(p)