def choose(a, choices, out=None, mode='raise'): """Construct an array from an index array and a set of arrays to choose from. Like `numpy.choose`. Masked indices in ``a`` will lead to masked output values and underlying data values are ignored if out of bounds (for ``mode='raise'``). Any values masked in ``choices`` will be propagated if chosen. """ from astropy.utils.masked import Masked a_data, a_mask = Masked._get_data_and_mask(a) if a_mask is not None and mode == 'raise': # Avoid raising on masked indices. a_data = a.filled(fill_value=0) kwargs = {'mode': mode} if out is not None: if not isinstance(out, Masked): raise NotImplementedError kwargs['out'] = out.unmasked data, masks = _get_data_and_masks(*choices) data_chosen = np.choose(a_data, data, **kwargs) if out is not None: kwargs['out'] = out.mask mask_chosen = np.choose(a_data, masks, **kwargs) if a_mask is not None: mask_chosen |= a_mask return Masked(data_chosen, mask_chosen) if out is None else out
def where(condition, *args): from astropy.utils.masked import Masked if not args: return condition.nonzero(), None, None condition, c_mask = Masked._get_data_and_mask(condition) data, masks = _get_data_and_masks(*args) unmasked = np.where(condition, *data) mask = np.where(condition, *masks) if c_mask is not None: mask |= c_mask return Masked(unmasked, mask=mask)
def nanfunc(a, *args, **kwargs): from astropy.utils.masked import Masked a, mask = Masked._get_data_and_mask(a) if issubclass(a.dtype.type, np.inexact): nans = np.isnan(a) mask = nans if mask is None else (nans | mask) if mask is not None: a = Masked(a, mask) if fill_value is not None: a = a.filled(fill_value) return np_func(a, *args, **kwargs)
def interp(x, xp, fp, *args, **kwargs): """One-dimensional linear interpolation. Like `numpy.interp`, but any masked points in ``xp`` and ``fp`` are ignored. Any masked values in ``x`` will still be evaluated, but masked on output. """ from astropy.utils.masked import Masked xd, xm = Masked._get_data_and_mask(x) if isinstance(xp, Masked) or isinstance(fp, Masked): (xp, fp), (xpm, fpm) = _get_data_and_masks(xp, fp) if xp.ndim == fp.ndim == 1: # Avoid making arrays 1-D; will just raise below. m = xpm | fpm xp = xp[~m] fp = fp[~m] result = np.interp(xd, xp, fp, *args, **kwargs) return result if xm is None else Masked(result, xm.copy())
def bincount(x, weights=None, minlength=0): """Count number of occurrences of each value in array of non-negative ints. Like `numpy.bincount`, but masked entries in ``x`` will be skipped. Any masked entries in ``weights`` will lead the corresponding bin to be masked. """ from astropy.utils.masked import Masked if weights is not None: weights = np.asanyarray(weights) if isinstance(x, Masked) and x.ndim <= 1: # let other dimensions lead to errors. if weights is not None and weights.ndim == x.ndim: weights = weights[~x.mask] x = x.unmasked[~x.mask] mask = None if weights is not None: weights, w_mask = Masked._get_data_and_mask(weights) if w_mask is not None: mask = np.bincount(x, w_mask.astype(int), minlength=minlength).astype(bool) result = np.bincount(x, weights, minlength=0) return result, mask, None