Dispatch on ProjectionStyle trait for adjoints

Author: Gaurav Arya

src/definitions.jl

@@ -345,16 +345,6 @@ function brfft_output_size(sz::Dims{N}, d::Integer, region) where {N}
     return ntuple(i -> i == d1 ? d : sz[i], Val(N))
 end
 
-function output_size(p::Plan)
-    if projection_style(p) == :none
-        return size(p)
-    elseif projection_style(p) == :real
-        return rfft_output_size(size(p), region(p))
-    elseif projection_style(p) == :real_inv
-        return brfft_output_size(size(p), irfft_dim(p), region(p))
-    end
-end
-
 plan_irfft(x::AbstractArray{Complex{T}}, d::Integer, region; kws...) where {T} =
     ScaledPlan(plan_brfft(x, d, region; kws...),
                normalization(T, brfft_output_size(x, d, region), region))
@@ -590,11 +580,19 @@ plan_brfft
 
 ##############################################################################
 
-# Projection style (:none, :real, or :real_inv) to handle real FFTs
-function projection_style end
-# Length of halved dimension, needed only for irfft 
+struct NoProjectionStyle end
+struct RealProjectionStyle end 
+struct RealInverseProjectionStyle end
+const ProjectionStyle = Union{NoProjectionStyle, RealProjectionStyle, RealInverseProjectionStyle}
+
 function irfft_dim end
 
+ProjectionStyle(p::Plan) = error("No projection style defined for plan")
+output_size(p::Plan) = _output_size(p, ProjectionStyle(p))
+_output_size(p::Plan, ::NoProjectionStyle) = size(p)
+_output_size(p::Plan, ::RealProjectionStyle) = rfft_output_size(size(p), region(p))
+_output_size(p::Plan, ::RealInverseProjectionStyle) = brfft_output_size(size(p), irfft_dim(p), region(p))
+
 mutable struct AdjointPlan{T,P} <: Plan{T}
     p::P
     pinv::Plan
@@ -611,31 +609,38 @@ Base.adjoint(p::Plan{T}) where {T} = AdjointPlan{T}(p)
 size(p::AdjointPlan) = output_size(p)
 output_size(p::AdjointPlan) = size(p)
 
-function Base.:*(p::AdjointPlan{T}, x::AbstractArray) where {T}
+Base.:*(p::AdjointPlan, x::AbstractArray) = _mul(p, x, ProjectionStyle(p.p))
+
+function _mul(p::AdjointPlan{T}, x::AbstractArray, ::NoProjectionStyle) where {T}
+    dims = region(p.p)
+    N = normalization(T, size(p.p), dims)
+    return 1/N * (p.p \ x)
+end
+
+function _mul(p::AdjointPlan{T}, x::AbstractArray, ::RealProjectionStyle) where {T}
     dims = region(p.p)
+    N = normalization(T, size(p.p), dims)
     halfdim = first(dims)
     d = size(p.p, halfdim)
     n = output_size(p.p, halfdim)
-    if projection_style(p.p) == :none
-        N = normalization(T, size(p.p), dims)
-        return 1/N * (p.p \ x)
-    elseif projection_style(p.p) == :real
-        N = normalization(T, size(p.p), dims)
-        scale = reshape(
-            [(i == 1 || (i == n && 2 * (i - 1)) == d) ? 1 : 2 for i in 1:n],
-            ntuple(i -> i == first(dims) ? n : 1, Val(ndims(x)))
-        )
-        return 1/N * (p.p \ (x ./ scale))
-    elseif projection_style(p.p) == :real_inv
-        N = normalization(real(T), output_size(p.p), dims)
-        scale = reshape(
-            [(i == 1 || (i == d && 2 * (i - 1)) == n) ? 1 : 2 for i in 1:d],
-            ntuple(i -> i == first(dims) ? d : 1, Val(ndims(x)))
-        )
-        return 1/N * scale .* (p.p \ x)
-    else
-        error("plan must define a valid projection style")
-    end
+    scale = reshape(
+        [(i == 1 || (i == n && 2 * (i - 1)) == d) ? 1 : 2 for i in 1:n],
+        ntuple(i -> i == first(dims) ? n : 1, Val(ndims(x)))
+    )
+    return 1/N * (p.p \ (x ./ scale))
+end
+
+function _mul(p::AdjointPlan{T}, x::AbstractArray, ::RealInverseProjectionStyle) where {T}
+    dims = region(p.p)
+    N = normalization(real(T), output_size(p.p), dims)
+    halfdim = first(dims)
+    n = size(p.p, halfdim)
+    d = output_size(p.p, halfdim)
+    scale = reshape(
+        [(i == 1 || (i == n && 2 * (i - 1)) == d) ? 1 : 2 for i in 1:n],
+        ntuple(i -> i == first(dims) ? n : 1, Val(ndims(x)))
+    )
+    return 1/N * scale .* (p.p \ x)
 end
 
 plan_inv(p::AdjointPlan) = AdjointPlan(plan_inv(p.p))

test/testplans.jl

@@ -21,8 +21,8 @@ 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
+AbstractFFTs.ProjectionStyle(::TestPlan) = AbstractFFTs.NoProjectionStyle()
+AbstractFFTs.ProjectionStyle(::InverseTestPlan) = AbstractFFTs.NoProjectionStyle()
 
 function AbstractFFTs.plan_fft(x::AbstractArray{T}, region; kwargs...) where {T}
     return TestPlan{T}(region, size(x))
@@ -111,8 +111,8 @@ mutable struct InverseTestRPlan{T,N,G} <: Plan{T}
     end
 end
 
-AbstractFFTs.projection_style(::TestRPlan) = :real
-AbstractFFTs.projection_style(::InverseTestRPlan) = :real_inv
+AbstractFFTs.ProjectionStyle(::TestRPlan) = AbstractFFTs.RealProjectionStyle()
+AbstractFFTs.ProjectionStyle(::InverseTestRPlan) = AbstractFFTs.RealInverseProjectionStyle()
 AbstractFFTs.irfft_dim(p::InverseTestRPlan) = p.d
 
 function AbstractFFTs.plan_rfft(x::AbstractArray{T}, region; kwargs...) where {T}