How to do equivalent of block_reduce on a masked array?

I'm calculating an aggregate value over smaller blocks in a 2D numpy array. I'd like to exclude values 0 from the aggregation operation in an efficient manner (rather than for and if statements).

I'm using skimage.measure.block_reduce and numpy.ma.masked_equal, but it looks like block_reduce ignores the mask.

import numpy as np
import skimage
a = np.array([[2,4,0,12,5,7],[6,0,8,4,3,9]])
zeros_included = skimage.measure.block_reduce(a,(2,2),np.mean)

includes 0s and (correctly) produces

zeros_included
array([[3., 6., 6.]])

I was hoping

masked = np.ma.masked_equal(a,0)
zeros_excluded = skimage.measure.block_reduce(masked,(2,2),np.mean)

would do the trick, but still produces

zeros_excluded
array([[3., 6., 6.]])

The desired result would be:

array([[4., 8., 6.]])

I'm looking for a pythonesque way to achieve the correct result, use of skimage is optional. Of course my actual arrays and blocks are much bigger than in this example, hence the need for efficiency.

Thanks for your interest.

Solution

You could use np.nanmean, but you'll have to modify original array or create a new one:

import numpy as np
import skimage

a = np.array([[2,4,0,12,5,7],[6,0,8,4,3,9]])


b = a.astype("float")
b[b==0] = np.nan
zeros_excluded = skimage.measure.block_reduce(b,(2,2), np.nanmean)
zeros_excluded

# array([[4., 8., 6.]])