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

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/exports.jl

@@ -489,6 +489,7 @@ export
     cat,
     checkbounds,
     checkindex,
+    circcopy!,
     circshift,
     circshift!,
     clamp!,

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

doc/stdlib/arrays.rst

@@ -647,6 +647,33 @@ Indexing, Assignment, and Concatenation
 
    See also ``circshift``\ .
 
+.. function:: circcopy!(dest, src)
+
+   .. Docstring generated from Julia source
+
+   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``\ .
+
+   .. code-block:: 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:: find(A)
 
    .. Docstring generated from Julia source

test/TestHelpers.jl

@@ -43,4 +43,101 @@ function with_fake_pty(f)
     close(master)
 end
 
+# OffsetArrays (arrays with indexing that doesn't start at 1)
+
+# This test file is designed to exercise support for generic indexing,
+# even though offset arrays aren't implemented in Base.
+
+module OAs
+
+using Base: Indices, LinearSlow, LinearFast, tail
+
+export OffsetArray
+
+immutable OffsetArray{T,N,AA<:AbstractArray} <: AbstractArray{T,N}
+    parent::AA
+    offsets::NTuple{N,Int}
+end
+typealias OffsetVector{T,AA<:AbstractArray} OffsetArray{T,1,AA}
+
+OffsetArray{T,N}(A::AbstractArray{T,N}, offsets::NTuple{N,Int}) = OffsetArray{T,N,typeof(A)}(A, offsets)
+OffsetArray{T,N}(A::AbstractArray{T,N}, offsets::Vararg{Int,N}) = OffsetArray(A, offsets)
+
+(::Type{OffsetArray{T,N}}){T,N}(inds::Indices{N}) = OffsetArray{T,N,Array{T,N}}(Array{T,N}(map(length, inds)), map(indsoffset, inds))
+(::Type{OffsetArray{T}}){T,N}(inds::Indices{N}) = OffsetArray{T,N}(inds)
+
+Base.linearindexing{T<:OffsetArray}(::Type{T}) = Base.linearindexing(parenttype(T))
+parenttype{T,N,AA}(::Type{OffsetArray{T,N,AA}}) = AA
+parenttype(A::OffsetArray) = parenttype(typeof(A))
+
+Base.parent(A::OffsetArray) = A.parent
+
+errmsg(A) = error("size not supported for arrays with indices $(indices(A)); see http://docs.julialang.org/en/latest/devdocs/offset-arrays/")
+Base.size(A::OffsetArray) = errmsg(A)
+Base.size(A::OffsetArray, d) = errmsg(A)
+Base.eachindex(::LinearSlow, A::OffsetArray) = CartesianRange(indices(A))
+Base.eachindex(::LinearFast, A::OffsetVector) = indices(A, 1)
+
+# Implementations of indices and indices1. Since bounds-checking is
+# performance-critical and relies on indices, these are usually worth
+# optimizing thoroughly.
+@inline Base.indices(A::OffsetArray, d) = 1 <= d <= length(A.offsets) ? indices(parent(A))[d] + A.offsets[d] : (1:1)
+@inline Base.indices(A::OffsetArray) = _indices(indices(parent(A)), A.offsets)  # would rather use ntuple, but see #15276
+@inline _indices(inds, offsets) = (inds[1]+offsets[1], _indices(tail(inds), tail(offsets))...)
+_indices(::Tuple{}, ::Tuple{}) = ()
+Base.indices1{T}(A::OffsetArray{T,0}) = 1:1  # we only need to specialize this one
+
+function Base.similar(A::OffsetArray, T::Type, dims::Dims)
+    B = similar(parent(A), T, dims)
+end
+function Base.similar(A::AbstractArray, T::Type, inds::Tuple{UnitRange,Vararg{UnitRange}})
+    B = similar(A, T, map(length, inds))
+    OffsetArray(B, map(indsoffset, inds))
+end
+
+Base.similar(f::Union{Function,DataType}, shape::Tuple{UnitRange,Vararg{UnitRange}}) = OffsetArray(f(map(length, shape)), map(indsoffset, shape))
+
+Base.reshape(A::AbstractArray, inds::Tuple{UnitRange,Vararg{UnitRange}}) = OffsetArray(reshape(A, map(length, inds)), map(indsoffset, inds))
+
+@inline function Base.getindex{T,N}(A::OffsetArray{T,N}, I::Vararg{Int,N})
+    checkbounds(A, I...)
+    @inbounds ret = parent(A)[offset(A.offsets, I)...]
+    ret
+end
+@inline function Base._getindex(::LinearFast, A::OffsetVector, i::Int)
+    checkbounds(A, i)
+    @inbounds ret = parent(A)[offset(A.offsets, (i,))[1]]
+    ret
+end
+@inline function Base._getindex(::LinearFast, A::OffsetArray, i::Int)
+    checkbounds(A, i)
+    @inbounds ret = parent(A)[i]
+    ret
+end
+@inline function Base.setindex!{T,N}(A::OffsetArray{T,N}, val, I::Vararg{Int,N})
+    checkbounds(A, I...)
+    @inbounds parent(A)[offset(A.offsets, I)...] = val
+    val
+end
+@inline function Base._setindex!(::LinearFast, A::OffsetVector, val, i::Int)
+    checkbounds(A, i)
+    @inbounds parent(A)[offset(A.offsets, (i,))[1]] = val
+    val
+end
+@inline function Base._setindex!(::LinearFast, A::OffsetArray, val, i::Int)
+    checkbounds(A, i)
+    @inbounds parent(A)[i] = val
+    val
+end
+
+# Computing a shifted index (subtracting the offset)
+offset{N}(offsets::NTuple{N,Int}, inds::NTuple{N,Int}) = _offset((), offsets, inds)
+_offset(out, ::Tuple{}, ::Tuple{}) = out
+@inline _offset(out, offsets, inds) = _offset((out..., inds[1]-offsets[1]), Base.tail(offsets), Base.tail(inds))
+
+indsoffset(r::Range) = first(r) - 1
+indsoffset(i::Integer) = 0
+
+end
+
 end

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