Skip to content

Commit

Permalink
Drop dimensions indexed by scalars
Browse files Browse the repository at this point in the history
  • Loading branch information
mbauman committed Nov 4, 2015
1 parent 6ad25fc commit f90f304
Show file tree
Hide file tree
Showing 10 changed files with 73 additions and 36 deletions.
5 changes: 4 additions & 1 deletion base/linalg/diagonal.jl
Original file line number Diff line number Diff line change
Expand Up @@ -121,7 +121,10 @@ function A_ldiv_B!{T}(D::Diagonal{T}, V::AbstractMatrix{T})
if d == zero(T)
throw(SingularException(i))
end
V[i,:] *= inv(d)
d⁻¹ = inv(d)
for j=1:size(V,2)
@inbounds V[i,j] *= d⁻¹
end
end
V
end
Expand Down
37 changes: 20 additions & 17 deletions base/multidimensional.jl
Original file line number Diff line number Diff line change
Expand Up @@ -166,15 +166,15 @@ index_lengths_dim(A, dim, ::Colon) = (trailingsize(A, dim),)
@inline index_lengths_dim(A, dim, i::AbstractArray{Bool}, I...) = (sum(i), index_lengths_dim(A, dim+1, I...)...)
@inline index_lengths_dim(A, dim, i::AbstractArray, I...) = (length(i), index_lengths_dim(A, dim+1, I...)...)

# shape of array to create for getindex() with indexes I, dropping trailing scalars
# shape of array to create for getindex() with indexes I, dropping scalars
index_shape(A::AbstractArray, I::AbstractArray) = size(I) # Linear index reshape
index_shape(A::AbstractArray, I::AbstractArray{Bool}) = (sum(I),) # Logical index
index_shape(A::AbstractArray, I::Colon) = (length(A),)
@inline index_shape(A::AbstractArray, I...) = index_shape_dim(A, 1, I...)
index_shape_dim(A, dim, I::Real...) = ()
index_shape_dim(A, dim, ::Colon) = (trailingsize(A, dim),)
@inline index_shape_dim(A, dim, ::Colon, i, I...) = (size(A, dim), index_shape_dim(A, dim+1, i, I...)...)
@inline index_shape_dim(A, dim, ::Real, I...) = (1, index_shape_dim(A, dim+1, I...)...)
@inline index_shape_dim(A, dim, ::Real, I...) = (index_shape_dim(A, dim+1, I...)...)
@inline index_shape_dim(A, dim, i::AbstractVector{Bool}, I...) = (sum(i), index_shape_dim(A, dim+1, I...)...)
@inline index_shape_dim(A, dim, i::AbstractVector, I...) = (length(i), index_shape_dim(A, dim+1, I...)...)

Expand Down Expand Up @@ -238,7 +238,8 @@ end
$(Expr(:meta, :inline))
D = eachindex(dest)
Ds = start(D)
@nloops $N i dest d->(j_d = unsafe_getindex(I[d], i_d)) begin
idxlens = index_lengths(src, I...) # TODO: unsplat?
@nloops $N i d->(1:idxlens[d]) d->(j_d = unsafe_getindex(I[d], i_d)) begin
d, Ds = next(D, Ds)
v = @ncall $N unsafe_getindex src j
unsafe_setindex!(dest, v, d)
Expand All @@ -248,18 +249,18 @@ end
end

# checksize ensures the output array A is the correct size for the given indices
checksize(A::AbstractArray, I::AbstractArray) = size(A) == size(I) || throw(DimensionMismatch("index 1 has size $(size(I)), but size(A) = $(size(A))"))
checksize(A::AbstractArray, I::AbstractArray{Bool}) = length(A) == sum(I) || throw(DimensionMismatch("index 1 selects $(sum(I)) elements, but length(A) = $(length(A))"))
@generated function checksize(A::AbstractArray, I...)
N = length(I)
quote
@nexprs $N d->(_checksize(A, d, I[d]) || throw(DimensionMismatch("index $d selects $(length(I[d])) elements, but size(A, $d) = $(size(A,d))")))
end
end
_checksize(A::AbstractArray, dim, I) = size(A, dim) == length(I)
_checksize(A::AbstractArray, dim, I::AbstractVector{Bool}) = size(A, dim) == sum(I)
_checksize(A::AbstractArray, dim, ::Colon) = true
_checksize(A::AbstractArray, dim, ::Real) = size(A, dim) == 1
@noinline throw_checksize_error(arr, dim, idx) = throw(DimensionMismatch("index $d selects $(length(I[d])) elements, but size(A, $d) = $(size(A,d))"))

checksize(A::AbstractArray, I::AbstractArray) = size(A) == size(I) || throw_checksize_error(A, 1, I)
checksize(A::AbstractArray, I::AbstractArray{Bool}) = length(A) == sum(I) || throw_checksize_error(A, 1, I)

checksize(A::AbstractArray, I...) = _checksize(A, 1, I...)
_checksize(A::AbstractArray, dim) = true
# Skip scalars
_checksize(A::AbstractArray, dim, ::Real, J...) = _checksize(A, dim, J...)
_checksize(A::AbstractArray, dim, I, J...) = (size(A, dim) == length(I) || throw_checksize_error(A, dim, I); _checksize(A, dim+1, J...))
_checksize(A::AbstractArray, dim, I::AbstractVector{Bool}, J...) = (size(A, dim) == sum(I) || throw_checksize_error(A, dim, I); _checksize(A, dim+1, J...))
_checksize(A::AbstractArray, dim, ::Colon, J...) = _checksize(A, dim+1, J...)

@inline unsafe_setindex!(v::BitArray, x::Bool, ind::Int) = (Base.unsafe_bitsetindex!(v.chunks, x, ind); v)
@inline unsafe_setindex!(v::BitArray, x, ind::Real) = (Base.unsafe_bitsetindex!(v.chunks, convert(Bool, x), to_index(ind)); v)
Expand Down Expand Up @@ -585,7 +586,8 @@ end

storeind = 1
Xc, Bc = X.chunks, B.chunks
@nloops($N, i, d->1:size(X, d+1),
idxlens = index_lengths(B, I0, I...) # TODO: unsplat?
@nloops($N, i, d->(1:idxlens[d+1]),
d->nothing, # PRE
d->(ind += stride_lst_d - gap_lst_d), # POST
begin # BODY
Expand All @@ -608,7 +610,8 @@ end
$(symbol(:offset_, N)) = 1
ind = 0
Xc, Bc = X.chunks, B.chunks
@nloops $N i X d->(offset_{d-1} = offset_d + (unsafe_getindex(I[d], i_d)-1)*stride_d) begin
idxlens = index_lengths(B, I...) # TODO: unsplat?
@nloops $N i d->(1:idxlens[d]) d->(offset_{d-1} = offset_d + (unsafe_getindex(I[d], i_d)-1)*stride_d) begin
ind += 1
unsafe_bitsetindex!(Xc, unsafe_bitgetindex(Bc, offset_0), ind)
end
Expand Down
2 changes: 0 additions & 2 deletions base/sparse/sparsematrix.jl
Original file line number Diff line number Diff line change
Expand Up @@ -1352,8 +1352,6 @@ function getindex{T}(A::SparseMatrixCSC{T}, i0::Integer, i1::Integer)
((r1 > r2) || (A.rowval[r1] != i0)) ? zero(T) : A.nzval[r1]
end

getindex{T<:Integer}(A::SparseMatrixCSC, i::Integer, J::AbstractVector{T}) = getindex(A,[i],J)

# Colon translation
getindex(A::SparseMatrixCSC, ::Colon, ::Colon) = copy(A)
getindex(A::SparseMatrixCSC, i, ::Colon) = getindex(A, i, 1:size(A, 2))
Expand Down
33 changes: 33 additions & 0 deletions base/sparse/sparsevector.jl
Original file line number Diff line number Diff line change
Expand Up @@ -345,6 +345,7 @@ sprandbool(r::AbstractRNG, n::Integer, p::AbstractFloat) = sprand(r, n, p, trueb

## Indexing into Matrices can return SparseVectors

# Column slices
function getindex(x::SparseMatrixCSC, ::Colon, j::Integer)
checkbounds(x, :, j)
r1 = convert(Int, x.colptr[j])
Expand All @@ -369,6 +370,38 @@ end
SparseVector(M.m, M.rowval, M.nzval)
end

# Row slices
getindex(A::SparseMatrixCSC, i::Integer, ::Colon) = A[i, 1:end]
getindex{Tv,Ti}(A::SparseMatrixCSC{Tv,Ti}, i::Integer, J::AbstractVector{Bool}) = A[i, find(J)]
function Base.getindex{Tv,Ti}(A::SparseMatrixCSC{Tv,Ti}, i::Integer, J::AbstractVector)
checkbounds(A, i, J)
nJ = length(J)
colptrA = A.colptr; rowvalA = A.rowval; nzvalA = A.nzval

nzinds = Array(Ti, 0)
nzvals = Array(Tv, 0)

# adapted from SparseMatrixCSC's sorted_bsearch_A
ptrI = 1
@inbounds for j = 1:nJ
col = J[j]
rowI = i
ptrA = colptrA[col]
stopA = colptrA[col+1]-1
if ptrA <= stopA
if rowvalA[ptrA] <= rowI
ptrA = searchsortedfirst(rowvalA, rowI, ptrA, stopA, Base.Order.Forward)
if ptrA <= stopA && rowvalA[ptrA] == rowI
push!(nzinds, ptrI)
push!(nzvals, nzvalA[ptrA])
end
end
ptrI += 1
end
end
return SparseVector(nJ, nzinds, nzvals)
end


# Logical and linear indexing into SparseMatrices
getindex{Tv}(A::SparseMatrixCSC{Tv}, I::AbstractVector{Bool}) = _logical_index(A, I) # Ambiguities
Expand Down
4 changes: 2 additions & 2 deletions base/subarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -604,8 +604,8 @@ end

# Indexing with non-scalars. For now, this returns a copy, but changing that
# is just a matter of deleting the explicit call to copy.
getindex{T,N,P,IV}(V::SubArray{T,N,P,IV}, I::ViewIndex...) = copy(sub(V, I...))
unsafe_getindex{T,N,P,IV}(V::SubArray{T,N,P,IV}, I::ViewIndex...) = copy(sub_unsafe(V, I))
getindex{T,N,P,IV}(V::SubArray{T,N,P,IV}, I::ViewIndex...) = copy(slice(V, I...))
unsafe_getindex{T,N,P,IV}(V::SubArray{T,N,P,IV}, I::ViewIndex...) = copy(slice_unsafe(V, I))

# Nonscalar setindex! falls back to the AbstractArray versions

Expand Down
12 changes: 6 additions & 6 deletions test/arrayops.jl
Original file line number Diff line number Diff line change
Expand Up @@ -93,7 +93,7 @@ a = rand(1, 1, 8, 8, 1)
sz = (5,8,7)
A = reshape(1:prod(sz),sz...)
@test A[2:6] == [2:6;]
@test A[1:3,2,2:4] == cat(3,46:48,86:88,126:128)
@test A[1:3,2,2:4] == cat(2,46:48,86:88,126:128)
@test A[:,7:-3:1,5] == [191 176 161; 192 177 162; 193 178 163; 194 179 164; 195 180 165]
@test A[:,3:9] == reshape(11:45,5,7)
rng = (2,2:3,2:2:5)
Expand All @@ -111,7 +111,7 @@ tmp = cat([1,3],blk,blk)

x = rand(2,2)
b = x[1,:]
@test isequal(size(b), (1, 2))
@test isequal(size(b), (2,))
b = x[:,1]
@test isequal(size(b), (2,))

Expand All @@ -129,7 +129,7 @@ B[[3,1],[2,4]] = [21 22; 23 24]
B[4,[2,3]] = 7
@test B == [0 23 1 24 0; 11 12 13 14 15; 0 21 3 22 0; 0 7 7 0 0]

@test isequal(reshape(1:27, 3, 3, 3)[1,:], [1 4 7 10 13 16 19 22 25])
@test isequal(reshape(1:27, 3, 3, 3)[1,:], [1, 4, 7, 10, 13, 16, 19, 22, 25])

a = [3, 5, -7, 6]
b = [4, 6, 2, -7, 1]
Expand Down Expand Up @@ -575,12 +575,12 @@ let

@test isequal(c[:,1], cv)
@test isequal(c[:,3], cv2)
@test isequal(c[4,:], [2.0 2.0 2.0 2.0]*1000)
@test isequal(c[4,:], [2.0, 2.0, 2.0, 2.0]*1000)

c = cumsum_kbn(A, 2)

@test isequal(c[1,:], cv2')
@test isequal(c[3,:], cv')
@test isequal(c[1,:], cv2)
@test isequal(c[3,:], cv)
@test isequal(c[:,4], [2.0,2.0,2.0,2.0]*1000)

end
Expand Down
6 changes: 3 additions & 3 deletions test/bitarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -171,10 +171,10 @@ function gen_getindex_data()
m1, m2 = rand_m1m2()
produce((m1, m2, Bool))
m1, m2 = rand_m1m2()
produce((m1, 1:m2, BitMatrix))
produce((m1, :, BitMatrix))
produce((m1, 1:m2, BitVector))
produce((m1, :, BitVector))
m1, m2 = rand_m1m2()
produce((m1, randperm(m2), BitMatrix))
produce((m1, randperm(m2), BitVector))
m1, m2 = rand_m1m2()
produce((1:m1, m2, BitVector))
produce((:, m2, BitVector))
Expand Down
2 changes: 1 addition & 1 deletion test/readdlm.jl
Original file line number Diff line number Diff line change
Expand Up @@ -205,7 +205,7 @@ let i18n_data = ["Origin (English)", "Name (English)", "Origin (Native)", "Name
writedlm(i18n_buff, i18n_arr, ',')
@test i18n_arr == readcsv(i18n_buff)

hdr = i18n_arr[1, :]
hdr = i18n_arr[1:1, :]
data = i18n_arr[2:end, :]
writedlm(i18n_buff, i18n_arr, ',')
@test (data, hdr) == readcsv(i18n_buff, header=true)
Expand Down
6 changes: 3 additions & 3 deletions test/sparsedir/sparse.jl
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,7 @@ de33 = eye(3)
# also side-checks sparse ref
for i = 1 : 10
a = sprand(5, 4, 0.5)
@test all([a[1:2,1:2] a[1:2,3:4]; a[3:5,1] [a[3:4,2:4]; a[5,2:4]]] == a)
@test all([a[1:2,1:2] a[1:2,3:4]; a[3:5,1] [a[3:4,2:4]; a[5:5,2:4]]] == a)
end

# sparse ref
Expand Down Expand Up @@ -460,13 +460,13 @@ let a = spzeros(Int, 10, 10)
@test countnz(a) == 0
a[1,:] = 1
@test countnz(a) == 10
@test a[1,:] == sparse(ones(Int,1,10))
@test a[1,:] == sparse(ones(Int,10))
a[:,2] = 2
@test countnz(a) == 19
@test a[:,2] == 2*sparse(ones(Int,10))

a[1,:] = 1:10
@test a[1,:] == sparse([1:10;]')
@test a[1,:] == sparse([1:10;])
a[:,2] = 1:10
@test a[:,2] == sparse([1:10;])
end
Expand Down
2 changes: 1 addition & 1 deletion test/subarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -388,7 +388,7 @@ sA[:] = -3
sA = sub(A, 1:3, 3, 2:5)
@test Base.parentdims(sA) == [1:3;]
@test size(sA) == (3,1,4)
@test sA == A[1:3,3,2:5]
@test sA == A[1:3,3:3,2:5]
@test sA[:] == A[1:3,3,2:5][:]
sA = sub(A, 1:2:3, 1:3:5, 1:2:8)
@test Base.parentdims(sA) == [1:3;]
Expand Down

0 comments on commit f90f304

Please sign in to comment.