0

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.

mathfux
  • 5,759
  • 1
  • 14
  • 34

2 Answers2

1

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))
chepner
  • 497,756
  • 71
  • 530
  • 681
1

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)
Divakar
  • 218,885
  • 19
  • 262
  • 358
  • Thanks. This is a quite beautiful approach I couldn't think of before – mathfux Sep 04 '20 at 16:16
  • Sounds like `Ellipsis` is an alternative for `slice()`. But what about performance? I mean using `x.reshape(1,-1,-1)` vs `x[Ellipsis, None, None]`? – mathfux Sep 04 '20 at 16:35
  • @mathfux Both are creating views. So, don't think performance should be an issue. [`Elllipsis is an alias for "..."`](https://stackoverflow.com/questions/772124/what-does-the-ellipsis-object-do). In our case, it becomes an equivalent to your ":". – Divakar Sep 04 '20 at 16:37