Merge pull request #17816 from JuliaLang/teh/indices_iters

More indices generalizations

Author: timholy

base/abstractarray.jl

@@ -625,6 +625,13 @@ convert{T,S,N}(::Type{AbstractArray{T  }}, A::AbstractArray{S,N}) = convert(Abst
 
 convert{T,N}(::Type{Array}, A::AbstractArray{T,N}) = convert(Array{T,N}, A)
 
+"""
+   of_indices(x, y)
+
+Represents the array `y` as an array having the same indices type as `x`.
+"""
+of_indices(x, y) = similar(dims->y, oftype(indices(x), indices(y)))
+
 full(x::AbstractArray) = x
 
 map(::Type{Integer},  a::Array) = map!(Integer, similar(a,typeof(Integer(one(eltype(a))))), a)

base/array.jl

@@ -211,7 +211,7 @@ The result has the same shape and number of dimensions as `collection`.
 collect{T}(::Type{T}, itr) = _collect(T, itr, iteratorsize(itr))
 
 _collect{T}(::Type{T}, itr, isz::HasLength) = copy!(Array{T,1}(Int(length(itr)::Integer)), itr)
-_collect{T}(::Type{T}, itr, isz::HasShape)  = copy!(Array{T}(convert(Dims,size(itr))), itr)
+_collect{T}(::Type{T}, itr, isz::HasShape)  = copy!(similar(Array{T}, indices(itr)), itr)
 function _collect{T}(::Type{T}, itr, isz::SizeUnknown)
     a = Array{T,1}(0)
     for x in itr
@@ -259,7 +259,7 @@ end
 _default_eltype{I,T}(::Type{Generator{I,Type{T}}}) = T
 
 _array_for(T, itr, ::HasLength) = Array{T,1}(Int(length(itr)::Integer))
-_array_for(T, itr, ::HasShape) = Array{T}(convert(Dims,size(itr)))
+_array_for(T, itr, ::HasShape) = similar(Array{T}, indices(itr))
 
 function collect(itr::Generator)
     isz = iteratorsize(itr.iter)
@@ -795,7 +795,7 @@ function findn(A::AbstractMatrix)
     I = similar(A, Int, nnzA)
     J = similar(A, Int, nnzA)
     count = 1
-    for j=1:size(A,2), i=1:size(A,1)
+    for j=indices(A,2), i=indices(A,1)
         if A[i,j] != 0
             I[count] = i
             J[count] = j
@@ -812,7 +812,7 @@ function findnz{T}(A::AbstractMatrix{T})
     NZs = Array{T,1}(nnzA)
     count = 1
     if nnzA > 0
-        for j=1:size(A,2), i=1:size(A,1)
+        for j=indices(A,2), i=indices(A,1)
             Aij = A[i,j]
             if Aij != 0
                 I[count] = i

base/arraymath.jl

@@ -413,8 +413,8 @@ function ctranspose(A::AbstractMatrix)
 end
 ctranspose{T<:Real}(A::AbstractVecOrMat{T}) = transpose(A)
 
-transpose(x::AbstractVector) = [ transpose(v) for i=1:1, v in x ]
-ctranspose{T}(x::AbstractVector{T}) = T[ ctranspose(v) for i=1:1, v in x ]
+transpose(x::AbstractVector) = [ transpose(v) for i=of_indices(x, OneTo(1)), v in x ]
+ctranspose{T}(x::AbstractVector{T}) = T[ ctranspose(v) for i=of_indices(x, OneTo(1)), v in x ]
 
 _cumsum_type{T<:Number}(v::AbstractArray{T}) = typeof(+zero(T))
 _cumsum_type(v) = typeof(v[1]+v[1])

base/iterator.jl

@@ -69,6 +69,7 @@ end
 zip(a) = Zip1(a)
 length(z::Zip1) = length(z.a)
 size(z::Zip1) = size(z.a)
+indices(z::Zip1) = indices(z.a)
 eltype{I}(::Type{Zip1{I}}) = Tuple{eltype(I)}
 @inline start(z::Zip1) = start(z.a)
 @inline function next(z::Zip1, st)
@@ -87,6 +88,7 @@ end
 zip(a, b) = Zip2(a, b)
 length(z::Zip2) = _min_length(z.a, z.b, iteratorsize(z.a), iteratorsize(z.b))
 size(z::Zip2) = promote_shape(size(z.a), size(z.b))
+indices(z::Zip2) = promote_shape(indices(z.a), indices(z.b))
 eltype{I1,I2}(::Type{Zip2{I1,I2}}) = Tuple{eltype(I1), eltype(I2)}
 @inline start(z::Zip2) = (start(z.a), start(z.b))
 @inline function next(z::Zip2, st)
@@ -137,6 +139,7 @@ julia> first(c)
 zip(a, b, c...) = Zip(a, zip(b, c...))
 length(z::Zip) = _min_length(z.a, z.z, iteratorsize(z.a), iteratorsize(z.z))
 size(z::Zip) = promote_shape(size(z.a), size(z.z))
+indices(z::Zip) = promote_shape(indices(z.a), indices(z.z))
 tuple_type_cons{S}(::Type{S}, ::Type{Union{}}) = Union{}
 function tuple_type_cons{S,T<:Tuple}(::Type{S}, ::Type{T})
     @_pure_meta
@@ -430,8 +433,10 @@ iteratoreltype{O}(::Type{Repeated{O}}) = HasEltype()
 abstract AbstractProdIterator
 
 length(p::AbstractProdIterator) = prod(size(p))
+_length(p::AbstractProdIterator) = prod(map(unsafe_length, indices(p)))
 size(p::AbstractProdIterator) = _prod_size(p.a, p.b, iteratorsize(p.a), iteratorsize(p.b))
-ndims(p::AbstractProdIterator) = length(size(p))
+indices(p::AbstractProdIterator) = _prod_indices(p.a, p.b, iteratorsize(p.a), iteratorsize(p.b))
+ndims(p::AbstractProdIterator) = length(indices(p))
 
 # generic methods to handle size of Prod* types
 _prod_size(a, ::HasShape)  = size(a)
@@ -445,6 +450,17 @@ _prod_size(a, b, ::HasShape,  ::HasShape)   = (size(a)..., size(b)...)
 _prod_size(a, b, A, B) =
     throw(ArgumentError("Cannot construct size for objects of types $(typeof(a)) and $(typeof(b))"))
 
+_prod_indices(a, ::HasShape)  = indices(a)
+_prod_indices(a, ::HasLength) = (OneTo(length(a)), )
+_prod_indices(a, A) =
+    throw(ArgumentError("Cannot compute indices for object of type $(typeof(a))"))
+_prod_indices(a, b, ::HasLength, ::HasLength)  = (OneTo(length(a)),  OneTo(length(b)))
+_prod_indices(a, b, ::HasLength, ::HasShape)   = (OneTo(length(a)),  indices(b)...)
+_prod_indices(a, b, ::HasShape,  ::HasLength)  = (indices(a)..., OneTo(length(b)))
+_prod_indices(a, b, ::HasShape,  ::HasShape)   = (indices(a)..., indices(b)...)
+_prod_indices(a, b, A, B) =
+    throw(ArgumentError("Cannot construct indices for objects of types $(typeof(a)) and $(typeof(b))"))
+
 # one iterator
 immutable Prod1{I} <: AbstractProdIterator
     a::I
@@ -453,6 +469,7 @@ product(a) = Prod1(a)
 
 eltype{I}(::Type{Prod1{I}}) = Tuple{eltype(I)}
 size(p::Prod1) = _prod_size(p.a, iteratorsize(p.a))
+indices(p::Prod1) = _prod_indices(p.a, iteratorsize(p.a))
 
 @inline start(p::Prod1) = start(p.a)
 @inline function next(p::Prod1, st)

base/test.jl

@@ -827,10 +827,11 @@ approx_full(x) = full(x)
 function test_approx_eq(va, vb, Eps, astr, bstr)
     va = approx_full(va)
     vb = approx_full(vb)
-    if length(va) != length(vb)
+    la, lb = length(linearindices(va)), length(linearindices(vb))
+    if la != lb
         error("lengths of ", astr, " and ", bstr, " do not match: ",
-              "\n  ", astr, " (length $(length(va))) = ", va,
-              "\n  ", bstr, " (length $(length(vb))) = ", vb)
+              "\n  ", astr, " (length $la) = ", va,
+              "\n  ", bstr, " (length $lb) = ", vb)
     end
     diff = real(zero(eltype(va)))
     for (xa, xb) = zip(va, vb)
@@ -856,7 +857,7 @@ array_eps{T}(a::AbstractArray{Complex{T}}) = eps(float(maximum(x->(isfinite(x) ?
 array_eps(a) = eps(float(maximum(x->(isfinite(x) ? abs(x) : oftype(x,NaN)), a)))
 
 test_approx_eq(va, vb, astr, bstr) =
-    test_approx_eq(va, vb, 1E4*length(va)*max(array_eps(va), array_eps(vb)), astr, bstr)
+    test_approx_eq(va, vb, 1E4*length(linearindices(va))*max(array_eps(va), array_eps(vb)), astr, bstr)
 
 """
     @test_approx_eq_eps(a, b, tol)
@@ -966,10 +967,10 @@ end
 # nothing.
 function test_approx_eq_modphase{S<:Real,T<:Real}(
         a::StridedVecOrMat{S}, b::StridedVecOrMat{T}, err=nothing)
-    m, n = size(a)
-    @test n==size(b, 2) && m==size(b, 1)
+    @test indices(a,1) == indices(b,1) && indices(a,2) == indices(b,2)
+    m = length(indices(a,1))
     err === nothing && (err=m^3*(eps(S)+eps(T)))
-    for i=1:n
+    for i in indices(a,2)
         v1, v2 = a[:, i], b[:, i]
         @test_approx_eq_eps min(abs(norm(v1-v2)), abs(norm(v1+v2))) 0.0 err
     end

test/offsetarray.jl

@@ -101,6 +101,12 @@ using OAs
 
 let
 # Basics
+v0 = rand(4)
+v = OffsetArray(v0, (-3,))
+@test indices(v) == (-2:1,)
+@test_throws ErrorException size(v)
+@test_throws ErrorException size(v, 1)
+
 A0 = [1 3; 2 4]
 A = OffsetArray(A0, (-1,2))                   # LinearFast
 S = OffsetArray(view(A0, 1:2, 1:2), (-1,2))   # LinearSlow
@@ -179,6 +185,10 @@ end
 
 # show
 io = IOBuffer()
+show(io, v)
+str = takebuf_string(io)
+show(io, v0)
+@test str == takebuf_string(io)
 show(io, A)
 str = takebuf_string(io)
 @test str == "[1 3; 2 4]"
@@ -261,7 +271,7 @@ v = view(A0, 1:1, i1)
 # logical indexing
 @test A[A .> 2] == [3,4]
 
-# copy!
+# copy! and fill!
 a = OffsetArray{Int}((-3:-1,))
 fill!(a, -1)
 copy!(a, (1,2))   # non-array iterables
@@ -316,6 +326,7 @@ copy!(am, b)
 @test am[1,8] == 2
 @test am[1,9] == -1
 
+# map
 dest = similar(am)
 map!(+, dest, am, am)
 @test dest[1,7] == 2
@@ -326,7 +337,13 @@ am = map(identity, a)
 @test isa(am, OffsetArray)
 @test am == a
 
+# other functions
+v = OffsetArray(v0, (-3,))
+@test_approx_eq v v
+@test parent(v') == v0'
+@test indices(v') === (1:1,-2:1)
 A = OffsetArray(rand(4,4), (-3,5))
+@test_approx_eq A A
 @test maximum(A) == maximum(parent(A))
 @test minimum(A) == minimum(parent(A))
 @test extrema(A) == extrema(parent(A))
@@ -352,6 +369,14 @@ pmax, ipmax = findmax(parent(A))
 @test amax == pmax
 @test A[iamax] == amax
 @test amax == parent(A)[ipmax]
+z = OffsetArray([0 0; 2 0; 0 0; 0 0], (-3,-1))
+I,J = findn(z)
+@test I == [-1]
+@test J == [0]
+I,J,N = findnz(z)
+@test I == [-1]
+@test J == [0]
+@test N == [2]
 
 v  = OffsetArray([1,1e100,1,-1e100], (-3,))*1000
 v2 = OffsetArray([1,-1e100,1,1e100], (5,))*1000