In-place rfft with padded arrays

Hi,

It seems Julia does not have direct support for something like rfft!
I found this question on StackOverFlow which has been asked some years ago by now.

Is it possible now to “easily” perform a in-place rfft with the usage of A_mult_B! ?

I am not a real programer myself nor very literate on the FFTW library.
My use would be for 3D arrays with Nx x Ny x Nz elements.

I am able to create the (Nx+2) x Ny x Nz (the needed padding). Am I able to perform the inplace fft with pure Julia code? If so, can someone give me some directions on that?
I guess it can be done by directly calling the proper fftw library function, but I’d like to stay away from that if possible.

1 Like

No, no in-place high-level rfft! function is currently defined, although this functionality is indeed supported by the underlying FFTW library.

However, you can do plan_rfft and then apply the plan in-place to a pre-allocated output array with A_mul_B!. (This is distinct from overwriting the input in-place.) Is that what you want?

Thank you for the reply!

That is what I wanted. But I might get away for now with this usage of A_mul_B! you presented.
I’ll try digging into the fftw library and Julia’s ccall and hopefully once the size of my array became too big I’ll be able to write a wrapper to call the in-place rfftw directly from the library.

I’ve come up with the following code to perform rfft! and irfft!.
It uses as a basis the FFTW of Julia, so I didn’t have to mess with ccall directly. I performed a few test and it seems to work… But I don’t guarantee anything!

module InPlaceRealFFT

import Base: size, linearindexing, getindex, setindex!, eltype, *

export TransformableField , plan_rfft!, rfft!, plan_irfft!, irfft!

const Float3264 = Union{Float32,Float64}

immutable TransformableField{T<:Float3264,N} <: AbstractArray{T,N}
  c::Array{Complex{T},N} # Complex view of the array
  r::SubArray # Real view skipping padding
  rr::Array{T,N} # Raw real data, including padding

  function (::Type{TransformableField}){T,N}(rr::Array{T,N},nx::Int)
    fsize = [size(rr)...]
    assert(rem(fsize[1],2)==0)
    assert(nx == fsize[1]-2 || nx == fsize[1]-1)
    fsize[1] = fsize[1]/2
    rsize = (fsize...)
    c = reinterpret(Complex{T}, rr, rsize)
    fsize[1] = nx
    r = view(rr,(1:l for l in fsize)...)
    return new{T,N}(c,r,rr)
  end # function

  function (::Type{TransformableField}){T,N,P,I,L}(r::SubArray{T,N,P,I,L})
    rr = parent(r)
    fsize = [size(r)...]
    fsize[1] = div(fsize[1],2) + 1
    assert(fsize[1]*2 == size(rr)[1])
    c = reinterpret(Complex{T}, rr, (fsize...))
    return new{T,N}(c,r,rr)
  end # function

end # immutable

Base.size(S::TransformableField) = size(S.c)
Base.linearindexing{T<:TransformableField}(::Type{T}) = Base.LinearFast()
Base.getindex(S::TransformableField,I) = S.c[I]
Base.setindex!(S::TransformableField,v,I) =  setindex!(S.c,v,I)
Base.eltype(S::TransformableField) = typeof(S[1])

function TransformableField(t::DataType,ndims::Tuple)
  fsize = [ndims...]
  mod(fsize[1],2) == 0 ? fsize[1]+=2 : fsize[1]+=1
  a = Array{t,length(ndims)}((fsize...))
  TransformableField(a,ndims[1])
end

TransformableField(ndims::Tuple) = TransformableField(Float64,ndims)

###########################################################################################

function plan_rfft!{T<:Float3264,N}(X::TransformableField{T,N}, region;
                   flags::Integer=FFTW.ESTIMATE,
                   timelimit::Real=FFTW.NO_TIMELIMIT)# where N
  assert(flags&FFTW.ESTIMATE != 0)
  assert(1 in region)
  FFTW.rFFTWPlan{T,FFTW.FORWARD,true,N}(X.r, X.c, region, flags, timelimit)
end

plan_rfft!(f::TransformableField;kws...) = plan_rfft!(f,1:ndims(f.r);kws...)

*{T<:Float3264,N}(p::FFTW.rFFTWPlan{T,FFTW.FORWARD,true,N},f::TransformableField{T,N}) = (A_mul_B!(f.c,p,f.r);f)

rfft!(f::TransformableField, region=1:ndims(f.r)) = plan_rfft!(f,region) * f


##########################################################################################

function plan_brfft!{T,N}(X::TransformableField{T,N}, region;
                    flags::Integer=FFTW.PRESERVE_INPUT,
                    timelimit::Real=FFTW.NO_TIMELIMIT)
  assert(1 in region)
  FFTW.rFFTWPlan{Complex{T},FFTW.BACKWARD,true,N}(X.c, X.r, region, flags,timelimit)
end

function plan_irfft!{T,N}(x::TransformableField{T,N}, region; kws...)
  Base.DFT.ScaledPlan(plan_brfft!(x, region; kws...),Base.DFT.normalization(T, size(x.r), region))
end

plan_irfft!(f::TransformableField;kws...) = plan_irfft!(f,1:ndims(f.r);kws...)

*{T<:Float3264,N}(p::Base.DFT.ScaledPlan,f::TransformableField{T,N}) = begin
  A_mul_B!(f.r,p.p,f.c)
  scale!(f.rr,p.scale)
  f
end

irfft!(f::TransformableField, region=1:ndims(f.r)) = plan_irfft!(f,region) * f

end

The first direction must always be transformed.
Do you guys have any comments on the code? Any possibly dangerous thing? (I just spotted now that I didn’t enforce the PRESERVE_INPUT flag, which must be “on” for things to work properly).