Support non-1 indices and fix type problems in DFT (fixes #17896)

(cherry picked from commit 996e275)
ref #17919

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!

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

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;]))

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