CODEWITHSUNDEEP
pythonarraysnumpy
Cupitor
Cupitor

Reputation: 11637

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

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.

Upvotes: 2

Views: 1746

Answers (2)

askewchan
askewchan

Reputation: 46530

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.

Upvotes: 1

Ashwini Chaudhary
Ashwini Chaudhary

Reputation: 250891

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

Upvotes: 2

Related Questions