Skip to content

Commit 366a8ed

Browse files
authored
Update indexing.py
1 parent 21817d5 commit 366a8ed

File tree

1 file changed

+11
-8
lines changed

1 file changed

+11
-8
lines changed

sparse/compressed/indexing.py

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -58,20 +58,23 @@ def getitem(x, key):
5858
count += 1
5959

6060
reordered_key = [key[i] for i in x.axis_order]
61-
61+
62+
# prepare for converting to flat indices
63+
for i, ind in enumerate(reordered_key[:x.axisptr]):
64+
if isinstance(ind,slice):
65+
reordered_key[i] = range(ind.start,ind.stop,ind.step)
66+
for i, ind in enumerate(reordered_key[x.axisptr:]):
67+
if isinstance(ind, Integral):
68+
reordered_key[i + x.axisptr] = [ind]
69+
elif isinstance(ind, slice):
70+
reordered_key[i + x.axisptr] = np.arange(ind.start, ind.stop, ind.step)
71+
6272
# find starts and ends of rows
6373
a = x.indptr[:-1].reshape(x.reordered_shape[:x.axisptr])
6474
b = x.indptr[1:].reshape(x.reordered_shape[:x.axisptr])
6575
starts = a[tuple(reordered_key[:x.axisptr])].flatten()
6676
ends = b[tuple(reordered_key[:x.axisptr])].flatten()
6777

68-
# prepare for converting to flat indices
69-
for i, ind in enumerate(reordered_key):
70-
if isinstance(ind, Integral):
71-
reordered_key[i] = [ind]
72-
elif isinstance(ind, slice):
73-
reordered_key[i] = np.arange(ind.start, ind.stop, ind.step)
74-
7578
shape = np.array(shape)
7679

7780
cols = convert_to_flat(reordered_key, x.reordered_shape, x.axisptr)

0 commit comments

Comments
 (0)