2

I have a 3-dimensional numpy array, and I'm trying to index it along its second and third dimensions using two single-dimension numpy arrays. To wit:

np.random.seed(0)
dims = (1,5,5)
test_array = np.arange(np.prod(dims)).reshape(*dims)
test_array

produces:

array([[[ 0,  1,  2,  3,  4],
        [ 5,  6,  7,  8,  9],
        [10, 11, 12, 13, 14],
        [15, 16, 17, 18, 19],
        [20, 21, 22, 23, 24]]])

if I then create the two arrays to index it:

idx_dim1 = np.array([0,2,4])
idx_dim2 = np.array([1,3])

I find that I can't apply these both at the same time:

test_array[:,idx_dim1,idx_dim2]
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
<ipython-input-193-95b23ed3210c> in <module>()
----> 1 test_array[:,idx_dim1,idx_dim2]

IndexError: shape mismatch: indexing arrays could not be broadcast together with shapes (3,) (2,)

I can split them out, by, say

test_array[:, idx_dim1,:][:,:,idx_dim2]

which gives

array([[[ 1,  3],
        [11, 13],
        [21, 23]]])

But this only works in a read-only sense - I can't assign values to test_array like this, since I'm assigning to a slice rather than the original. I'm also not quite sure why using the two indexers together doesn't work. Is there some conceptual nuance of numpy that I'm not understanding here? And is there a good solution that allows me to assign values?

Thanks!

denis
  • 21,378
  • 10
  • 65
  • 88
Chris J Harris
  • 1,597
  • 2
  • 14
  • 26
  • 1
    I replaced your example with something more reproducible. – cs95 Nov 30 '18 at 04:57
  • @coldspeed much obliged. – Chris J Harris Nov 30 '18 at 05:00
  • Are you coming from a MATLAB background? There selecting a block is easy, selecting a scatted set of items awkward. In `numpy` the block indexing is a bit more awkward, at least on the surface. But notice that the error talks about `broadcasting`. How would you add a shape (3,) array to a shape (2,)? The same broadcasting rules apply. – hpaulj Nov 30 '18 at 05:38

2 Answers2

4

Try

test_array[:, idx_dim1[:,None], idx_dim2] = some_value

Where, 

>>> test_array[:, idx_dim1[:,None], idx_dim2]
array([[[ 1,  3],
        [11, 13],
        [21, 23]]])
cs95
  • 379,657
  • 97
  • 704
  • 746
  • Thanks, that's right on the money. But why is that necessary? – Chris J Harris Nov 30 '18 at 04:56
  • @Chrisper It's a little complicated. Thankfully, the docs should be able to explain it much better than I can. I would recommend starting [here](https://docs.scipy.org/doc/numpy-1.15.1/user/basics.indexing.html). – cs95 Nov 30 '18 at 04:59
1

To index cross products (rectangular subsets) of arrays, use np.ix_ —

A = np.zeros(( 1, 5, 5 ))
A[np.ix_( [0], [0,2,4], [1,3] )] = 1
# print dot0( A )
[[[. 1 . 1 .]
  [. . . . .]
  [. 1 . 1 .]
  [. . . . .]
  [. 1 . 1 .]]]

You can also set cross products with ix_ —

B = np.arange( 3 * 2 ).reshape( 3, 2 )
A[np.ix_( [0], [0,2,4], [1,3] )] = B
denis
  • 21,378
  • 10
  • 65
  • 88