How to return a subarray of a multidimensional array in Python?

Question

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.

What exactly do you want that second thing to return? In particular, what do you want a[b,:,c].shape to be? — user2357112
– user2357112, Commented Nov 25, 2013 at 0:58
@user2357112 The same thing, i.e. [1,1,1] just with this new way of calling — Cupitor
– Cupitor, Commented Nov 25, 2013 at 0:59
What's your use case? Why do you have these b and c arrays? — user2357112
– user2357112, Commented Nov 25, 2013 at 1: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
– Cupitor, Commented Nov 25, 2013 at 1:09

Ashwini Chaudhary · Accepted Answer · 2013-11-25 00:47:43Z

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.])

edited Nov 25, 2013 at 0:47

answered Nov 25, 2013 at 0:37

Ashwini Chaudhary

252k60 gold badges478 silver badges519 bronze badges

Sign up to request clarification or add additional context in comments.

2 Comments

Cupitor Over a year ago

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 Over a year ago

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?

askewchan · Accepted Answer · 2013-11-25 04:11:16Z

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.

Collectives™ on Stack Overflow

How to return a subarray of a multidimensional array in Python?

2 Answers 2

2 Comments

Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

2 Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related