From 05197a51fa015dfcbec05ddba9f2e2e2a2945150 Mon Sep 17 00:00:00 2001 From: Tim Holy Date: Tue, 9 Aug 2016 17:47:48 -0500 Subject: [PATCH] Support non-1 indices and fix type problems in DFT (fixes #17896) (cherry picked from commit 996e2757cb39b88f112e03f602339aa2b7f0ee65) ref #17919 --- base/dft.jl | 39 ++++++++++++++------ base/multidimensional.jl | 78 ++++++++++++++++++++++++++++++++++++++++ test/fft.jl | 8 +++++ test/offsetarray.jl | 33 ++++++++++++++++- 4 files changed, 147 insertions(+), 11 deletions(-) diff --git a/base/dft.jl b/base/dft.jl index 04b9c5bf7875b..700bbc68c2dec 100644 --- a/base/dft.jl +++ b/base/dft.jl @@ -20,22 +20,41 @@ export fft, ifft, bfft, fft!, ifft!, bfft!, plan_fft, plan_ifft, plan_bfft, plan_fft!, plan_ifft!, plan_bfft!, rfft, irfft, brfft, plan_rfft, plan_irfft, plan_brfft -complexfloat{T<:AbstractFloat}(x::AbstractArray{Complex{T}}) = x +typealias FFTWFloat Union{Float32,Float64} +fftwfloat(x) = _fftwfloat(float(x)) +_fftwfloat{T<:FFTWFloat}(::Type{T}) = T +_fftwfloat(::Type{Float16}) = Float32 +_fftwfloat{T}(::Type{T}) = error("type $T not supported") +_fftwfloat{T}(x::T) = _fftwfloat(T)(x) + +complexfloat{T<:FFTWFloat}(x::StridedArray{Complex{T}}) = x +realfloat{T<:FFTWFloat}(x::StridedArray{T}) = x # return an Array, rather than similar(x), to avoid an extra copy for FFTW # (which only works on StridedArray types). -complexfloat{T<:Complex}(x::AbstractArray{T}) = copy!(Array{typeof(float(one(T)))}(size(x)), x) -complexfloat{T<:AbstractFloat}(x::AbstractArray{T}) = copy!(Array{typeof(complex(one(T)))}(size(x)), x) -complexfloat{T<:Real}(x::AbstractArray{T}) = copy!(Array{typeof(complex(float(one(T))))}(size(x)), x) +complexfloat{T<:Complex}(x::AbstractArray{T}) = copy1(typeof(fftwfloat(one(T))), x) +complexfloat{T<:Real}(x::AbstractArray{T}) = copy1(typeof(complex(fftwfloat(one(T)))), x) + +realfloat{T<:Real}(x::AbstractArray{T}) = copy1(typeof(fftwfloat(one(T))), x) + +# copy to a 1-based array, using circular permutation +function copy1{T}(::Type{T}, x) + y = Array{T}(map(length, indices(x))) + Base.circcopy!(y, x) +end + +to1(x::AbstractArray) = _to1(indices(x), x) +_to1(::Tuple{Base.OneTo,Vararg{Base.OneTo}}, x) = x +_to1(::Tuple, x) = copy1(eltype(x), x) # implementations only need to provide plan_X(x, region) # for X in (:fft, :bfft, ...): for f in (:fft, :bfft, :ifft, :fft!, :bfft!, :ifft!, :rfft) pf = Symbol("plan_", f) @eval begin - $f(x::AbstractArray) = $pf(x) * x - $f(x::AbstractArray, region) = $pf(x, region) * x - $pf(x::AbstractArray; kws...) = $pf(x, 1:ndims(x); kws...) + $f(x::AbstractArray) = (y = to1(x); $pf(y) * y) + $f(x::AbstractArray, region) = (y = to1(x); $pf(y, region) * y) + $pf(x::AbstractArray; kws...) = (y = to1(x); $pf(y, 1:ndims(y); kws...)) end end @@ -187,11 +206,11 @@ for f in (:fft, :bfft, :ifft) $pf{T<:Union{Integer,Rational}}(x::AbstractArray{Complex{T}}, region; kws...) = $pf(complexfloat(x), region; kws...) end end -rfft{T<:Union{Integer,Rational}}(x::AbstractArray{T}, region=1:ndims(x)) = rfft(float(x), region) -plan_rfft{T<:Union{Integer,Rational}}(x::AbstractArray{T}, region; kws...) = plan_rfft(float(x), region; kws...) +rfft{T<:Union{Integer,Rational}}(x::AbstractArray{T}, region=1:ndims(x)) = rfft(realfloat(x), region) +plan_rfft(x::AbstractArray, region; kws...) = plan_rfft(realfloat(x), region; kws...) # only require implementation to provide *(::Plan{T}, ::Array{T}) -*{T}(p::Plan{T}, x::AbstractArray) = p * copy!(Array{T}(size(x)), x) +*{T}(p::Plan{T}, x::AbstractArray) = p * copy1(T, x) # Implementations should also implement A_mul_B!(Y, plan, X) so as to support # pre-allocated output arrays. We don't define * in terms of A_mul_B! diff --git a/base/multidimensional.jl b/base/multidimensional.jl index a6e22e481601b..3773db8701964 100644 --- a/base/multidimensional.jl +++ b/base/multidimensional.jl @@ -645,6 +645,22 @@ See also `circshift`. end circshift!(dest::AbstractArray, src, shiftamt) = circshift!(dest, src, (shiftamt...,)) +# For each dimension, we copy the first half of src to the second half +# of dest, and the second half of src to the first half of dest. This +# uses a recursive bifurcation strategy so that these splits can be +# encoded by ranges, which means that we need only one call to `mod` +# per dimension rather than one call per index. +# `rdest` and `rsrc` are tuples-of-ranges that grow one dimension at a +# time; when all the dimensions have been filled in, you call `copy!` +# for that block. In other words, in two dimensions schematically we +# have the following call sequence (--> means a call): +# circshift!(dest, src, shiftamt) --> +# _circshift!(dest, src, ("first half of dim1",)) --> +# _circshift!(dest, src, ("first half of dim1", "first half of dim2")) --> copy! +# _circshift!(dest, src, ("first half of dim1", "second half of dim2")) --> copy! +# _circshift!(dest, src, ("second half of dim1",)) --> +# _circshift!(dest, src, ("second half of dim1", "first half of dim2")) --> copy! +# _circshift!(dest, src, ("second half of dim1", "second half of dim2")) --> copy! @inline function _circshift!(dest, rdest, src, rsrc, inds::Tuple{AbstractUnitRange,Vararg{Any}}, shiftamt::Tuple{Integer,Vararg{Any}}) @@ -662,6 +678,68 @@ function _circshift!(dest, rdest, src, rsrc, inds, shiftamt) copy!(dest, CartesianRange(rdest), src, CartesianRange(rsrc)) end +# circcopy! +""" + circcopy!(dest, src) + +Copy `src` to `dest`, indexing each dimension modulo its length. +`src` and `dest` must have the same size, but can be offset in +their indices; any offset results in a (circular) wraparound. If the +arrays have overlapping indices, then on the domain of the overlap +`dest` agrees with `src`. + +```julia +julia> src = reshape(collect(1:16), (4,4)) +4×4 Array{Int64,2}: + 1 5 9 13 + 2 6 10 14 + 3 7 11 15 + 4 8 12 16 + +julia> dest = OffsetArray{Int}((0:3,2:5)) + +julia> circcopy!(dest, src) +OffsetArrays.OffsetArray{Int64,2,Array{Int64,2}} with indices 0:3×2:5: + 8 12 16 4 + 5 9 13 1 + 6 10 14 2 + 7 11 15 3 + +julia> dest[1:3,2:4] == src[1:3,2:4] +true +``` +""" +function circcopy!(dest, src) + dest === src && throw(ArgumentError("dest and src must be separate arrays")) + indssrc, indsdest = indices(src), indices(dest) + if (szsrc = map(length, indssrc)) != (szdest = map(length, indsdest)) + throw(DimensionMismatch("src and dest must have the same sizes (got $szsrc and $szdest)")) + end + shift = map((isrc, idest)->first(isrc)-first(idest), indssrc, indsdest) + all(x->x==0, shift) && return copy!(dest, src) + _circcopy!(dest, (), indsdest, src, (), indssrc) +end + +# This uses the same strategy described above for _circshift! +@inline function _circcopy!(dest, rdest, indsdest::Tuple{AbstractUnitRange,Vararg{Any}}, + src, rsrc, indssrc::Tuple{AbstractUnitRange,Vararg{Any}}) + indd1, inds1 = indsdest[1], indssrc[1] + l = length(indd1) + s = mod(first(inds1)-first(indd1), l) + sdf = first(indd1)+s + rd1, rd2 = first(indd1):sdf-1, sdf:last(indd1) + ssf = last(inds1)-s + rs1, rs2 = first(inds1):ssf, ssf+1:last(inds1) + tindsd, tindss = tail(indsdest), tail(indssrc) + _circcopy!(dest, (rdest..., rd1), tindsd, src, (rsrc..., rs2), tindss) + _circcopy!(dest, (rdest..., rd2), tindsd, src, (rsrc..., rs1), tindss) +end + +# At least one of indsdest, indssrc are empty (and both should be, since we've checked) +function _circcopy!(dest, rdest, indsdest, src, rsrc, indssrc) + copy!(dest, CartesianRange(rdest), src, CartesianRange(rsrc)) +end + ### BitArrays ## getindex diff --git a/test/fft.jl b/test/fft.jl index 6905bc85ffb30..b30e05ff09a5a 100644 --- a/test/fft.jl +++ b/test/fft.jl @@ -326,3 +326,11 @@ for x in (randn(10),randn(10,12)) # note: inference doesn't work for plan_fft_ since the # algorithm steps are included in the CTPlan type end + +# issue #17896 +a = rand(5) +@test fft(a) == fft(view(a,:)) == fft(view(a, 1:5)) == fft(view(a, [1:5;])) +@test rfft(a) == rfft(view(a,:)) == rfft(view(a, 1:5)) == rfft(view(a, [1:5;])) +a16 = convert(Vector{Float16}, a) +@test fft(a16) == fft(view(a16,:)) == fft(view(a16, 1:5)) == fft(view(a16, [1:5;])) +@test rfft(a16) == rfft(view(a16,:)) == rfft(view(a16, 1:5)) == rfft(view(a16, [1:5;])) diff --git a/test/offsetarray.jl b/test/offsetarray.jl index a509e17d3a8fb..dac5aad7995ca 100644 --- a/test/offsetarray.jl +++ b/test/offsetarray.jl @@ -411,10 +411,41 @@ v = OffsetArray(rand(8), (-2,)) @test rotr90(A) == OffsetArray(rotr90(parent(A)), A.offsets[[2,1]]) @test flipdim(A, 1) == OffsetArray(flipdim(parent(A), 1), A.offsets) @test flipdim(A, 2) == OffsetArray(flipdim(parent(A), 2), A.offsets) -@test circshift(A, (-1,2)) == OffsetArray(circshift(parent(A), (-1,2)), A.offsets) @test A+1 == OffsetArray(parent(A)+1, A.offsets) @test 2*A == OffsetArray(2*parent(A), A.offsets) @test A+A == OffsetArray(parent(A)+parent(A), A.offsets) @test A.*A == OffsetArray(parent(A).*parent(A), A.offsets) + +@test circshift(A, (-1,2)) == OffsetArray(circshift(parent(A), (-1,2)), A.offsets) + +src = reshape(collect(1:16), (4,4)) +dest = OffsetArray(Array{Int}(4,4), (-1,1)) +circcopy!(dest, src) +@test parent(dest) == [8 12 16 4; 5 9 13 1; 6 10 14 2; 7 11 15 3] +@test dest[1:3,2:4] == src[1:3,2:4] + +e = eye(5) +a = [e[:,1], e[:,2], e[:,3], e[:,4], e[:,5]] +a1 = zeros(5) +c = [ones(Complex{Float64}, 5), + exp(-2*pi*im*(0:4)/5), + exp(-4*pi*im*(0:4)/5), + exp(-6*pi*im*(0:4)/5), + exp(-8*pi*im*(0:4)/5)] +for s = -5:5 + for i = 1:5 + thisa = OffsetArray(a[i], (s,)) + thisc = c[mod1(i+s+5,5)] + @test_approx_eq fft(thisa) thisc + @test_approx_eq fft(thisa, 1) thisc + @test_approx_eq ifft(fft(thisa)) circcopy!(a1, thisa) + @test_approx_eq ifft(fft(thisa, 1), 1) circcopy!(a1, thisa) + @test_approx_eq rfft(thisa) thisc[1:3] + @test_approx_eq rfft(thisa, 1) thisc[1:3] + @test_approx_eq irfft(rfft(thisa, 1), 5, 1) a1 + @test_approx_eq irfft(rfft(thisa, 1), 5, 1) a1 + end end + +end # let