CODEWITHSUNDEEP
pythonarraysnumpy
lvella
lvella

Reputation: 13413

How to iterate over initial dimensions of a Numpy array?

I have a Numpy array with shape [1000, 1000, 1000, 3], being the last dimension, sized 3, is contains the triplets of 3D spatial vectors components. How can I use nditer to iterate over each triplet? Like this:

 for vec in np.nditer(my_array, op_flags=['writeonly', <???>]):
     vec = np.array(something)

Upvotes: 2

Views: 416

Answers (1)

hpaulj
hpaulj

Reputation: 231355

I've addressed this question before, but here's a short example:

vec=np.arange(2*2*2*3).reshape(2,2,2,3)
it=np.ndindex(2,2,2)
for i in it:
    print(vec[i])

producing:

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

ndindex constructs a multi-index iterator around a dummy array of the size you give it (here (2,2,2)), and returns it along with a next method.

So you can use ndindex as is, or use it as a model for constructing your on nditer.

Upvotes: 2

Related Questions