real fft

realfft.jl

using Base: OneTo
using AbstractFFTs: ScaledPlan, normalization
using FFTW
using LinearAlgebra

struct RealFFT{T,N,V<:AbstractArray{T,N}}
    data::V
    nx::Int
end

to_real(A::RealFFT{T,N}) where {T,N} = view(A.data, OneTo(A.nx), ntuple(i -> :, N - 1)...)
to_complex(A::RealFFT{T,N}) where {T,N} = reinterpret(complex(T), A.data)

rfft_plan!(A::RealFFT{T,N}, flags::Integer=FFTW.ESTIMATE, timelimit::Real=FFTW.NO_TIMELIMIT) where {T,N} =
    FFTW.rFFTWPlan{T,FFTW.FORWARD,true,N}(to_real(A), to_complex(A), 1:N, flags, timelimit)

brfft_plan!(A::RealFFT{T,N}, flags::Integer=FFTW.ESTIMATE, timelimit::Real=FFTW.NO_TIMELIMIT) where {T,N} =
    FFTW.rFFTWPlan{Complex{T},FFTW.BACKWARD,true,N}(to_complex(A), to_real(A), 1:N, flags, timelimit)

irfft_plan!(A::RealFFT{T,N}, args...) where {T,N} =
    ScaledPlan(brfft_plan!(A, args...), normalization(T, size(to_real(A)), 1:N))

function example(n = 10, T = Float64, p = 2)
    FFTW.set_num_threads(p)
    @assert iseven(n)
    @time A = RealFFT(Array{T}(undef, n + 2, n, n), n)
    @time to_real(A) .= randn.()
    @time mul!(to_complex(A), rfft_plan!(A), to_real(A))
    @time mul!(to_real(A), irfft_plan!(A), to_complex(A))

    return to_real(A)
end

function example2(n = 10, T = Float64, p = 2)
    FFTW.set_num_threads(p)
    @assert iseven(n)
    @time A = randn(T, n, n, n)
    @time B = rfft(A)
    @time C = irfft(B, n)

    return C
end