def _compare_median2d(a, mask):
actual = median_utils.median2d(a.astype('f4'), mask.astype(np.uint8))
mask_array = mask == 0
n_good_pixels = mask_array.sum(axis=1)[0]
d = a[mask == 0].reshape(a.shape[0], n_good_pixels)
expected = np.median(d.astype('f4'), axis=1)
np.testing.assert_allclose(actual, np.float32(expected))
def _compare_median2d(a, mask):
actual = median_utils.median2d(a.astype('f4'), mask.astype(np.uint8))
mask_array = mask == 0
n_good_pixels = mask_array.sum(axis=1)[0]
d = a[mask == 0].reshape(a.shape[0], n_good_pixels)
expected = np.median(d.astype('f4'), axis=1)
np.testing.assert_allclose(actual, np.float32(expected))
def median(d, axis=None, mask=None):
"""
Find the median of a numpy array. If an axis is provided, then find the median along
the given axis. If mask is included, elements in d that have a non-zero mask value will
be ignored.
Parameters
----------
d : float32 numpy array
Input array to find the median.
axis : int (default is None)
Index of the array to take the median
mask : unit8 or boolean numpy array (default is None)
Numpy array of bitmask values. Non-zero values are ignored when calculating the median.
Returns
-------
med : float32 numpy array
The median value. If axis is None, then we return a single float.
Notes
-----
Makes extensive use of the quick select algorithm written in C, included in quick_select.c.
If all of the elements in the array are masked (or all of the elements of the axis of interest
are masked), we return zero.
"""
if axis is None:
if mask is not None:
median_mask = mask.ravel()
else:
median_mask = np.zeros(d.size, dtype=np.uint8)
output_median = median_utils.median1d(np.ascontiguousarray(d.ravel(), dtype=np.float32),
np.ascontiguousarray(median_mask, dtype=np.uint8))
else:
nx = d.shape[axis]
ny = d.size // nx
output_shape = np.delete(d.shape, axis)
if mask is not None:
median_mask = np.rollaxis(mask, axis, len(d.shape)).reshape(ny, nx).astype(np.uint8)
else:
median_mask = np.zeros((ny, nx), dtype=np.uint8)
med = median_utils.median2d(np.ascontiguousarray(np.rollaxis(d, axis, len(d.shape)).reshape(ny, nx), dtype=np.float32),
mask=np.ascontiguousarray(median_mask, dtype=np.uint8))
median_array = np.array(med)
output_median = median_array.reshape(output_shape)
return output_median