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
def test_median1d_all_masks_returns_zero():
size = 1000
a = np.zeros(size, dtype=np.float32)
mask = np.ones(size, dtype=np.uint8)
assert median_utils.median1d(a, mask) == 0.0
def _compare_median1d(a, mask):
a_copy = a.copy()
actual = median_utils.median1d(a.astype('f4'), mask=mask.astype(np.uint8))
expected = np.median(a_copy[mask == 0].astype('f4'))
assert actual == np.float32(expected)
def test_median1d_all_masks_returns_zero():
size = 1000
a = np.zeros(size, dtype=np.float32)
mask = np.ones(size, dtype=np.uint8)
assert median_utils.median1d(a, mask) == 0.0
def _compare_median1d(a, mask):
a_copy = a.copy()
actual = median_utils.median1d(a.astype('f4'), mask=mask.astype(np.uint8))
expected = np.median(a_copy[mask == 0].astype('f4'))
assert actual == np.float32(expected)