Merge pull request #10337 from Jutho/jh/stagedsub2ind

improved ind2sub/sub2ind

Author: timholy

base/abstractarray.jl

@@ -1009,97 +1009,81 @@ function repmat(a::AbstractVector, m::Int)
     return b
 end
 
-sub2ind(dims) = 1
-sub2ind(dims, i::Integer) = Int(i)
-sub2ind(dims, i::Integer, j::Integer) = sub2ind(dims, Int(i), Int(j))
-sub2ind(dims, i::Int, j::Int) = (j-1)*dims[1] + i
-sub2ind(dims, i0::Integer, i1::Integer, i2::Integer) = sub2ind(dims, Int(i0),Int(i1),Int(i2))
-sub2ind(dims, i0::Int, i1::Int, i2::Int) =
-    i0 + dims[1]*((i1-1) + dims[2]*(i2-1))
-sub2ind(dims, i0::Integer, i1::Integer, i2::Integer, i3::Integer) =
-    sub2ind(dims, Int(i0),Int(i1),Int(i2),Int(i3))
-sub2ind(dims, i0::Int, i1::Int, i2::Int, i3::Int) =
-    i0 + dims[1]*((i1-1) + dims[2]*((i2-1) + dims[3]*(i3-1)))
-
-function sub2ind(dims, I::Integer...)
-    ndims = length(dims)
-    index = Int(I[1])
-    stride = 1
-    for k=2:ndims
-        stride = stride * dims[k-1]
-        index += (Int(I[k])-1) * stride
+sub2ind(dims::Tuple{Vararg{Integer}}) = 1
+sub2ind(dims::Tuple{Vararg{Integer}}, I::Integer...) = _sub2ind(dims,I)
+@generated function _sub2ind{N,M}(dims::NTuple{N,Integer}, I::NTuple{M,Integer})
+    meta = Expr(:meta,:inline)
+    ex = :(I[$M] - 1)
+    for i = M-1:-1:1
+        if i > N
+            ex = :(I[$i] - 1 + $ex)
+        else
+            ex = :(I[$i] - 1 + dims[$i]*$ex)
+        end
     end
-    return index
+    Expr(:block, meta,:($ex + 1))
 end
 
-function sub2ind{T<:Integer}(dims::Array{T}, sub::Array{T})
-    ndims = length(dims)
-    ind = sub[1]
-    stride = 1
-    for k in 2:ndims
-        stride = stride * dims[k - 1]
-        ind += (sub[k] - 1) * stride
+@generated function ind2sub{N}(dims::NTuple{N,Integer}, ind::Integer)
+    meta = Expr(:meta,:inline)
+    N==0 && return :($meta; ind==1 ? () : throw(BoundsError()))
+    exprs = Expr[:(ind = ind-1)]
+    for i = 1:N-1
+        push!(exprs,:(ind2 = div(ind,dims[$i])))
+        push!(exprs,Expr(:(=),symbol(:s,i),:(ind-dims[$i]*ind2+1)))
+        push!(exprs,:(ind=ind2))
     end
-    return ind
+    push!(exprs,Expr(:(=),symbol(:s,N),:(ind+1)))
+    Expr(:block,meta,exprs...,Expr(:tuple,[symbol(:s,i) for i=1:N]...))
 end
 
-sub2ind{T<:Integer}(dims, I::AbstractVector{T}...) =
-    [ sub2ind(dims, map(X->X[i], I)...)::Int for i=1:length(I[1]) ]
+# TODO in v0.5: either deprecate line 1 or add line 2
+ind2sub(a::AbstractArray, ind::Integer) = ind2sub(size(a), ind)
+# sub2ind(a::AbstractArray, I::Integer...) = sub2ind(size(a), I...)
 
-function ind2sub(dims::Tuple{Integer,Vararg{Integer}}, ind::Int)
-    ndims = length(dims)
-    stride = dims[1]
-    for i=2:ndims-1
-        stride *= dims[i]
-    end
-
-    sub = ()
-    for i=(ndims-1):-1:1
-        rest = rem(ind-1, stride) + 1
-        sub = tuple(div(ind - rest, stride) + 1, sub...)
-        ind = rest
-        stride = div(stride, dims[i])
-    end
-    return tuple(ind, sub...)
-end
-
-ind2sub(dims::Tuple{Vararg{Integer}}, ind::Integer) = ind2sub(dims, Int(ind))
-ind2sub(dims::Tuple{}, ind::Integer) = ind==1 ? () : throw(BoundsError())
-ind2sub(dims::Tuple{Integer,}, ind::Int) = (ind,)
-ind2sub(dims::Tuple{Integer,Integer}, ind::Int) =
-    (rem(ind-1,dims[1])+1, div(ind-1,dims[1])+1)
-ind2sub(dims::Tuple{Integer,Integer,Integer}, ind::Int) =
-    (rem(ind-1,dims[1])+1, div(rem(ind-1,dims[1]*dims[2]), dims[1])+1,
-     div(rem(ind-1,dims[1]*dims[2]*dims[3]), dims[1]*dims[2])+1)
-ind2sub(a::AbstractArray, ind::Integer) = ind2sub(size(a), Int(ind))
-
-function ind2sub{T<:Integer}(dims::Tuple{Integer,Vararg{Integer}}, ind::AbstractVector{T})
-    n = length(dims)
-    l = length(ind)
-    t = ntuple(n, x->Array(Int, l))
-    for i = 1:l
-        s = ind2sub(dims, ind[i])
-        for j = 1:n
-            t[j][i] = s[j]
+function sub2ind{T<:Integer}(dims::Tuple{Vararg{Integer}}, I::AbstractVector{T}...)
+    N = length(dims)
+    M = length(I[1])
+    indices = Array{T}(length(I[1]))
+    copy!(indices,I[1])
+
+    s = dims[1]
+    for j=2:length(I)
+        Ij = I[j]
+        for i=1:M
+            indices[i] += s*(Ij[i]-1)
+        end
+        s*= (j <= N ? dims[j] : 1)
+    end
+    return indices
+end
+
+function ind2sub{N,T<:Integer}(dims::NTuple{N,Integer}, ind::AbstractVector{T})
+    M = length(ind)
+    t = NTuple{N,Vector{T}}(ntuple(N,n->Array{T}(M)))
+    copy!(t[1],ind)
+    for j = 1:N-1
+        d = dims[j]
+        tj = t[j]
+        tj2 = t[j+1]
+        for i = 1:M
+            ind2 = div(tj[i]-1, d)
+            tj[i] -= d*ind2
+            tj2[i] = ind2+1
         end
     end
     return t
 end
 
-function ind2sub!{T<:Integer}(sub::Array{T}, dims::Array{T}, ind::T)
+function ind2sub!{T<:Integer}(sub::Array{T}, dims::Tuple{Vararg{T}}, ind::T)
     ndims = length(dims)
-    stride = dims[1]
-    for i in 2:(ndims - 1)
-        stride *= dims[i]
-    end
-    for i in (ndims - 1):-1:1
-        rest = rem1(ind, stride)
-        sub[i + 1] = div(ind - rest, stride) + 1
-        ind = rest
-        stride = div(stride, dims[i])
+    for i=1:ndims-1
+        ind2 = div(ind-1,dims[i])+1
+        sub[i] = ind - dims[i]*(ind2-1)
+        ind = ind2
     end
-    sub[1] = ind
-    return
+    sub[ndims] = ind
+    return sub
 end
 
 # Generalized repmat
@@ -1115,26 +1099,18 @@ function repeat{T}(A::Array{T};
         throw(ArgumentError("inner/outer repetitions must be set for all input dimensions"))
     end
 
-    size_in = Array(Int, ndims_in)
-    size_out = Array(Int, ndims_out)
-    inner_size_out = Array(Int, ndims_out)
+    inner = vcat(inner, ones(Int,ndims_out-length_inner))
+    outer = vcat(outer, ones(Int,ndims_out-length_outer))
 
-    for i in 1:ndims_in
-        size_in[i] = size(A, i)
-    end
-    for i in 1:ndims_out
-        t1 = ndims_in < i ? 1 : size_in[i]
-        t2 = length_inner < i ? 1 : inner[i]
-        t3 = length_outer < i ? 1 : outer[i]
-        size_out[i] = t1 * t2 * t3
-        inner_size_out[i] = t1 * t2
-    end
+    size_in = size(A)
+    size_out = ntuple(i->inner[i]*size(A,i)*outer[i],ndims_out)::Dims
+    inner_size_out = ntuple(i->inner[i]*size(A,i),ndims_out)::Dims
 
     indices_in = Array(Int, ndims_in)
     indices_out = Array(Int, ndims_out)
 
     length_out = prod(size_out)
-    R = Array(T, size_out...)
+    R = Array(T, size_out)
 
     for index_out in 1:length_out
         ind2sub!(indices_out, size_out, index_out)
@@ -1146,7 +1122,7 @@ function repeat{T}(A::Array{T};
                 indices_in[t] = fld1(indices_in[t], inner[t])
             end
         end
-        index_in = sub2ind(size_in, indices_in)
+        index_in = sub2ind(size_in, indices_in...)
         R[index_out] = A[index_in]
     end
 

base/deprecated.jl

@@ -321,6 +321,9 @@ end
 @deprecate flipud(A::AbstractArray) flipdim(A, 1)
 @deprecate fliplr(A::AbstractArray) flipdim(A, 2)
 
+@deprecate sub2ind{T<:Integer}(dims::Array{T}, sub::Array{T}) sub2ind(tuple(dims...), sub...)
+@deprecate ind2sub!{T<:Integer}(sub::Array{T}, dims::Array{T}, ind::T) ind2sub!(sub, tuple(dims...), ind)
+
 @deprecate strftime     Libc.strftime
 @deprecate strptime     Libc.strptime
 @deprecate flush_cstdio Libc.flush_cstdio