Iterating over numpy array keeping the indices

Question

Consider the following snippet of code:

n_samples, n_rows, n_cols, n_boxes_per_cell, _ = Y_pred.shape
    for example in range(n_samples):
        for x_grid in range(n_rows):
            for y_grid in range(n_cols):
                for bnd_box in range(n_boxes_per_cell):
                    bnd_box_label = Y_pred[example, x_grid, y_grid, bnd_box]
                    do_some_stuff(bnd_box_label, x_grid, y_grid)

How can I get functionally equivalent code with at most one explicit iteration? Notice that I need the indices x_grid and y_grid.

Paul Panzer · Accepted Answer · 2018-11-24 22:01:11Z

3

You can use np.ndindex:

for example, x_grid, y_grid, bnd_box in np.ndindex(Y_pred.shape[:4]):
    etc.

answered Nov 24, 2018 at 22:01

Paul Panzer

53.3k3 gold badges60 silver badges103 bronze badges

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

Comments

Halbeard · Accepted Answer · 2018-11-24 22:12:32Z

1

I'm not sure if this is what you're looking for, but you can always build a generator out of multiple iterables:

all_combinations = ((a, b, c, d) for a in range(n_samples)
                                 for b in range(n_rows) 
                                 for c in range(n_cols) 
                                 for d in range(n_boxes_per_cell))

for examples, x_grid, y_grid, bnd_box in all_combinations:
    do stuff

This is the same as using itertools.product(*iterables) and valid for any iterable, not just iteration over indices/integers.

answered Nov 24, 2018 at 22:12

Halbeard

1215 bronze badges

Collectives™ on Stack Overflow

Iterating over numpy array keeping the indices

2 Answers 2

Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related