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