2

I need to be able to return a part of a multidimensional array, but I don't know how to do this in a correct way. The way I do it seems very naive:

import numpy as np
a=np.ones([3,3,3,3,3])
b=np.asarray([2,2])
c=np.asarray([2,2])
print a[b[0],b[1],:,c[0],c[1]]

and will return

[1,1,1]

However what I want is something like this:

a=np.ones([3,3,3,3,3])
b=np.asarray([2,2])
c=np.asarray([2,2])
print a[b,:,c]

Which returns the a itself, Although I want it to return [1,1,1] instead.

And I don't know why. How can I read part of an array without specifying element by element but giving the indices of the array I want as a pack?

P.S. Thanks to @hcwhsa, I updated the question to address more specifically what I want.

Cupitor
  • 11,007
  • 19
  • 65
  • 91
  • What exactly do you want that second thing to return? In particular, what do you want `a[b,:,c].shape` to be? – user2357112 Nov 25 '13 at 00:58
  • @user2357112 The same thing, i.e. `[1,1,1]` just with this new way of calling – Cupitor Nov 25 '13 at 00:59
  • What's your use case? Why do you have these `b` and `c` arrays? – user2357112 Nov 25 '13 at 01:01
  • @user2357112, Well I have a giant joint distribution and sometimes I need to marginalize it. Then I need to partly sum up. Or sample from a specific part of it. Obviously in a discrete case. – Cupitor Nov 25 '13 at 01:09
  • Marginalizing sounds like a job for `sum` with the optional `axis` parameter. It can take a tuple of axes to sum over. If you want to select specific indices on specific axes, where both indices and axes are determined at runtime, I'm not sure what the best way to go about that is, but you can construct a tuple of `slice` objects and pass that to the indexing operator easily enough. – user2357112 Nov 25 '13 at 01:17
  • @user2357112, Thank you! Yes but that is not the only thing I do. I need matrix multiplication in sub-matrices which needs this kind of indexing. What is slice object? – Cupitor Nov 25 '13 at 01:38
  • 2
    `a[tuple(b) + (slice(None),) + tuple(c)]` – askewchan Nov 25 '13 at 01:58
  • Oh. Wow! The old mate! We had a chat here:http://stackoverflow.com/questions/15415455/plotting-probability-density-function-by-sample-with-matplotlib Thanks dude. – Cupitor Nov 25 '13 at 02:01
  • @Naji, haha we meet again. I added a couple more options in my answer. – askewchan Nov 25 '13 at 04:11

2 Answers2

2

Define b as a tuple:

>>> b = (2, 2)
>>> a[b]
array([ 1.,  1.,  1.])

Or convert it to a tuple before passing it to a[]

>>> b = np.asarray([2,2])
>>> a[tuple(b)]
array([ 1.,  1.,  1.])
Ashwini Chaudhary
  • 244,495
  • 58
  • 464
  • 504
  • Thank you. Vote up. But in my code b is array and not a tuple and the "tuple(map(tuple, arr))" doesn't work here. – Cupitor Nov 25 '13 at 00:46
  • And what about the case that my dimension is higher and I need the subarray for example: `a[b,:,:,c]` Where `b` and `c` are arrays? – Cupitor Nov 25 '13 at 00:50
1

I can think of two ways to do this, neither is perfect. One is to roll the axis you want to get all of to the end:

ax = 2 # the axis you want to have all values in
np.rollaxis(a, ax, a.ndim)[tuple(np.r_[b,c])]

This works for a[b,:,:,c] if you move two axes to the back (be careful in the index shift for axis number!)

np.rollaxis(np.rollaxis(a, ax, a.ndim), ax, a.ndim)[tuple(np.r_[b,c])]

where np.rollaxis(a, ax, a.ndim) moves the axis ax you want to keep all of to the end:

a = np.zeros((1,2,3,4,5))
a.shape
#(1,2,3,4,5)
np.rollaxis(a, ax, a.ndim).shape
#(1,2,4,5,3)

And the np.r_[b,c] just concatentes the two arrays. You could also do: tuple(np.concatenate([b,c]))

Or, you can use the one from my comment:

a[tuple(b) + (slice(None),) + tuple(c)]

where slice is the object that the start:end:step syntax creates. None gives you the :, but you can create it dynamically (without having to type the : in the right spot). So, a[1:3] is equivalent to a[slice(1,3)], and a[:3] is a[slice(None,3)]. I've wrapped it inside a tuple so that it can be "added" to the other two tuples to create one long tuple.

askewchan
  • 45,161
  • 17
  • 118
  • 134