CODEWITHSUNDEEP
pythonnumpyslicenumpy-slicing
mathfux
mathfux

Reputation: 5949

How to determine slices that contains "None" and ":" in numpy?

I'm doing a substitute of itertools.product as an exercise and my approach is to use np.broadcast:

import numpy as np
x = np.array([4,3,1])
y = np.array([2,4,0])
z = np.array([0,1,2])
ix = np.broadcast(x[:,None,None], y[None,:,None], z[None, None, :])
print(*ix)

It works OK for this time but how to create all these 'eyed-arranged` new axis in an automatic way if I want, say 7 dimensions like this:

[:,None,None,None,None,None,None]
[None,:,None,None,None,None,None]
[None,None,:,None,None,None,None]
[None,None,None,:,None,None,None]
[None,None,None,None,:,None,None]
[None,None,None,None,None,:,None]
[None,None,None,None,None,None,:]

I expect something like np.ix_ that allows to use all these Nones and :s in assignment of these slices.

Upvotes: 0

Views: 180

Answers (2)

chepner
chepner

Reputation: 531808

Each index is an n-element tuple. The diagonal element of each tuple, that is, the ith element of the ith tuple, is a slice object. Specifically, : is short for slice(None). (In general, x:y:z is the object slice(x, y, z), with any missing elements being None, though at least one argument must be provided).

arguments = [(None,)*i + (slice(None),) + (None,)*(6-i) for i in range(7)]
# E.g., i == 0 -> () + (slice(None),) + (None, None, None, None, None, None)
#              == (slice(None), None, None, None, None, None, None)
ix = np.broadcast(*(x[arg] for arg in arguments))

Upvotes: 1

Divakar
Divakar

Reputation: 221624

Approach #1

That same eye-arrangement could be use to reshape all participating arrays -

A = [x,y,z] # all arrays
s = 1-2*np.eye(len(A), dtype=int)
out = [a.reshape(i) for i,a in zip(s,A)]

Here's the trick part of shaping with s :

In [53]: s
Out[53]: 
array([[-1,  1,  1], # reshape to keep all axes singleton except first
       [ 1, -1,  1], #                                       .. second 
       [ 1,  1, -1]]) #                                      ... third

Thus, considering [-1, 1, 1] is to reshape to keep all axes singleton except first. This is same as [:,None,None] and so on.

Approach #2

Using the same eye-arrangement to get an indexer -

idx = np.where(np.eye(len(A)), Ellipsis, None)
out = [a[tuple(i)] for a,i in zip(A,idx)]

The indexer is :

In [77]: idx
Out[77]: 
array([[Ellipsis, None, None],
       [None, Ellipsis, None],
       [None, None, Ellipsis]], dtype=object)

Upvotes: 1

Related Questions