numpy advanced indexing: transparently optimising ranges?

Question

This is inspired by this post.

Consider a function f that returns a 1d np.ndarray idx of indices which the user will typically use to index other arrays. Assume further that a frequent outcome is for f to return the full range of legal indices. In the linked post it is suggested this be special-cased by f returning slice(None) instead of np.arange(maxind).

Since advanced indexing comes at a cost

>>> a = np.arange(1_000_000)
>>> direct = lambda: np.sum(a[:])
>>> indirect = lambda: np.sum(a[a])
>>> timeit(direct, number=100)
0.07656216900795698
>>> timeit(indirect, number=100)
0.2885982050211169

this looks like a reasonable optimisation at first sight.

Unfortunately, it is not "correct". Imagine, for example, the user wants to create a one-hot representation of idx. One straightforward way of going about this is

result = np.zeros((k, maxind), dtype=int)
result[np.arange(k), idx] = 1

This breaks if np.arange(maxind) is substituted by slice(None) (It will fill the whole of result with ones).

So my question is: Can one have one's cake and eat it here, i.e.:

Is there anything f could return that faithfully mimics the semantics of np.arange(maxind) while avoiding advanced indexing where possible?

Since I'm almost resigned to the answer being no:

What's the next best thing?

Maybe return an "enhanced np.s_", i.e. an object with an engineered __getitem__?

class smart_idx:
    def __init__(self, n):
        self.n = n
    def __getitem__(self, idx):
        idx = idx if isinstance(idx, tuple) else (idx,)
        if idx:
            count = idx.count('X')
            need_adv = count > 1
            if count == 1:
                for i in idx:
                    if not isinstance(i, slice) and i != Ellipsis:
                        need_adv = True
                        break
            repl = np.arange(self.n) if need_adv else slice(None)
            return tuple(repl if i == 'X' else i for i in idx)
        return slice(None)

The user would have to use it like

data[idx[3, 4:9, 'X', [1,3,2,6]]]
data[idx['X', ..., :4:-1]]
data[idx[]]

and __getitem__, detecting the advanced index would decide to replace 'X' with np.arange(4) in the first and with slice(None) in the other two examples.

But that's rather clunky, not to mention that the overhead added might eat up whatever speed we gained.

Are there simpler strategies?

If idx is a list/array you are going to get advanced indexing, regardless of whether you use a slice or arange for the to the other index. In the [arange...,idx] indexing it is picking one item from each row. Indexing on a flattened array is faster, but counterbalanced by the cost of computing the flat index. — hpaulj
@hpaulj "If idx is a list/array you are going to get advanced indexing, regardless". That is why we are trying to substitute idx with slice(None) in special cases where that makes sense (basically idx==np.arange(maxind) + no other advanced indices around). — Paul Panzer
What @hpaulj was trying to say, is that regardless of what idx is, when you do indexing for arbitrary positions in >1D (2D in your example), e.g. [something..., idx] it is advanced indexing regardless of what idx could possibly be, unless both something and idx are slice objects, but that case would never mimic your desired behaviour. — Imanol Luengo
@ImanolLuengo Yes, I know that. Still, no advanced indexing in any axis is a common enough use-case, so it would be nice if one could optimise that in a more or less transparent way. — Paul Panzer

hpaulj hpaulj · Accepted Answer · 2017-02-22T21:15:14

In [104]: x=np.arange(12).reshape(4,3)

These look the same, though one is a copy, the other a view:

In [107]: x[np.arange(0,4,2),:]
Out[107]: 
array([[0, 1, 2],
       [6, 7, 8]])
In [108]: x[0:4:2,:]
Out[108]: 
array([[0, 1, 2],
       [6, 7, 8]])

But if the 2nd index is an array, the arange and slice aren't substitutes.

In [109]: idx=np.array([0,2])
In [110]: x[np.arange(0,4,2),idx]
Out[110]: array([0, 8])
In [111]: x[0:4:2,idx]
Out[111]: 
array([[0, 2],
       [6, 8]])

To match the sliced version I have to add a dimension to the arange.

In [113]: x[np.ix_(np.arange(0,4,2),idx)]
Out[113]: 
array([[0, 2],
       [6, 8]])
In [114]: x[np.arange(0,4,2)[:,None],idx]
Out[114]: 
array([[0, 2],
       [6, 8]])

I'm not aware of a slice expression that produces Out[110].

So apart from replacing an arange with a slice, we need to pay attention as to how the advanced indexing arrays broadcast against each other, and what broadcasting is implied by slicing.

With 3 or more dimensions, mixing slices and advanced indexing gets even more complicated, as described in https://docs.scipy.org/doc/numpy/reference/arrays.indexing.html#combining-advanced-and-basic-indexing

numpy advanced indexing: transparently optimising ranges?

1 Answers