CODEWITHSUNDEEP
pythonarraysnumpyslice
DilithiumMatrix
DilithiumMatrix

Reputation: 18667

Store multidimensional numpy array slice with newaxis to object

I have some code where I repeatedly need to repeatedly broadcast arrays in complex ways, for example:

a = b[np.newaxis, ..., :, np.newaxis] * c[..., np.newaxis, np.newaxis, :]

Is there an object to which I can store these slicing specifications?

i.e. (but obviously this doesn't work):

s1 = magic([np.newaxis, ..., :, np.newaxis])
s2 = magic([..., np.newaxis, np.newaxis, :])

Edit: perhaps this could be done with numpy.broadcast_to, but it's unclear how exactly while making sure that the correct axes are broadcast over...

Upvotes: 0

Views: 200

Answers (1)

user2357112
user2357112

Reputation: 281758

You can construct the index tuple manually, but NumPy includes a helper for it:

slice_tuple = np.s_[np.newaxis, ..., :, np.newaxis]

Then b[np.newaxis, ..., :, np.newaxis] is equivalent to b[slicetuple].

For reference, constructing the tuple manually would be (np.newaxis, Ellipsis, slice(None), np.newaxis). Also, np.newaxis is None, so (None, Ellipsis, slice(None), None) would be equivalent.

np.s_ can be reimplemented yourself as follows:

class IndexHelper(object):
    def __getitem__(self, arg):
        return arg

s_ = IndexHelper()

Upvotes: 5

Related Questions