I have an array with five axes:
colors = numpy.zeros(3, 3, 3, 6, 3))
I want to iterate over it using multi_index like in the second example from this link, but instead of iterating over the entire 5 dimensions, I want to iterate through the first three. A Pythonic way of doing it (without involving Numpy) would be something like this:
indexes = itertools.product(range(3), repeat=3)
for coordinates in indexes:
colors[coordinates]
How can I implement this in pure Numpy?
us numpy.ndindex():
for idx in np.ndindex(*colors.shape[:3]):
data = colors[coordinates]
As I understand the question, what you mainly want is a numpy substitute for itertools.product(). The nearest analogue in numpy would be numpy.indices(). If we modify the code sample in the question just slightly, in order to show us what output it is that we will need to be able to reproduce when working purely in numpy:
indexes = itertools.product(range(3), repeat=3)
for coordinates in indexes:
print(coordinates)
we get the following result:
(0, 0, 0)
(0, 0, 1)
(0, 0, 2)
(0, 1, 0)
(0, 1, 1)
(0, 1, 2)
(0, 2, 0)
(0, 2, 1)
(0, 2, 2)
(1, 0, 0)
(1, 0, 1)
(1, 0, 2)
(1, 1, 0)
(1, 1, 1)
(1, 1, 2)
(1, 2, 0)
(1, 2, 1)
(1, 2, 2)
(2, 0, 0)
(2, 0, 1)
(2, 0, 2)
(2, 1, 0)
(2, 1, 1)
(2, 1, 2)
(2, 2, 0)
(2, 2, 1)
(2, 2, 2)
The following code sample will reproduce this result line-for-line exactly, using numpy.indices() instead of itertools.product():
import numpy
a, b, c = numpy.indices((3,3,3))
indexes = numpy.transpose(numpy.asarray([a.flatten(), b.flatten(), c.flatten()]))
for coordinates in indexes:
print(tuple(coordinates))
