Average chunks of numpy arrays

I just found a great tool in numpy: ufunc.reduceat(). In fact, I think it's so good that I started this blog to tell the world about it!

I found it while searching for a way to average 2-d chunks or blocks of a numpy array. Given a 2-d array physiography.elevation, and chunks defined by (row_segments, col_segments), where the chunks are elements in physiography.elevation between row_segments[0] and row_segments[1], row_segments[2] and row_segments[3], row_segments[4] and row_segments[5],and so on (same for col_segments), chunk averages can be calculated like this:

# Average over all segments
# Mask out no_data and missing_data values
elevation = np.ma.array(physiography.elevetion, fill_value=0,
                      mask=((physiography.elevation ==
                             physiography.no_data) *
                            (physiography.elevation ==
                             physiography.missing_data)))
# Sum chunks defined by (row_segments, col_segments),
# discarding every other (odd) sum, as they are inbetween segments
sum = np.add.reduceat
mean_elevation = sum(sum(elevation.filled(), row_segments)[::2],
                   col_segments, axis=1)[:, ::2]
# Divide by total number of pixels in each segment
mean_elevation /= sum(sum(~elevation.mask, row_segments)[::2],
                    col_segments, axis=1)[:, ::2]

Pretty neat.

The inspiration was taken from a message on the Numpy-discussion mailing list.