Added rfft! and irfft! functionality through PaddedRFFTArray type.

src/FFTW.jl

@@ -44,5 +44,6 @@ end
 
 include("fft.jl")
 include("dct.jl")
+include("rfft!.jl")
 
 end # module

src/rfft!.jl

@@ -0,0 +1,234 @@
+import Base: IndexStyle, getindex, setindex!, eltype, \, similar, copy, real, read!
+
+export PaddedRFFTArray, plan_rfft!, rfft!, plan_irfft!, plan_brfft!, brfft!, irfft!
+
+
+
+# As the time this code was written the new `ReinterpretArray` introduced in
+# Julia v0.7 had major performace issues. Those issues were bypassed with the usage of the
+# custom getindex and setindex! below. Hopefully, once the performance issues with ReinterpretArray
+# are solved we can just index the reinterpret array directly.
+
+struct PaddedRFFTArray{T<:fftwReal,N,L} <: DenseArray{Complex{T},N}
+    data::Array{T,N}
+    r::SubArray{T,N,Array{T,N},NTuple{N,UnitRange{Int}},L} # Real view skipping padding
+    c::Base.ReinterpretArray{Complex{T},N,T,Array{T,N}}
+
+    function PaddedRFFTArray{T,N}(rr::Array{T,N},nx::Int) where {T<:fftwReal,N}
+        rrsize = size(rr)
+        fsize = rrsize[1]
+        iseven(fsize) || throw(
+            ArgumentError("First dimension of allocated array must have even number of elements"))
+        (nx == fsize-2 || nx == fsize-1) || throw(
+            ArgumentError("Number of elements on the first dimension of array must be either 1 or 2 less than the number of elements on the first dimension of the allocated array"))
+        fsize = fsize÷2
+        csize = (fsize, rrsize[2:end]...)
+        c = reinterpret(Complex{T}, rr)
+        rsize = (nx,rrsize[2:end]...)
+        r = view(rr,(1:l for l in rsize)...)
+        return  new{T, N, N === 1 ? true : false}(rr,r,c)
+    end # function
+end # struct
+
+@inline real(S::PaddedRFFTArray) = S.r
+
+@inline complex_view(S::PaddedRFFTArray) = S.c
+
+@inline data(S::PaddedRFFTArray) = S.data
+
+copy(S::PaddedRFFTArray) = PaddedRFFTArray(copy(data(S)),size(real(S),1))
+
+similar(f::PaddedRFFTArray,::Type{T},dims::Tuple{Vararg{Int,N}}) where {T, N} =
+    PaddedRFFTArray{T}(dims) 
+similar(f::PaddedRFFTArray{T,N,L},dims::NTuple{N2,Int}) where {T,N,L,N2} =
+    PaddedRFFTArray{T}(dims) 
+similar(f::PaddedRFFTArray,::Type{T}) where {T} =
+    PaddedRFFTArray{T}(size(real(f))) 
+similar(f::PaddedRFFTArray{T,N}) where {T,N} = 
+    PaddedRFFTArray{T,N}(similar(data(f)), size(real(f),1)) 
+
+size(S::PaddedRFFTArray) =
+    size(complex_view(S))
+
+IndexStyle(::Type{T}) where {T<:PaddedRFFTArray} = 
+    IndexLinear()
+
+@inline function getindex(A::PaddedRFFTArray{T,N}, i2::Integer) where {T,N}
+    d = data(A)
+    i = 2i2
+    @boundscheck checkbounds(d,i)
+    @inbounds begin 
+        return Complex{T}(d[i-1],d[i])
+    end
+end    
+
+@inline @generated function getindex(A::PaddedRFFTArray{T,N}, I2::Vararg{Integer,N}) where {T,N}
+    ip = :(2*I2[1])
+    t = Expr(:tuple)
+    for i=2:N
+        push!(t.args,:(I2[$i]))
+    end
+    quote
+        d = data(A)
+        i = $ip
+        I = $t
+        @boundscheck checkbounds(d,i,I...)
+        @inbounds begin 
+            return Complex{T}(d[i-1,I...],d[i,I...])
+        end
+    end
+end
+
+@inline function setindex!(A::PaddedRFFTArray{T,N},x, i2::Integer) where {T,N}
+    d = data(A)
+    i = 2i2
+    @boundscheck checkbounds(d,i)
+    @inbounds begin 
+        d[i-1] = real(x)
+        d[i] = imag(x)
+    end
+    A
+end
+
+@inline @generated function setindex!(A::PaddedRFFTArray{T,N}, x, I2::Vararg{Integer,N}) where {T,N}
+    ip = :(2*I2[1])
+    t = Expr(:tuple)
+    for i=2:N
+        push!(t.args,:(I2[$i]))
+    end
+    quote
+        d = data(A)
+        i = $ip
+        I = $t
+        @boundscheck checkbounds(d,i,I...)
+        @inbounds begin 
+            d[i-1,I...] = real(x)
+            d[i,I...] = imag(x)
+        end
+        A
+    end
+end
+
+PaddedRFFTArray(rr::Array{T,N},nx::Int) where {T<:fftwReal,N} = PaddedRFFTArray{T,N}(rr,nx)
+
+function PaddedRFFTArray{T}(ndims::Vararg{Integer,N}) where {T,N}
+    fsize = (ndims[1]÷2 + 1)*2
+    a = zeros(T,(fsize, ndims[2:end]...))
+    PaddedRFFTArray{T,N}(a, ndims[1])
+end
+
+PaddedRFFTArray{T}(ndims::NTuple{N,Integer}) where {T,N} =
+    PaddedRFFTArray{T}(ndims...)
+
+PaddedRFFTArray(ndims::Vararg{Integer,N}) where N = 
+    PaddedRFFTArray{Float64}(ndims...)
+
+PaddedRFFTArray(ndims::NTuple{N,Integer}) where N = 
+    PaddedRFFTArray{Float64}(ndims...)
+
+function PaddedRFFTArray{T}(a::AbstractArray{<:Real,N}) where {T<:fftwReal,N}
+    t = PaddedRFFTArray{T}(size(a))
+    @inbounds copyto!(t.r, a) 
+    return t
+end
+
+PaddedRFFTArray(a::AbstractArray{<:Real}) = PaddedRFFTArray{Float64}(a)
+
+function PaddedRFFTArray(stream, dims)
+    field = PaddedRFFTArray(dims)
+    return read!(stream,field)
+end
+
+function PaddedRFFTArray{T}(stream, dims) where T
+    field = PaddedRFFTArray{T}(dims)
+    return read!(stream,field)
+end
+
+function read!(file::AbstractString, field::PaddedRFFTArray)
+    open(file) do io 
+       return read!(io,field) 
+    end
+end
+
+# Read a binary file of an unpaded array directly to a PaddedRFFT array, without the need
+# of the creation of a intermediary Array. If the data is already padded then the user
+# should just use PaddedRFFTArray{T}(read("file",unpaddeddim),nx)
+function read!(stream::IO, field::PaddedRFFTArray{T,N,L}) where {T,N,L}
+    rr = data(field)
+    dims = size(real(field))
+    nx = dims[1]
+    nb = sizeof(T)*nx
+    npencils = prod(dims)÷nx
+    npad = iseven(nx) ? 2 : 1
+    for i=0:(npencils-1)
+        unsafe_read(stream,Ref(rr,Int((nx+npad)*i+1)),nb)
+    end
+    return field
+end
+
+
+###########################################################################################
+# Foward plans
+
+function plan_rfft!(X::PaddedRFFTArray{T,N}, region;
+                   flags::Integer=ESTIMATE,
+                   timelimit::Real=NO_TIMELIMIT) where {T<:fftwReal,N}
+
+    (1 in region) || throw(ArgumentError("The first dimension must always be transformed"))
+    if flags&ESTIMATE != 0
+        p = rFFTWPlan{T,FORWARD,true,N}(real(X), complex_view(X), region, flags, timelimit)
+    else
+        x = similar(X)
+        p = rFFTWPlan{T,FORWARD,true,N}(real(x), complex_view(x), region, flags, timelimit)
+    end
+    return p
+end
+
+plan_rfft!(f::PaddedRFFTArray;kws...) = plan_rfft!(f, 1:ndims(f); kws...)
+
+*(p::rFFTWPlan{T,FORWARD,true,N},f::PaddedRFFTArray{T,N}) where {T<:fftwReal,N} = 
+    (mul!(complex_view(f), p, real(f)); f)
+
+rfft!(f::PaddedRFFTArray, region=1:ndims(f)) = plan_rfft!(f, region) * f
+
+function \(p::rFFTWPlan{T,FORWARD,true,N},f::PaddedRFFTArray{T,N}) where {T<:fftwReal,N}
+    isdefined(p,:pinv) || (p.pinv = plan_irfft!(f,p.region))
+    return p.pinv * f
+end
+
+
+##########################################################################################
+# Inverse plans
+
+function plan_brfft!(X::PaddedRFFTArray{T,N}, region;
+                    flags::Integer=ESTIMATE,
+                    timelimit::Real=NO_TIMELIMIT) where {T<:fftwReal,N}
+    (1 in region) || throw(ArgumentError("The first dimension must always be transformed"))
+    if flags&ESTIMATE != 0
+        return rFFTWPlan{Complex{T},BACKWARD,true,N}(complex_view(X), real(X), region, flags,timelimit)
+    else
+        a = similar(X)
+        return rFFTWPlan{Complex{T},BACKWARD,true,N}(complex_view(a), real(a), region, flags,timelimit)
+    end
+end
+
+plan_brfft!(f::PaddedRFFTArray;kws...) = plan_brfft!(f,1:ndims(f);kws...)
+
+*(p::rFFTWPlan{Complex{T},BACKWARD,true,N},f::PaddedRFFTArray{T,N}) where {T<:fftwReal,N} = 
+    (mul!(real(f), p, complex_view(f)); real(f))
+
+brfft!(f::PaddedRFFTArray, region=1:ndims(f)) = plan_brfft!(f, region) * f
+
+function plan_irfft!(x::PaddedRFFTArray{T,N}, region; kws...) where {T,N}
+    ScaledPlan(plan_brfft!(x, region; kws...),normalization(T, size(real(x)), region))
+end
+
+plan_irfft!(f::PaddedRFFTArray;kws...) = plan_irfft!(f,1:ndims(f);kws...)
+
+*(p::ScaledPlan,f::PaddedRFFTArray) = begin
+    p.p * f
+    rmul!(data(f), p.scale)
+    real(f)
+end
+
+irfft!(f::PaddedRFFTArray, region=1:ndims(f)) = plan_irfft!(f,region) * f

test/runtests.jl

@@ -512,3 +512,5 @@ let A = rand(Float32, 35), Ac = rand(Complex{Float32}, 35)
     @test_throws ArgumentError plan_rfft(Array{Float32}(undef, 32)) * view(A, 2:33)
     @test_throws ArgumentError plan_fft(Array{Complex{Float32}}(undef, 32)) * view(Ac, 2:33)
 end
+
+include("test_rfft!.jl")

test/test_rfft!.jl

@@ -0,0 +1,66 @@
+let a = rand(Float64,(8,4,4)), b = PaddedRFFTArray(a), c = copy(b)
+
+@testset "PaddedRFFTArray creation" begin
+  @test a == real(b)
+  @test c == b
+  @test c.r == b.r
+  @test typeof(similar(b)) === typeof(b)
+  @test size(similar(b,Float32)) === size(b)
+  @test size(similar(b,Float32).r) === size(b.r)
+  @test size(similar(b,(4,4,4)).r) === (4,4,4)
+  @test size(similar(b,Float32,(4,4,4)).r) === (4,4,4) 
+end
+
+@testset "rfft! and irfft!" begin
+  @test rfft(a) ≈ rfft!(b) 
+  @test a ≈ irfft!(b)
+  @test rfft(a,1:2) ≈ rfft!(b,1:2) 
+  @test a ≈ irfft!(b,1:2)
+  @test rfft(a,(1,3)) ≈ rfft!(b,(1,3)) 
+  @test a ≈ irfft!(b,(1,3))
+
+  p = plan_rfft!(c)
+  @test p*c ≈ rfft!(b)
+  @test p\c ≈ irfft!(b)
+
+  a = rand(Float64,(9,4,4))
+  b = PaddedRFFTArray(a)
+  @test a == real(b)
+  @test rfft(a) ≈ rfft!(b) 
+  @test a ≈ irfft!(b)
+  @test rfft(a,1:2) ≈ rfft!(b,1:2) 
+  @test a ≈ irfft!(b,1:2)
+  @test rfft(a,(1,3)) ≈ rfft!(b,(1,3)) 
+  @test a ≈ irfft!(b,(1,3))
+end
+
+@testset "Read binary file to PaddedRFFTArray" begin
+  for s in ((8,4,4),(9,4,4),(8,),(9,))
+    aa = rand(Float64,s)
+    f = IOBuffer()
+    write(f,aa)
+    @test aa == real(PaddedRFFTArray(seekstart(f),s))
+    aa = rand(Float32,s)
+    f = IOBuffer()
+    write(f,aa)
+    @test aa == real(PaddedRFFTArray{Float32}(seekstart(f),s))
+  end
+end
+
+@testset "brfft!" begin
+  a = rand(Float64,(4,4))
+  b = PaddedRFFTArray(a)
+  rfft!(b)
+  @test (brfft!(b) ./ 16) ≈ a
+end
+
+@testset "FFTW MEASURE flag" begin
+  c = similar(b)
+  p = plan_rfft!(c,flags=FFTW.MEASURE)
+  p.pinv = plan_irfft!(c,flags=FFTW.MEASURE)
+  c .= b 
+  @test c == b
+  @test p*c ≈ rfft!(b)
+  @test p\c ≈ irfft!(b)
+end
+end #let block