Skip to content

Commit

Permalink
Merge pull request #17919 from JuliaLang/teh/fft
Browse files Browse the repository at this point in the history
Support non-1 indices and fix type problems in DFT (fixes #17896)
  • Loading branch information
timholy authored Aug 11, 2016
2 parents 8cc4a9b + e95b5e2 commit b066714
Show file tree
Hide file tree
Showing 7 changed files with 279 additions and 114 deletions.
39 changes: 29 additions & 10 deletions base/dft.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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!
Expand Down
1 change: 1 addition & 0 deletions base/exports.jl
Original file line number Diff line number Diff line change
Expand Up @@ -489,6 +489,7 @@ export
cat,
checkbounds,
checkindex,
circcopy!,
circshift,
circshift!,
clamp!,
Expand Down
78 changes: 78 additions & 0 deletions base/multidimensional.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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}})
Expand All @@ -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
Expand Down
27 changes: 27 additions & 0 deletions doc/stdlib/arrays.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
97 changes: 97 additions & 0 deletions test/TestHelpers.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
8 changes: 8 additions & 0 deletions test/fft.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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;]))
Loading

0 comments on commit b066714

Please sign in to comment.