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).