Chris@87: """ Chris@87: Masked arrays add-ons. Chris@87: Chris@87: A collection of utilities for `numpy.ma`. Chris@87: Chris@87: :author: Pierre Gerard-Marchant Chris@87: :contact: pierregm_at_uga_dot_edu Chris@87: :version: $Id: extras.py 3473 2007-10-29 15:18:13Z jarrod.millman $ Chris@87: Chris@87: """ Chris@87: from __future__ import division, absolute_import, print_function Chris@87: Chris@87: __author__ = "Pierre GF Gerard-Marchant ($Author: jarrod.millman $)" Chris@87: __version__ = '1.0' Chris@87: __revision__ = "$Revision: 3473 $" Chris@87: __date__ = '$Date: 2007-10-29 17:18:13 +0200 (Mon, 29 Oct 2007) $' Chris@87: Chris@87: __all__ = ['apply_along_axis', 'apply_over_axes', 'atleast_1d', 'atleast_2d', Chris@87: 'atleast_3d', 'average', Chris@87: 'clump_masked', 'clump_unmasked', 'column_stack', 'compress_cols', Chris@87: 'compress_rowcols', 'compress_rows', 'count_masked', 'corrcoef', Chris@87: 'cov', Chris@87: 'diagflat', 'dot', 'dstack', Chris@87: 'ediff1d', Chris@87: 'flatnotmasked_contiguous', 'flatnotmasked_edges', Chris@87: 'hsplit', 'hstack', Chris@87: 'in1d', 'intersect1d', Chris@87: 'mask_cols', 'mask_rowcols', 'mask_rows', 'masked_all', Chris@87: 'masked_all_like', 'median', 'mr_', Chris@87: 'notmasked_contiguous', 'notmasked_edges', Chris@87: 'polyfit', Chris@87: 'row_stack', Chris@87: 'setdiff1d', 'setxor1d', Chris@87: 'unique', 'union1d', Chris@87: 'vander', 'vstack', Chris@87: ] Chris@87: Chris@87: import itertools Chris@87: import warnings Chris@87: Chris@87: from . import core as ma Chris@87: from .core import MaskedArray, MAError, add, array, asarray, concatenate, count, \ Chris@87: filled, getmask, getmaskarray, make_mask_descr, masked, masked_array, \ Chris@87: mask_or, nomask, ones, sort, zeros Chris@87: #from core import * Chris@87: Chris@87: import numpy as np Chris@87: from numpy import ndarray, array as nxarray Chris@87: import numpy.core.umath as umath Chris@87: from numpy.lib.index_tricks import AxisConcatenator Chris@87: from numpy.linalg import lstsq Chris@87: Chris@87: Chris@87: #............................................................................... Chris@87: def issequence(seq): Chris@87: """Is seq a sequence (ndarray, list or tuple)?""" Chris@87: if isinstance(seq, (ndarray, tuple, list)): Chris@87: return True Chris@87: return False Chris@87: Chris@87: def count_masked(arr, axis=None): Chris@87: """ Chris@87: Count the number of masked elements along the given axis. Chris@87: Chris@87: Parameters Chris@87: ---------- Chris@87: arr : array_like Chris@87: An array with (possibly) masked elements. Chris@87: axis : int, optional Chris@87: Axis along which to count. If None (default), a flattened Chris@87: version of the array is used. Chris@87: Chris@87: Returns Chris@87: ------- Chris@87: count : int, ndarray Chris@87: The total number of masked elements (axis=None) or the number Chris@87: of masked elements along each slice of the given axis. Chris@87: Chris@87: See Also Chris@87: -------- Chris@87: MaskedArray.count : Count non-masked elements. Chris@87: Chris@87: Examples Chris@87: -------- Chris@87: >>> import numpy.ma as ma Chris@87: >>> a = np.arange(9).reshape((3,3)) Chris@87: >>> a = ma.array(a) Chris@87: >>> a[1, 0] = ma.masked Chris@87: >>> a[1, 2] = ma.masked Chris@87: >>> a[2, 1] = ma.masked Chris@87: >>> a Chris@87: masked_array(data = Chris@87: [[0 1 2] Chris@87: [-- 4 --] Chris@87: [6 -- 8]], Chris@87: mask = Chris@87: [[False False False] Chris@87: [ True False True] Chris@87: [False True False]], Chris@87: fill_value=999999) Chris@87: >>> ma.count_masked(a) Chris@87: 3 Chris@87: Chris@87: When the `axis` keyword is used an array is returned. Chris@87: Chris@87: >>> ma.count_masked(a, axis=0) Chris@87: array([1, 1, 1]) Chris@87: >>> ma.count_masked(a, axis=1) Chris@87: array([0, 2, 1]) Chris@87: Chris@87: """ Chris@87: m = getmaskarray(arr) Chris@87: return m.sum(axis) Chris@87: Chris@87: def masked_all(shape, dtype=float): Chris@87: """ Chris@87: Empty masked array with all elements masked. Chris@87: Chris@87: Return an empty masked array of the given shape and dtype, where all the Chris@87: data are masked. Chris@87: Chris@87: Parameters Chris@87: ---------- Chris@87: shape : tuple Chris@87: Shape of the required MaskedArray. Chris@87: dtype : dtype, optional Chris@87: Data type of the output. Chris@87: Chris@87: Returns Chris@87: ------- Chris@87: a : MaskedArray Chris@87: A masked array with all data masked. Chris@87: Chris@87: See Also Chris@87: -------- Chris@87: masked_all_like : Empty masked array modelled on an existing array. Chris@87: Chris@87: Examples Chris@87: -------- Chris@87: >>> import numpy.ma as ma Chris@87: >>> ma.masked_all((3, 3)) Chris@87: masked_array(data = Chris@87: [[-- -- --] Chris@87: [-- -- --] Chris@87: [-- -- --]], Chris@87: mask = Chris@87: [[ True True True] Chris@87: [ True True True] Chris@87: [ True True True]], Chris@87: fill_value=1e+20) Chris@87: Chris@87: The `dtype` parameter defines the underlying data type. Chris@87: Chris@87: >>> a = ma.masked_all((3, 3)) Chris@87: >>> a.dtype Chris@87: dtype('float64') Chris@87: >>> a = ma.masked_all((3, 3), dtype=np.int32) Chris@87: >>> a.dtype Chris@87: dtype('int32') Chris@87: Chris@87: """ Chris@87: a = masked_array(np.empty(shape, dtype), Chris@87: mask=np.ones(shape, make_mask_descr(dtype))) Chris@87: return a Chris@87: Chris@87: def masked_all_like(arr): Chris@87: """ Chris@87: Empty masked array with the properties of an existing array. Chris@87: Chris@87: Return an empty masked array of the same shape and dtype as Chris@87: the array `arr`, where all the data are masked. Chris@87: Chris@87: Parameters Chris@87: ---------- Chris@87: arr : ndarray Chris@87: An array describing the shape and dtype of the required MaskedArray. Chris@87: Chris@87: Returns Chris@87: ------- Chris@87: a : MaskedArray Chris@87: A masked array with all data masked. Chris@87: Chris@87: Raises Chris@87: ------ Chris@87: AttributeError Chris@87: If `arr` doesn't have a shape attribute (i.e. not an ndarray) Chris@87: Chris@87: See Also Chris@87: -------- Chris@87: masked_all : Empty masked array with all elements masked. Chris@87: Chris@87: Examples Chris@87: -------- Chris@87: >>> import numpy.ma as ma Chris@87: >>> arr = np.zeros((2, 3), dtype=np.float32) Chris@87: >>> arr Chris@87: array([[ 0., 0., 0.], Chris@87: [ 0., 0., 0.]], dtype=float32) Chris@87: >>> ma.masked_all_like(arr) Chris@87: masked_array(data = Chris@87: [[-- -- --] Chris@87: [-- -- --]], Chris@87: mask = Chris@87: [[ True True True] Chris@87: [ True True True]], Chris@87: fill_value=1e+20) Chris@87: Chris@87: The dtype of the masked array matches the dtype of `arr`. Chris@87: Chris@87: >>> arr.dtype Chris@87: dtype('float32') Chris@87: >>> ma.masked_all_like(arr).dtype Chris@87: dtype('float32') Chris@87: Chris@87: """ Chris@87: a = np.empty_like(arr).view(MaskedArray) Chris@87: a._mask = np.ones(a.shape, dtype=make_mask_descr(a.dtype)) Chris@87: return a Chris@87: Chris@87: Chris@87: #####-------------------------------------------------------------------------- Chris@87: #---- --- Standard functions --- Chris@87: #####-------------------------------------------------------------------------- Chris@87: class _fromnxfunction: Chris@87: """ Chris@87: Defines a wrapper to adapt NumPy functions to masked arrays. Chris@87: Chris@87: Chris@87: An instance of `_fromnxfunction` can be called with the same parameters Chris@87: as the wrapped NumPy function. The docstring of `newfunc` is adapted from Chris@87: the wrapped function as well, see `getdoc`. Chris@87: Chris@87: Parameters Chris@87: ---------- Chris@87: funcname : str Chris@87: The name of the function to be adapted. The function should be Chris@87: in the NumPy namespace (i.e. ``np.funcname``). Chris@87: Chris@87: """ Chris@87: Chris@87: def __init__(self, funcname): Chris@87: self.__name__ = funcname Chris@87: self.__doc__ = self.getdoc() Chris@87: Chris@87: def getdoc(self): Chris@87: """ Chris@87: Retrieve the docstring and signature from the function. Chris@87: Chris@87: The ``__doc__`` attribute of the function is used as the docstring for Chris@87: the new masked array version of the function. A note on application Chris@87: of the function to the mask is appended. Chris@87: Chris@87: .. warning:: Chris@87: If the function docstring already contained a Notes section, the Chris@87: new docstring will have two Notes sections instead of appending a note Chris@87: to the existing section. Chris@87: Chris@87: Parameters Chris@87: ---------- Chris@87: None Chris@87: Chris@87: """ Chris@87: npfunc = getattr(np, self.__name__, None) Chris@87: doc = getattr(npfunc, '__doc__', None) Chris@87: if doc: Chris@87: sig = self.__name__ + ma.get_object_signature(npfunc) Chris@87: locdoc = "Notes\n-----\nThe function is applied to both the _data"\ Chris@87: " and the _mask, if any." Chris@87: return '\n'.join((sig, doc, locdoc)) Chris@87: return Chris@87: Chris@87: Chris@87: def __call__(self, *args, **params): Chris@87: func = getattr(np, self.__name__) Chris@87: if len(args) == 1: Chris@87: x = args[0] Chris@87: if isinstance(x, ndarray): Chris@87: _d = func(x.__array__(), **params) Chris@87: _m = func(getmaskarray(x), **params) Chris@87: return masked_array(_d, mask=_m) Chris@87: elif isinstance(x, tuple) or isinstance(x, list): Chris@87: _d = func(tuple([np.asarray(a) for a in x]), **params) Chris@87: _m = func(tuple([getmaskarray(a) for a in x]), **params) Chris@87: return masked_array(_d, mask=_m) Chris@87: else: Chris@87: arrays = [] Chris@87: args = list(args) Chris@87: while len(args) > 0 and issequence(args[0]): Chris@87: arrays.append(args.pop(0)) Chris@87: res = [] Chris@87: for x in arrays: Chris@87: _d = func(np.asarray(x), *args, **params) Chris@87: _m = func(getmaskarray(x), *args, **params) Chris@87: res.append(masked_array(_d, mask=_m)) Chris@87: return res Chris@87: Chris@87: atleast_1d = _fromnxfunction('atleast_1d') Chris@87: atleast_2d = _fromnxfunction('atleast_2d') Chris@87: atleast_3d = _fromnxfunction('atleast_3d') Chris@87: #atleast_1d = np.atleast_1d Chris@87: #atleast_2d = np.atleast_2d Chris@87: #atleast_3d = np.atleast_3d Chris@87: Chris@87: vstack = row_stack = _fromnxfunction('vstack') Chris@87: hstack = _fromnxfunction('hstack') Chris@87: column_stack = _fromnxfunction('column_stack') Chris@87: dstack = _fromnxfunction('dstack') Chris@87: Chris@87: hsplit = _fromnxfunction('hsplit') Chris@87: Chris@87: diagflat = _fromnxfunction('diagflat') Chris@87: Chris@87: Chris@87: #####-------------------------------------------------------------------------- Chris@87: #---- Chris@87: #####-------------------------------------------------------------------------- Chris@87: def flatten_inplace(seq): Chris@87: """Flatten a sequence in place.""" Chris@87: k = 0 Chris@87: while (k != len(seq)): Chris@87: while hasattr(seq[k], '__iter__'): Chris@87: seq[k:(k + 1)] = seq[k] Chris@87: k += 1 Chris@87: return seq Chris@87: Chris@87: Chris@87: def apply_along_axis(func1d, axis, arr, *args, **kwargs): Chris@87: """ Chris@87: (This docstring should be overwritten) Chris@87: """ Chris@87: arr = array(arr, copy=False, subok=True) Chris@87: nd = arr.ndim Chris@87: if axis < 0: Chris@87: axis += nd Chris@87: if (axis >= nd): Chris@87: raise ValueError("axis must be less than arr.ndim; axis=%d, rank=%d." Chris@87: % (axis, nd)) Chris@87: ind = [0] * (nd - 1) Chris@87: i = np.zeros(nd, 'O') Chris@87: indlist = list(range(nd)) Chris@87: indlist.remove(axis) Chris@87: i[axis] = slice(None, None) Chris@87: outshape = np.asarray(arr.shape).take(indlist) Chris@87: i.put(indlist, ind) Chris@87: j = i.copy() Chris@87: res = func1d(arr[tuple(i.tolist())], *args, **kwargs) Chris@87: # if res is a number, then we have a smaller output array Chris@87: asscalar = np.isscalar(res) Chris@87: if not asscalar: Chris@87: try: Chris@87: len(res) Chris@87: except TypeError: Chris@87: asscalar = True Chris@87: # Note: we shouldn't set the dtype of the output from the first result... Chris@87: #...so we force the type to object, and build a list of dtypes Chris@87: #...we'll just take the largest, to avoid some downcasting Chris@87: dtypes = [] Chris@87: if asscalar: Chris@87: dtypes.append(np.asarray(res).dtype) Chris@87: outarr = zeros(outshape, object) Chris@87: outarr[tuple(ind)] = res Chris@87: Ntot = np.product(outshape) Chris@87: k = 1 Chris@87: while k < Ntot: Chris@87: # increment the index Chris@87: ind[-1] += 1 Chris@87: n = -1 Chris@87: while (ind[n] >= outshape[n]) and (n > (1 - nd)): Chris@87: ind[n - 1] += 1 Chris@87: ind[n] = 0 Chris@87: n -= 1 Chris@87: i.put(indlist, ind) Chris@87: res = func1d(arr[tuple(i.tolist())], *args, **kwargs) Chris@87: outarr[tuple(ind)] = res Chris@87: dtypes.append(asarray(res).dtype) Chris@87: k += 1 Chris@87: else: Chris@87: res = array(res, copy=False, subok=True) Chris@87: j = i.copy() Chris@87: j[axis] = ([slice(None, None)] * res.ndim) Chris@87: j.put(indlist, ind) Chris@87: Ntot = np.product(outshape) Chris@87: holdshape = outshape Chris@87: outshape = list(arr.shape) Chris@87: outshape[axis] = res.shape Chris@87: dtypes.append(asarray(res).dtype) Chris@87: outshape = flatten_inplace(outshape) Chris@87: outarr = zeros(outshape, object) Chris@87: outarr[tuple(flatten_inplace(j.tolist()))] = res Chris@87: k = 1 Chris@87: while k < Ntot: Chris@87: # increment the index Chris@87: ind[-1] += 1 Chris@87: n = -1 Chris@87: while (ind[n] >= holdshape[n]) and (n > (1 - nd)): Chris@87: ind[n - 1] += 1 Chris@87: ind[n] = 0 Chris@87: n -= 1 Chris@87: i.put(indlist, ind) Chris@87: j.put(indlist, ind) Chris@87: res = func1d(arr[tuple(i.tolist())], *args, **kwargs) Chris@87: outarr[tuple(flatten_inplace(j.tolist()))] = res Chris@87: dtypes.append(asarray(res).dtype) Chris@87: k += 1 Chris@87: max_dtypes = np.dtype(np.asarray(dtypes).max()) Chris@87: if not hasattr(arr, '_mask'): Chris@87: result = np.asarray(outarr, dtype=max_dtypes) Chris@87: else: Chris@87: result = asarray(outarr, dtype=max_dtypes) Chris@87: result.fill_value = ma.default_fill_value(result) Chris@87: return result Chris@87: apply_along_axis.__doc__ = np.apply_along_axis.__doc__ Chris@87: Chris@87: Chris@87: def apply_over_axes(func, a, axes): Chris@87: """ Chris@87: (This docstring will be overwritten) Chris@87: """ Chris@87: val = asarray(a) Chris@87: N = a.ndim Chris@87: if array(axes).ndim == 0: Chris@87: axes = (axes,) Chris@87: for axis in axes: Chris@87: if axis < 0: axis = N + axis Chris@87: args = (val, axis) Chris@87: res = func(*args) Chris@87: if res.ndim == val.ndim: Chris@87: val = res Chris@87: else: Chris@87: res = ma.expand_dims(res, axis) Chris@87: if res.ndim == val.ndim: Chris@87: val = res Chris@87: else: Chris@87: raise ValueError("function is not returning " Chris@87: "an array of the correct shape") Chris@87: return val Chris@87: Chris@87: if apply_over_axes.__doc__ is not None: Chris@87: apply_over_axes.__doc__ = np.apply_over_axes.__doc__[ Chris@87: :np.apply_over_axes.__doc__.find('Notes')].rstrip() + \ Chris@87: """ Chris@87: Chris@87: Examples Chris@87: -------- Chris@87: >>> a = ma.arange(24).reshape(2,3,4) Chris@87: >>> a[:,0,1] = ma.masked Chris@87: >>> a[:,1,:] = ma.masked Chris@87: >>> print a Chris@87: [[[0 -- 2 3] Chris@87: [-- -- -- --] Chris@87: [8 9 10 11]] Chris@87: Chris@87: [[12 -- 14 15] Chris@87: [-- -- -- --] Chris@87: [20 21 22 23]]] Chris@87: >>> print ma.apply_over_axes(ma.sum, a, [0,2]) Chris@87: [[[46] Chris@87: [--] Chris@87: [124]]] Chris@87: Chris@87: Tuple axis arguments to ufuncs are equivalent: Chris@87: Chris@87: >>> print ma.sum(a, axis=(0,2)).reshape((1,-1,1)) Chris@87: [[[46] Chris@87: [--] Chris@87: [124]]] Chris@87: """ Chris@87: Chris@87: Chris@87: def average(a, axis=None, weights=None, returned=False): Chris@87: """ Chris@87: Return the weighted average of array over the given axis. Chris@87: Chris@87: Parameters Chris@87: ---------- Chris@87: a : array_like Chris@87: Data to be averaged. Chris@87: Masked entries are not taken into account in the computation. Chris@87: axis : int, optional Chris@87: Axis along which the average is computed. The default is to compute Chris@87: the average of the flattened array. Chris@87: weights : array_like, optional Chris@87: The importance that each element has in the computation of the average. Chris@87: The weights array can either be 1-D (in which case its length must be Chris@87: the size of `a` along the given axis) or of the same shape as `a`. Chris@87: If ``weights=None``, then all data in `a` are assumed to have a Chris@87: weight equal to one. If `weights` is complex, the imaginary parts Chris@87: are ignored. Chris@87: returned : bool, optional Chris@87: Flag indicating whether a tuple ``(result, sum of weights)`` Chris@87: should be returned as output (True), or just the result (False). Chris@87: Default is False. Chris@87: Chris@87: Returns Chris@87: ------- Chris@87: average, [sum_of_weights] : (tuple of) scalar or MaskedArray Chris@87: The average along the specified axis. When returned is `True`, Chris@87: return a tuple with the average as the first element and the sum Chris@87: of the weights as the second element. The return type is `np.float64` Chris@87: if `a` is of integer type, otherwise it is of the same type as `a`. Chris@87: If returned, `sum_of_weights` is of the same type as `average`. Chris@87: Chris@87: Examples Chris@87: -------- Chris@87: >>> a = np.ma.array([1., 2., 3., 4.], mask=[False, False, True, True]) Chris@87: >>> np.ma.average(a, weights=[3, 1, 0, 0]) Chris@87: 1.25 Chris@87: Chris@87: >>> x = np.ma.arange(6.).reshape(3, 2) Chris@87: >>> print x Chris@87: [[ 0. 1.] Chris@87: [ 2. 3.] Chris@87: [ 4. 5.]] Chris@87: >>> avg, sumweights = np.ma.average(x, axis=0, weights=[1, 2, 3], Chris@87: ... returned=True) Chris@87: >>> print avg Chris@87: [2.66666666667 3.66666666667] Chris@87: Chris@87: """ Chris@87: a = asarray(a) Chris@87: mask = a.mask Chris@87: ash = a.shape Chris@87: if ash == (): Chris@87: ash = (1,) Chris@87: if axis is None: Chris@87: if mask is nomask: Chris@87: if weights is None: Chris@87: n = a.sum(axis=None) Chris@87: d = float(a.size) Chris@87: else: Chris@87: w = filled(weights, 0.0).ravel() Chris@87: n = umath.add.reduce(a._data.ravel() * w) Chris@87: d = umath.add.reduce(w) Chris@87: del w Chris@87: else: Chris@87: if weights is None: Chris@87: n = a.filled(0).sum(axis=None) Chris@87: d = float(umath.add.reduce((~mask).ravel())) Chris@87: else: Chris@87: w = array(filled(weights, 0.0), float, mask=mask).ravel() Chris@87: n = add.reduce(a.ravel() * w) Chris@87: d = add.reduce(w) Chris@87: del w Chris@87: else: Chris@87: if mask is nomask: Chris@87: if weights is None: Chris@87: d = ash[axis] * 1.0 Chris@87: n = add.reduce(a._data, axis) Chris@87: else: Chris@87: w = filled(weights, 0.0) Chris@87: wsh = w.shape Chris@87: if wsh == (): Chris@87: wsh = (1,) Chris@87: if wsh == ash: Chris@87: w = np.array(w, float, copy=0) Chris@87: n = add.reduce(a * w, axis) Chris@87: d = add.reduce(w, axis) Chris@87: del w Chris@87: elif wsh == (ash[axis],): Chris@87: ni = ash[axis] Chris@87: r = [None] * len(ash) Chris@87: r[axis] = slice(None, None, 1) Chris@87: w = eval ("w[" + repr(tuple(r)) + "] * ones(ash, float)") Chris@87: n = add.reduce(a * w, axis) Chris@87: d = add.reduce(w, axis, dtype=float) Chris@87: del w, r Chris@87: else: Chris@87: raise ValueError('average: weights wrong shape.') Chris@87: else: Chris@87: if weights is None: Chris@87: n = add.reduce(a, axis) Chris@87: d = umath.add.reduce((~mask), axis=axis, dtype=float) Chris@87: else: Chris@87: w = filled(weights, 0.0) Chris@87: wsh = w.shape Chris@87: if wsh == (): Chris@87: wsh = (1,) Chris@87: if wsh == ash: Chris@87: w = array(w, dtype=float, mask=mask, copy=0) Chris@87: n = add.reduce(a * w, axis) Chris@87: d = add.reduce(w, axis, dtype=float) Chris@87: elif wsh == (ash[axis],): Chris@87: ni = ash[axis] Chris@87: r = [None] * len(ash) Chris@87: r[axis] = slice(None, None, 1) Chris@87: w = eval ("w[" + repr(tuple(r)) + \ Chris@87: "] * masked_array(ones(ash, float), mask)") Chris@87: n = add.reduce(a * w, axis) Chris@87: d = add.reduce(w, axis, dtype=float) Chris@87: else: Chris@87: raise ValueError('average: weights wrong shape.') Chris@87: del w Chris@87: if n is masked or d is masked: Chris@87: return masked Chris@87: result = n / d Chris@87: del n Chris@87: Chris@87: if isinstance(result, MaskedArray): Chris@87: if ((axis is None) or (axis == 0 and a.ndim == 1)) and \ Chris@87: (result.mask is nomask): Chris@87: result = result._data Chris@87: if returned: Chris@87: if not isinstance(d, MaskedArray): Chris@87: d = masked_array(d) Chris@87: if isinstance(d, ndarray) and (not d.shape == result.shape): Chris@87: d = ones(result.shape, dtype=float) * d Chris@87: if returned: Chris@87: return result, d Chris@87: else: Chris@87: return result Chris@87: Chris@87: Chris@87: def median(a, axis=None, out=None, overwrite_input=False): Chris@87: """ Chris@87: Compute the median along the specified axis. Chris@87: Chris@87: Returns the median of the array elements. Chris@87: Chris@87: Parameters Chris@87: ---------- Chris@87: a : array_like Chris@87: Input array or object that can be converted to an array. Chris@87: axis : int, optional Chris@87: Axis along which the medians are computed. The default (None) is Chris@87: to compute the median along a flattened version of the array. Chris@87: out : ndarray, optional Chris@87: Alternative output array in which to place the result. It must Chris@87: have the same shape and buffer length as the expected output Chris@87: but the type will be cast if necessary. Chris@87: overwrite_input : bool, optional Chris@87: If True, then allow use of memory of input array (a) for Chris@87: calculations. The input array will be modified by the call to Chris@87: median. This will save memory when you do not need to preserve Chris@87: the contents of the input array. Treat the input as undefined, Chris@87: but it will probably be fully or partially sorted. Default is Chris@87: False. Note that, if `overwrite_input` is True, and the input Chris@87: is not already an `ndarray`, an error will be raised. Chris@87: Chris@87: Returns Chris@87: ------- Chris@87: median : ndarray Chris@87: A new array holding the result is returned unless out is Chris@87: specified, in which case a reference to out is returned. Chris@87: Return data-type is `float64` for integers and floats smaller than Chris@87: `float64`, or the input data-type, otherwise. Chris@87: Chris@87: See Also Chris@87: -------- Chris@87: mean Chris@87: Chris@87: Notes Chris@87: ----- Chris@87: Given a vector ``V`` with ``N`` non masked values, the median of ``V`` Chris@87: is the middle value of a sorted copy of ``V`` (``Vs``) - i.e. Chris@87: ``Vs[(N-1)/2]``, when ``N`` is odd, or ``{Vs[N/2 - 1] + Vs[N/2]}/2`` Chris@87: when ``N`` is even. Chris@87: Chris@87: Examples Chris@87: -------- Chris@87: >>> x = np.ma.array(np.arange(8), mask=[0]*4 + [1]*4) Chris@87: >>> np.ma.extras.median(x) Chris@87: 1.5 Chris@87: Chris@87: >>> x = np.ma.array(np.arange(10).reshape(2, 5), mask=[0]*6 + [1]*4) Chris@87: >>> np.ma.extras.median(x) Chris@87: 2.5 Chris@87: >>> np.ma.extras.median(x, axis=-1, overwrite_input=True) Chris@87: masked_array(data = [ 2. 5.], Chris@87: mask = False, Chris@87: fill_value = 1e+20) Chris@87: Chris@87: """ Chris@87: if not hasattr(a, 'mask') or np.count_nonzero(a.mask) == 0: Chris@87: return masked_array(np.median(getattr(a, 'data', a), axis=axis, Chris@87: out=out, overwrite_input=overwrite_input), copy=False) Chris@87: if overwrite_input: Chris@87: if axis is None: Chris@87: asorted = a.ravel() Chris@87: asorted.sort() Chris@87: else: Chris@87: a.sort(axis=axis) Chris@87: asorted = a Chris@87: else: Chris@87: asorted = sort(a, axis=axis) Chris@87: if axis is None: Chris@87: axis = 0 Chris@87: elif axis < 0: Chris@87: axis += a.ndim Chris@87: Chris@87: counts = asorted.shape[axis] - (asorted.mask).sum(axis=axis) Chris@87: h = counts // 2 Chris@87: # create indexing mesh grid for all but reduced axis Chris@87: axes_grid = [np.arange(x) for i, x in enumerate(asorted.shape) Chris@87: if i != axis] Chris@87: ind = np.meshgrid(*axes_grid, sparse=True, indexing='ij') Chris@87: # insert indices of low and high median Chris@87: ind.insert(axis, h - 1) Chris@87: low = asorted[ind] Chris@87: ind[axis] = h Chris@87: high = asorted[ind] Chris@87: # duplicate high if odd number of elements so mean does nothing Chris@87: odd = counts % 2 == 1 Chris@87: if asorted.ndim == 1: Chris@87: if odd: Chris@87: low = high Chris@87: else: Chris@87: low[odd] = high[odd] Chris@87: Chris@87: if np.issubdtype(asorted.dtype, np.inexact): Chris@87: # avoid inf / x = masked Chris@87: s = np.ma.sum([low, high], axis=0, out=out) Chris@87: np.true_divide(s.data, 2., casting='unsafe', out=s.data) Chris@87: else: Chris@87: s = np.ma.mean([low, high], axis=0, out=out) Chris@87: return s Chris@87: Chris@87: Chris@87: #.............................................................................. Chris@87: def compress_rowcols(x, axis=None): Chris@87: """ Chris@87: Suppress the rows and/or columns of a 2-D array that contain Chris@87: masked values. Chris@87: Chris@87: The suppression behavior is selected with the `axis` parameter. Chris@87: Chris@87: - If axis is None, both rows and columns are suppressed. Chris@87: - If axis is 0, only rows are suppressed. Chris@87: - If axis is 1 or -1, only columns are suppressed. Chris@87: Chris@87: Parameters Chris@87: ---------- Chris@87: axis : int, optional Chris@87: Axis along which to perform the operation. Default is None. Chris@87: Chris@87: Returns Chris@87: ------- Chris@87: compressed_array : ndarray Chris@87: The compressed array. Chris@87: Chris@87: Examples Chris@87: -------- Chris@87: >>> x = np.ma.array(np.arange(9).reshape(3, 3), mask=[[1, 0, 0], Chris@87: ... [1, 0, 0], Chris@87: ... [0, 0, 0]]) Chris@87: >>> x Chris@87: masked_array(data = Chris@87: [[-- 1 2] Chris@87: [-- 4 5] Chris@87: [6 7 8]], Chris@87: mask = Chris@87: [[ True False False] Chris@87: [ True False False] Chris@87: [False False False]], Chris@87: fill_value = 999999) Chris@87: Chris@87: >>> np.ma.extras.compress_rowcols(x) Chris@87: array([[7, 8]]) Chris@87: >>> np.ma.extras.compress_rowcols(x, 0) Chris@87: array([[6, 7, 8]]) Chris@87: >>> np.ma.extras.compress_rowcols(x, 1) Chris@87: array([[1, 2], Chris@87: [4, 5], Chris@87: [7, 8]]) Chris@87: Chris@87: """ Chris@87: x = asarray(x) Chris@87: if x.ndim != 2: Chris@87: raise NotImplementedError("compress2d works for 2D arrays only.") Chris@87: m = getmask(x) Chris@87: # Nothing is masked: return x Chris@87: if m is nomask or not m.any(): Chris@87: return x._data Chris@87: # All is masked: return empty Chris@87: if m.all(): Chris@87: return nxarray([]) Chris@87: # Builds a list of rows/columns indices Chris@87: (idxr, idxc) = (list(range(len(x))), list(range(x.shape[1]))) Chris@87: masked = m.nonzero() Chris@87: if not axis: Chris@87: for i in np.unique(masked[0]): Chris@87: idxr.remove(i) Chris@87: if axis in [None, 1, -1]: Chris@87: for j in np.unique(masked[1]): Chris@87: idxc.remove(j) Chris@87: return x._data[idxr][:, idxc] Chris@87: Chris@87: def compress_rows(a): Chris@87: """ Chris@87: Suppress whole rows of a 2-D array that contain masked values. Chris@87: Chris@87: This is equivalent to ``np.ma.extras.compress_rowcols(a, 0)``, see Chris@87: `extras.compress_rowcols` for details. Chris@87: Chris@87: See Also Chris@87: -------- Chris@87: extras.compress_rowcols Chris@87: Chris@87: """ Chris@87: return compress_rowcols(a, 0) Chris@87: Chris@87: def compress_cols(a): Chris@87: """ Chris@87: Suppress whole columns of a 2-D array that contain masked values. Chris@87: Chris@87: This is equivalent to ``np.ma.extras.compress_rowcols(a, 1)``, see Chris@87: `extras.compress_rowcols` for details. Chris@87: Chris@87: See Also Chris@87: -------- Chris@87: extras.compress_rowcols Chris@87: Chris@87: """ Chris@87: return compress_rowcols(a, 1) Chris@87: Chris@87: def mask_rowcols(a, axis=None): Chris@87: """ Chris@87: Mask rows and/or columns of a 2D array that contain masked values. Chris@87: Chris@87: Mask whole rows and/or columns of a 2D array that contain Chris@87: masked values. The masking behavior is selected using the Chris@87: `axis` parameter. Chris@87: Chris@87: - If `axis` is None, rows *and* columns are masked. Chris@87: - If `axis` is 0, only rows are masked. Chris@87: - If `axis` is 1 or -1, only columns are masked. Chris@87: Chris@87: Parameters Chris@87: ---------- Chris@87: a : array_like, MaskedArray Chris@87: The array to mask. If not a MaskedArray instance (or if no array Chris@87: elements are masked). The result is a MaskedArray with `mask` set Chris@87: to `nomask` (False). Must be a 2D array. Chris@87: axis : int, optional Chris@87: Axis along which to perform the operation. If None, applies to a Chris@87: flattened version of the array. Chris@87: Chris@87: Returns Chris@87: ------- Chris@87: a : MaskedArray Chris@87: A modified version of the input array, masked depending on the value Chris@87: of the `axis` parameter. Chris@87: Chris@87: Raises Chris@87: ------ Chris@87: NotImplementedError Chris@87: If input array `a` is not 2D. Chris@87: Chris@87: See Also Chris@87: -------- Chris@87: mask_rows : Mask rows of a 2D array that contain masked values. Chris@87: mask_cols : Mask cols of a 2D array that contain masked values. Chris@87: masked_where : Mask where a condition is met. Chris@87: Chris@87: Notes Chris@87: ----- Chris@87: The input array's mask is modified by this function. Chris@87: Chris@87: Examples Chris@87: -------- Chris@87: >>> import numpy.ma as ma Chris@87: >>> a = np.zeros((3, 3), dtype=np.int) Chris@87: >>> a[1, 1] = 1 Chris@87: >>> a Chris@87: array([[0, 0, 0], Chris@87: [0, 1, 0], Chris@87: [0, 0, 0]]) Chris@87: >>> a = ma.masked_equal(a, 1) Chris@87: >>> a Chris@87: masked_array(data = Chris@87: [[0 0 0] Chris@87: [0 -- 0] Chris@87: [0 0 0]], Chris@87: mask = Chris@87: [[False False False] Chris@87: [False True False] Chris@87: [False False False]], Chris@87: fill_value=999999) Chris@87: >>> ma.mask_rowcols(a) Chris@87: masked_array(data = Chris@87: [[0 -- 0] Chris@87: [-- -- --] Chris@87: [0 -- 0]], Chris@87: mask = Chris@87: [[False True False] Chris@87: [ True True True] Chris@87: [False True False]], Chris@87: fill_value=999999) Chris@87: Chris@87: """ Chris@87: a = array(a, subok=False) Chris@87: if a.ndim != 2: Chris@87: raise NotImplementedError("mask_rowcols works for 2D arrays only.") Chris@87: m = getmask(a) Chris@87: # Nothing is masked: return a Chris@87: if m is nomask or not m.any(): Chris@87: return a Chris@87: maskedval = m.nonzero() Chris@87: a._mask = a._mask.copy() Chris@87: if not axis: Chris@87: a[np.unique(maskedval[0])] = masked Chris@87: if axis in [None, 1, -1]: Chris@87: a[:, np.unique(maskedval[1])] = masked Chris@87: return a Chris@87: Chris@87: def mask_rows(a, axis=None): Chris@87: """ Chris@87: Mask rows of a 2D array that contain masked values. Chris@87: Chris@87: This function is a shortcut to ``mask_rowcols`` with `axis` equal to 0. Chris@87: Chris@87: See Also Chris@87: -------- Chris@87: mask_rowcols : Mask rows and/or columns of a 2D array. Chris@87: masked_where : Mask where a condition is met. Chris@87: Chris@87: Examples Chris@87: -------- Chris@87: >>> import numpy.ma as ma Chris@87: >>> a = np.zeros((3, 3), dtype=np.int) Chris@87: >>> a[1, 1] = 1 Chris@87: >>> a Chris@87: array([[0, 0, 0], Chris@87: [0, 1, 0], Chris@87: [0, 0, 0]]) Chris@87: >>> a = ma.masked_equal(a, 1) Chris@87: >>> a Chris@87: masked_array(data = Chris@87: [[0 0 0] Chris@87: [0 -- 0] Chris@87: [0 0 0]], Chris@87: mask = Chris@87: [[False False False] Chris@87: [False True False] Chris@87: [False False False]], Chris@87: fill_value=999999) Chris@87: >>> ma.mask_rows(a) Chris@87: masked_array(data = Chris@87: [[0 0 0] Chris@87: [-- -- --] Chris@87: [0 0 0]], Chris@87: mask = Chris@87: [[False False False] Chris@87: [ True True True] Chris@87: [False False False]], Chris@87: fill_value=999999) Chris@87: Chris@87: """ Chris@87: return mask_rowcols(a, 0) Chris@87: Chris@87: def mask_cols(a, axis=None): Chris@87: """ Chris@87: Mask columns of a 2D array that contain masked values. Chris@87: Chris@87: This function is a shortcut to ``mask_rowcols`` with `axis` equal to 1. Chris@87: Chris@87: See Also Chris@87: -------- Chris@87: mask_rowcols : Mask rows and/or columns of a 2D array. Chris@87: masked_where : Mask where a condition is met. Chris@87: Chris@87: Examples Chris@87: -------- Chris@87: >>> import numpy.ma as ma Chris@87: >>> a = np.zeros((3, 3), dtype=np.int) Chris@87: >>> a[1, 1] = 1 Chris@87: >>> a Chris@87: array([[0, 0, 0], Chris@87: [0, 1, 0], Chris@87: [0, 0, 0]]) Chris@87: >>> a = ma.masked_equal(a, 1) Chris@87: >>> a Chris@87: masked_array(data = Chris@87: [[0 0 0] Chris@87: [0 -- 0] Chris@87: [0 0 0]], Chris@87: mask = Chris@87: [[False False False] Chris@87: [False True False] Chris@87: [False False False]], Chris@87: fill_value=999999) Chris@87: >>> ma.mask_cols(a) Chris@87: masked_array(data = Chris@87: [[0 -- 0] Chris@87: [0 -- 0] Chris@87: [0 -- 0]], Chris@87: mask = Chris@87: [[False True False] Chris@87: [False True False] Chris@87: [False True False]], Chris@87: fill_value=999999) Chris@87: Chris@87: """ Chris@87: return mask_rowcols(a, 1) Chris@87: Chris@87: Chris@87: def dot(a, b, strict=False): Chris@87: """ Chris@87: Return the dot product of two arrays. Chris@87: Chris@87: .. note:: Chris@87: Works only with 2-D arrays at the moment. Chris@87: Chris@87: This function is the equivalent of `numpy.dot` that takes masked values Chris@87: into account, see `numpy.dot` for details. Chris@87: Chris@87: Parameters Chris@87: ---------- Chris@87: a, b : ndarray Chris@87: Inputs arrays. Chris@87: strict : bool, optional Chris@87: Whether masked data are propagated (True) or set to 0 (False) for the Chris@87: computation. Default is False. Chris@87: Propagating the mask means that if a masked value appears in a row or Chris@87: column, the whole row or column is considered masked. Chris@87: Chris@87: See Also Chris@87: -------- Chris@87: numpy.dot : Equivalent function for ndarrays. Chris@87: Chris@87: Examples Chris@87: -------- Chris@87: >>> a = ma.array([[1, 2, 3], [4, 5, 6]], mask=[[1, 0, 0], [0, 0, 0]]) Chris@87: >>> b = ma.array([[1, 2], [3, 4], [5, 6]], mask=[[1, 0], [0, 0], [0, 0]]) Chris@87: >>> np.ma.dot(a, b) Chris@87: masked_array(data = Chris@87: [[21 26] Chris@87: [45 64]], Chris@87: mask = Chris@87: [[False False] Chris@87: [False False]], Chris@87: fill_value = 999999) Chris@87: >>> np.ma.dot(a, b, strict=True) Chris@87: masked_array(data = Chris@87: [[-- --] Chris@87: [-- 64]], Chris@87: mask = Chris@87: [[ True True] Chris@87: [ True False]], Chris@87: fill_value = 999999) Chris@87: Chris@87: """ Chris@87: #!!!: Works only with 2D arrays. There should be a way to get it to run with higher dimension Chris@87: if strict and (a.ndim == 2) and (b.ndim == 2): Chris@87: a = mask_rows(a) Chris@87: b = mask_cols(b) Chris@87: # Chris@87: d = np.dot(filled(a, 0), filled(b, 0)) Chris@87: # Chris@87: am = (~getmaskarray(a)) Chris@87: bm = (~getmaskarray(b)) Chris@87: m = ~np.dot(am, bm) Chris@87: return masked_array(d, mask=m) Chris@87: Chris@87: #####-------------------------------------------------------------------------- Chris@87: #---- --- arraysetops --- Chris@87: #####-------------------------------------------------------------------------- Chris@87: Chris@87: def ediff1d(arr, to_end=None, to_begin=None): Chris@87: """ Chris@87: Compute the differences between consecutive elements of an array. Chris@87: Chris@87: This function is the equivalent of `numpy.ediff1d` that takes masked Chris@87: values into account, see `numpy.ediff1d` for details. Chris@87: Chris@87: See Also Chris@87: -------- Chris@87: numpy.ediff1d : Equivalent function for ndarrays. Chris@87: Chris@87: """ Chris@87: arr = ma.asanyarray(arr).flat Chris@87: ed = arr[1:] - arr[:-1] Chris@87: arrays = [ed] Chris@87: # Chris@87: if to_begin is not None: Chris@87: arrays.insert(0, to_begin) Chris@87: if to_end is not None: Chris@87: arrays.append(to_end) Chris@87: # Chris@87: if len(arrays) != 1: Chris@87: # We'll save ourselves a copy of a potentially large array in the common Chris@87: # case where neither to_begin or to_end was given. Chris@87: ed = hstack(arrays) Chris@87: # Chris@87: return ed Chris@87: Chris@87: Chris@87: def unique(ar1, return_index=False, return_inverse=False): Chris@87: """ Chris@87: Finds the unique elements of an array. Chris@87: Chris@87: Masked values are considered the same element (masked). The output array Chris@87: is always a masked array. See `numpy.unique` for more details. Chris@87: Chris@87: See Also Chris@87: -------- Chris@87: numpy.unique : Equivalent function for ndarrays. Chris@87: Chris@87: """ Chris@87: output = np.unique(ar1, Chris@87: return_index=return_index, Chris@87: return_inverse=return_inverse) Chris@87: if isinstance(output, tuple): Chris@87: output = list(output) Chris@87: output[0] = output[0].view(MaskedArray) Chris@87: output = tuple(output) Chris@87: else: Chris@87: output = output.view(MaskedArray) Chris@87: return output Chris@87: Chris@87: Chris@87: def intersect1d(ar1, ar2, assume_unique=False): Chris@87: """ Chris@87: Returns the unique elements common to both arrays. Chris@87: Chris@87: Masked values are considered equal one to the other. Chris@87: The output is always a masked array. Chris@87: Chris@87: See `numpy.intersect1d` for more details. Chris@87: Chris@87: See Also Chris@87: -------- Chris@87: numpy.intersect1d : Equivalent function for ndarrays. Chris@87: Chris@87: Examples Chris@87: -------- Chris@87: >>> x = array([1, 3, 3, 3], mask=[0, 0, 0, 1]) Chris@87: >>> y = array([3, 1, 1, 1], mask=[0, 0, 0, 1]) Chris@87: >>> intersect1d(x, y) Chris@87: masked_array(data = [1 3 --], Chris@87: mask = [False False True], Chris@87: fill_value = 999999) Chris@87: Chris@87: """ Chris@87: if assume_unique: Chris@87: aux = ma.concatenate((ar1, ar2)) Chris@87: else: Chris@87: # Might be faster than unique( intersect1d( ar1, ar2 ) )? Chris@87: aux = ma.concatenate((unique(ar1), unique(ar2))) Chris@87: aux.sort() Chris@87: return aux[:-1][aux[1:] == aux[:-1]] Chris@87: Chris@87: Chris@87: def setxor1d(ar1, ar2, assume_unique=False): Chris@87: """ Chris@87: Set exclusive-or of 1-D arrays with unique elements. Chris@87: Chris@87: The output is always a masked array. See `numpy.setxor1d` for more details. Chris@87: Chris@87: See Also Chris@87: -------- Chris@87: numpy.setxor1d : Equivalent function for ndarrays. Chris@87: Chris@87: """ Chris@87: if not assume_unique: Chris@87: ar1 = unique(ar1) Chris@87: ar2 = unique(ar2) Chris@87: Chris@87: aux = ma.concatenate((ar1, ar2)) Chris@87: if aux.size == 0: Chris@87: return aux Chris@87: aux.sort() Chris@87: auxf = aux.filled() Chris@87: # flag = ediff1d( aux, to_end = 1, to_begin = 1 ) == 0 Chris@87: flag = ma.concatenate(([True], (auxf[1:] != auxf[:-1]), [True])) Chris@87: # flag2 = ediff1d( flag ) == 0 Chris@87: flag2 = (flag[1:] == flag[:-1]) Chris@87: return aux[flag2] Chris@87: Chris@87: def in1d(ar1, ar2, assume_unique=False, invert=False): Chris@87: """ Chris@87: Test whether each element of an array is also present in a second Chris@87: array. Chris@87: Chris@87: The output is always a masked array. See `numpy.in1d` for more details. Chris@87: Chris@87: See Also Chris@87: -------- Chris@87: numpy.in1d : Equivalent function for ndarrays. Chris@87: Chris@87: Notes Chris@87: ----- Chris@87: .. versionadded:: 1.4.0 Chris@87: Chris@87: """ Chris@87: if not assume_unique: Chris@87: ar1, rev_idx = unique(ar1, return_inverse=True) Chris@87: ar2 = unique(ar2) Chris@87: Chris@87: ar = ma.concatenate((ar1, ar2)) Chris@87: # We need this to be a stable sort, so always use 'mergesort' Chris@87: # here. The values from the first array should always come before Chris@87: # the values from the second array. Chris@87: order = ar.argsort(kind='mergesort') Chris@87: sar = ar[order] Chris@87: if invert: Chris@87: bool_ar = (sar[1:] != sar[:-1]) Chris@87: else: Chris@87: bool_ar = (sar[1:] == sar[:-1]) Chris@87: flag = ma.concatenate((bool_ar, [invert])) Chris@87: indx = order.argsort(kind='mergesort')[:len(ar1)] Chris@87: Chris@87: if assume_unique: Chris@87: return flag[indx] Chris@87: else: Chris@87: return flag[indx][rev_idx] Chris@87: Chris@87: Chris@87: def union1d(ar1, ar2): Chris@87: """ Chris@87: Union of two arrays. Chris@87: Chris@87: The output is always a masked array. See `numpy.union1d` for more details. Chris@87: Chris@87: See also Chris@87: -------- Chris@87: numpy.union1d : Equivalent function for ndarrays. Chris@87: Chris@87: """ Chris@87: return unique(ma.concatenate((ar1, ar2))) Chris@87: Chris@87: Chris@87: def setdiff1d(ar1, ar2, assume_unique=False): Chris@87: """ Chris@87: Set difference of 1D arrays with unique elements. Chris@87: Chris@87: The output is always a masked array. See `numpy.setdiff1d` for more Chris@87: details. Chris@87: Chris@87: See Also Chris@87: -------- Chris@87: numpy.setdiff1d : Equivalent function for ndarrays. Chris@87: Chris@87: Examples Chris@87: -------- Chris@87: >>> x = np.ma.array([1, 2, 3, 4], mask=[0, 1, 0, 1]) Chris@87: >>> np.ma.extras.setdiff1d(x, [1, 2]) Chris@87: masked_array(data = [3 --], Chris@87: mask = [False True], Chris@87: fill_value = 999999) Chris@87: Chris@87: """ Chris@87: if not assume_unique: Chris@87: ar1 = unique(ar1) Chris@87: ar2 = unique(ar2) Chris@87: aux = in1d(ar1, ar2, assume_unique=True) Chris@87: if aux.size == 0: Chris@87: return aux Chris@87: else: Chris@87: return ma.asarray(ar1)[aux == 0] Chris@87: Chris@87: Chris@87: #####-------------------------------------------------------------------------- Chris@87: #---- --- Covariance --- Chris@87: #####-------------------------------------------------------------------------- Chris@87: Chris@87: Chris@87: Chris@87: Chris@87: def _covhelper(x, y=None, rowvar=True, allow_masked=True): Chris@87: """ Chris@87: Private function for the computation of covariance and correlation Chris@87: coefficients. Chris@87: Chris@87: """ Chris@87: x = ma.array(x, ndmin=2, copy=True, dtype=float) Chris@87: xmask = ma.getmaskarray(x) Chris@87: # Quick exit if we can't process masked data Chris@87: if not allow_masked and xmask.any(): Chris@87: raise ValueError("Cannot process masked data...") Chris@87: # Chris@87: if x.shape[0] == 1: Chris@87: rowvar = True Chris@87: # Make sure that rowvar is either 0 or 1 Chris@87: rowvar = int(bool(rowvar)) Chris@87: axis = 1 - rowvar Chris@87: if rowvar: Chris@87: tup = (slice(None), None) Chris@87: else: Chris@87: tup = (None, slice(None)) Chris@87: # Chris@87: if y is None: Chris@87: xnotmask = np.logical_not(xmask).astype(int) Chris@87: else: Chris@87: y = array(y, copy=False, ndmin=2, dtype=float) Chris@87: ymask = ma.getmaskarray(y) Chris@87: if not allow_masked and ymask.any(): Chris@87: raise ValueError("Cannot process masked data...") Chris@87: if xmask.any() or ymask.any(): Chris@87: if y.shape == x.shape: Chris@87: # Define some common mask Chris@87: common_mask = np.logical_or(xmask, ymask) Chris@87: if common_mask is not nomask: Chris@87: x.unshare_mask() Chris@87: y.unshare_mask() Chris@87: xmask = x._mask = y._mask = ymask = common_mask Chris@87: x = ma.concatenate((x, y), axis) Chris@87: xnotmask = np.logical_not(np.concatenate((xmask, ymask), axis)).astype(int) Chris@87: x -= x.mean(axis=rowvar)[tup] Chris@87: return (x, xnotmask, rowvar) Chris@87: Chris@87: Chris@87: def cov(x, y=None, rowvar=True, bias=False, allow_masked=True, ddof=None): Chris@87: """ Chris@87: Estimate the covariance matrix. Chris@87: Chris@87: Except for the handling of missing data this function does the same as Chris@87: `numpy.cov`. For more details and examples, see `numpy.cov`. Chris@87: Chris@87: By default, masked values are recognized as such. If `x` and `y` have the Chris@87: same shape, a common mask is allocated: if ``x[i,j]`` is masked, then Chris@87: ``y[i,j]`` will also be masked. Chris@87: Setting `allow_masked` to False will raise an exception if values are Chris@87: missing in either of the input arrays. Chris@87: Chris@87: Parameters Chris@87: ---------- Chris@87: x : array_like Chris@87: A 1-D or 2-D array containing multiple variables and observations. Chris@87: Each row of `x` represents a variable, and each column a single Chris@87: observation of all those variables. Also see `rowvar` below. Chris@87: y : array_like, optional Chris@87: An additional set of variables and observations. `y` has the same Chris@87: form as `x`. Chris@87: rowvar : bool, optional Chris@87: If `rowvar` is True (default), then each row represents a Chris@87: variable, with observations in the columns. Otherwise, the relationship Chris@87: is transposed: each column represents a variable, while the rows Chris@87: contain observations. Chris@87: bias : bool, optional Chris@87: Default normalization (False) is by ``(N-1)``, where ``N`` is the Chris@87: number of observations given (unbiased estimate). If `bias` is True, Chris@87: then normalization is by ``N``. This keyword can be overridden by Chris@87: the keyword ``ddof`` in numpy versions >= 1.5. Chris@87: allow_masked : bool, optional Chris@87: If True, masked values are propagated pair-wise: if a value is masked Chris@87: in `x`, the corresponding value is masked in `y`. Chris@87: If False, raises a `ValueError` exception when some values are missing. Chris@87: ddof : {None, int}, optional Chris@87: .. versionadded:: 1.5 Chris@87: If not ``None`` normalization is by ``(N - ddof)``, where ``N`` is Chris@87: the number of observations; this overrides the value implied by Chris@87: ``bias``. The default value is ``None``. Chris@87: Chris@87: Chris@87: Raises Chris@87: ------ Chris@87: ValueError Chris@87: Raised if some values are missing and `allow_masked` is False. Chris@87: Chris@87: See Also Chris@87: -------- Chris@87: numpy.cov Chris@87: Chris@87: """ Chris@87: # Check inputs Chris@87: if ddof is not None and ddof != int(ddof): Chris@87: raise ValueError("ddof must be an integer") Chris@87: # Set up ddof Chris@87: if ddof is None: Chris@87: if bias: Chris@87: ddof = 0 Chris@87: else: Chris@87: ddof = 1 Chris@87: Chris@87: (x, xnotmask, rowvar) = _covhelper(x, y, rowvar, allow_masked) Chris@87: if not rowvar: Chris@87: fact = np.dot(xnotmask.T, xnotmask) * 1. - ddof Chris@87: result = (dot(x.T, x.conj(), strict=False) / fact).squeeze() Chris@87: else: Chris@87: fact = np.dot(xnotmask, xnotmask.T) * 1. - ddof Chris@87: result = (dot(x, x.T.conj(), strict=False) / fact).squeeze() Chris@87: return result Chris@87: Chris@87: Chris@87: def corrcoef(x, y=None, rowvar=True, bias=False, allow_masked=True, ddof=None): Chris@87: """ Chris@87: Return correlation coefficients of the input array. Chris@87: Chris@87: Except for the handling of missing data this function does the same as Chris@87: `numpy.corrcoef`. For more details and examples, see `numpy.corrcoef`. Chris@87: Chris@87: Parameters Chris@87: ---------- Chris@87: x : array_like Chris@87: A 1-D or 2-D array containing multiple variables and observations. Chris@87: Each row of `x` represents a variable, and each column a single Chris@87: observation of all those variables. Also see `rowvar` below. Chris@87: y : array_like, optional Chris@87: An additional set of variables and observations. `y` has the same Chris@87: shape as `x`. Chris@87: rowvar : bool, optional Chris@87: If `rowvar` is True (default), then each row represents a Chris@87: variable, with observations in the columns. Otherwise, the relationship Chris@87: is transposed: each column represents a variable, while the rows Chris@87: contain observations. Chris@87: bias : bool, optional Chris@87: Default normalization (False) is by ``(N-1)``, where ``N`` is the Chris@87: number of observations given (unbiased estimate). If `bias` is 1, Chris@87: then normalization is by ``N``. This keyword can be overridden by Chris@87: the keyword ``ddof`` in numpy versions >= 1.5. Chris@87: allow_masked : bool, optional Chris@87: If True, masked values are propagated pair-wise: if a value is masked Chris@87: in `x`, the corresponding value is masked in `y`. Chris@87: If False, raises an exception. Chris@87: ddof : {None, int}, optional Chris@87: .. versionadded:: 1.5 Chris@87: If not ``None`` normalization is by ``(N - ddof)``, where ``N`` is Chris@87: the number of observations; this overrides the value implied by Chris@87: ``bias``. The default value is ``None``. Chris@87: Chris@87: See Also Chris@87: -------- Chris@87: numpy.corrcoef : Equivalent function in top-level NumPy module. Chris@87: cov : Estimate the covariance matrix. Chris@87: Chris@87: """ Chris@87: # Check inputs Chris@87: if ddof is not None and ddof != int(ddof): Chris@87: raise ValueError("ddof must be an integer") Chris@87: # Set up ddof Chris@87: if ddof is None: Chris@87: if bias: Chris@87: ddof = 0 Chris@87: else: Chris@87: ddof = 1 Chris@87: Chris@87: # Get the data Chris@87: (x, xnotmask, rowvar) = _covhelper(x, y, rowvar, allow_masked) Chris@87: # Compute the covariance matrix Chris@87: if not rowvar: Chris@87: fact = np.dot(xnotmask.T, xnotmask) * 1. - ddof Chris@87: c = (dot(x.T, x.conj(), strict=False) / fact).squeeze() Chris@87: else: Chris@87: fact = np.dot(xnotmask, xnotmask.T) * 1. - ddof Chris@87: c = (dot(x, x.T.conj(), strict=False) / fact).squeeze() Chris@87: # Check whether we have a scalar Chris@87: try: Chris@87: diag = ma.diagonal(c) Chris@87: except ValueError: Chris@87: return 1 Chris@87: # Chris@87: if xnotmask.all(): Chris@87: _denom = ma.sqrt(ma.multiply.outer(diag, diag)) Chris@87: else: Chris@87: _denom = diagflat(diag) Chris@87: n = x.shape[1 - rowvar] Chris@87: if rowvar: Chris@87: for i in range(n - 1): Chris@87: for j in range(i + 1, n): Chris@87: _x = mask_cols(vstack((x[i], x[j]))).var(axis=1, ddof=ddof) Chris@87: _denom[i, j] = _denom[j, i] = ma.sqrt(ma.multiply.reduce(_x)) Chris@87: else: Chris@87: for i in range(n - 1): Chris@87: for j in range(i + 1, n): Chris@87: _x = mask_cols( Chris@87: vstack((x[:, i], x[:, j]))).var(axis=1, ddof=ddof) Chris@87: _denom[i, j] = _denom[j, i] = ma.sqrt(ma.multiply.reduce(_x)) Chris@87: return c / _denom Chris@87: Chris@87: #####-------------------------------------------------------------------------- Chris@87: #---- --- Concatenation helpers --- Chris@87: #####-------------------------------------------------------------------------- Chris@87: Chris@87: class MAxisConcatenator(AxisConcatenator): Chris@87: """ Chris@87: Translate slice objects to concatenation along an axis. Chris@87: Chris@87: For documentation on usage, see `mr_class`. Chris@87: Chris@87: See Also Chris@87: -------- Chris@87: mr_class Chris@87: Chris@87: """ Chris@87: Chris@87: def __init__(self, axis=0): Chris@87: AxisConcatenator.__init__(self, axis, matrix=False) Chris@87: Chris@87: def __getitem__(self, key): Chris@87: if isinstance(key, str): Chris@87: raise MAError("Unavailable for masked array.") Chris@87: if not isinstance(key, tuple): Chris@87: key = (key,) Chris@87: objs = [] Chris@87: scalars = [] Chris@87: final_dtypedescr = None Chris@87: for k in range(len(key)): Chris@87: scalar = False Chris@87: if isinstance(key[k], slice): Chris@87: step = key[k].step Chris@87: start = key[k].start Chris@87: stop = key[k].stop Chris@87: if start is None: Chris@87: start = 0 Chris@87: if step is None: Chris@87: step = 1 Chris@87: if isinstance(step, complex): Chris@87: size = int(abs(step)) Chris@87: newobj = np.linspace(start, stop, num=size) Chris@87: else: Chris@87: newobj = np.arange(start, stop, step) Chris@87: elif isinstance(key[k], str): Chris@87: if (key[k] in 'rc'): Chris@87: self.matrix = True Chris@87: self.col = (key[k] == 'c') Chris@87: continue Chris@87: try: Chris@87: self.axis = int(key[k]) Chris@87: continue Chris@87: except (ValueError, TypeError): Chris@87: raise ValueError("Unknown special directive") Chris@87: elif type(key[k]) in np.ScalarType: Chris@87: newobj = asarray([key[k]]) Chris@87: scalars.append(k) Chris@87: scalar = True Chris@87: else: Chris@87: newobj = key[k] Chris@87: objs.append(newobj) Chris@87: if isinstance(newobj, ndarray) and not scalar: Chris@87: if final_dtypedescr is None: Chris@87: final_dtypedescr = newobj.dtype Chris@87: elif newobj.dtype > final_dtypedescr: Chris@87: final_dtypedescr = newobj.dtype Chris@87: if final_dtypedescr is not None: Chris@87: for k in scalars: Chris@87: objs[k] = objs[k].astype(final_dtypedescr) Chris@87: res = concatenate(tuple(objs), axis=self.axis) Chris@87: return self._retval(res) Chris@87: Chris@87: class mr_class(MAxisConcatenator): Chris@87: """ Chris@87: Translate slice objects to concatenation along the first axis. Chris@87: Chris@87: This is the masked array version of `lib.index_tricks.RClass`. Chris@87: Chris@87: See Also Chris@87: -------- Chris@87: lib.index_tricks.RClass Chris@87: Chris@87: Examples Chris@87: -------- Chris@87: >>> np.ma.mr_[np.ma.array([1,2,3]), 0, 0, np.ma.array([4,5,6])] Chris@87: array([1, 2, 3, 0, 0, 4, 5, 6]) Chris@87: Chris@87: """ Chris@87: def __init__(self): Chris@87: MAxisConcatenator.__init__(self, 0) Chris@87: Chris@87: mr_ = mr_class() Chris@87: Chris@87: #####-------------------------------------------------------------------------- Chris@87: #---- Find unmasked data --- Chris@87: #####-------------------------------------------------------------------------- Chris@87: Chris@87: def flatnotmasked_edges(a): Chris@87: """ Chris@87: Find the indices of the first and last unmasked values. Chris@87: Chris@87: Expects a 1-D `MaskedArray`, returns None if all values are masked. Chris@87: Chris@87: Parameters Chris@87: ---------- Chris@87: arr : array_like Chris@87: Input 1-D `MaskedArray` Chris@87: Chris@87: Returns Chris@87: ------- Chris@87: edges : ndarray or None Chris@87: The indices of first and last non-masked value in the array. Chris@87: Returns None if all values are masked. Chris@87: Chris@87: See Also Chris@87: -------- Chris@87: flatnotmasked_contiguous, notmasked_contiguous, notmasked_edges, Chris@87: clump_masked, clump_unmasked Chris@87: Chris@87: Notes Chris@87: ----- Chris@87: Only accepts 1-D arrays. Chris@87: Chris@87: Examples Chris@87: -------- Chris@87: >>> a = np.ma.arange(10) Chris@87: >>> flatnotmasked_edges(a) Chris@87: [0,-1] Chris@87: Chris@87: >>> mask = (a < 3) | (a > 8) | (a == 5) Chris@87: >>> a[mask] = np.ma.masked Chris@87: >>> np.array(a[~a.mask]) Chris@87: array([3, 4, 6, 7, 8]) Chris@87: Chris@87: >>> flatnotmasked_edges(a) Chris@87: array([3, 8]) Chris@87: Chris@87: >>> a[:] = np.ma.masked Chris@87: >>> print flatnotmasked_edges(ma) Chris@87: None Chris@87: Chris@87: """ Chris@87: m = getmask(a) Chris@87: if m is nomask or not np.any(m): Chris@87: return np.array([0, a.size - 1]) Chris@87: unmasked = np.flatnonzero(~m) Chris@87: if len(unmasked) > 0: Chris@87: return unmasked[[0, -1]] Chris@87: else: Chris@87: return None Chris@87: Chris@87: Chris@87: def notmasked_edges(a, axis=None): Chris@87: """ Chris@87: Find the indices of the first and last unmasked values along an axis. Chris@87: Chris@87: If all values are masked, return None. Otherwise, return a list Chris@87: of two tuples, corresponding to the indices of the first and last Chris@87: unmasked values respectively. Chris@87: Chris@87: Parameters Chris@87: ---------- Chris@87: a : array_like Chris@87: The input array. Chris@87: axis : int, optional Chris@87: Axis along which to perform the operation. Chris@87: If None (default), applies to a flattened version of the array. Chris@87: Chris@87: Returns Chris@87: ------- Chris@87: edges : ndarray or list Chris@87: An array of start and end indexes if there are any masked data in Chris@87: the array. If there are no masked data in the array, `edges` is a Chris@87: list of the first and last index. Chris@87: Chris@87: See Also Chris@87: -------- Chris@87: flatnotmasked_contiguous, flatnotmasked_edges, notmasked_contiguous, Chris@87: clump_masked, clump_unmasked Chris@87: Chris@87: Examples Chris@87: -------- Chris@87: >>> a = np.arange(9).reshape((3, 3)) Chris@87: >>> m = np.zeros_like(a) Chris@87: >>> m[1:, 1:] = 1 Chris@87: Chris@87: >>> am = np.ma.array(a, mask=m) Chris@87: >>> np.array(am[~am.mask]) Chris@87: array([0, 1, 2, 3, 6]) Chris@87: Chris@87: >>> np.ma.extras.notmasked_edges(ma) Chris@87: array([0, 6]) Chris@87: Chris@87: """ Chris@87: a = asarray(a) Chris@87: if axis is None or a.ndim == 1: Chris@87: return flatnotmasked_edges(a) Chris@87: m = getmaskarray(a) Chris@87: idx = array(np.indices(a.shape), mask=np.asarray([m] * a.ndim)) Chris@87: return [tuple([idx[i].min(axis).compressed() for i in range(a.ndim)]), Chris@87: tuple([idx[i].max(axis).compressed() for i in range(a.ndim)]), ] Chris@87: Chris@87: Chris@87: def flatnotmasked_contiguous(a): Chris@87: """ Chris@87: Find contiguous unmasked data in a masked array along the given axis. Chris@87: Chris@87: Parameters Chris@87: ---------- Chris@87: a : narray Chris@87: The input array. Chris@87: Chris@87: Returns Chris@87: ------- Chris@87: slice_list : list Chris@87: A sorted sequence of slices (start index, end index). Chris@87: Chris@87: See Also Chris@87: -------- Chris@87: flatnotmasked_edges, notmasked_contiguous, notmasked_edges, Chris@87: clump_masked, clump_unmasked Chris@87: Chris@87: Notes Chris@87: ----- Chris@87: Only accepts 2-D arrays at most. Chris@87: Chris@87: Examples Chris@87: -------- Chris@87: >>> a = np.ma.arange(10) Chris@87: >>> np.ma.extras.flatnotmasked_contiguous(a) Chris@87: slice(0, 10, None) Chris@87: Chris@87: >>> mask = (a < 3) | (a > 8) | (a == 5) Chris@87: >>> a[mask] = np.ma.masked Chris@87: >>> np.array(a[~a.mask]) Chris@87: array([3, 4, 6, 7, 8]) Chris@87: Chris@87: >>> np.ma.extras.flatnotmasked_contiguous(a) Chris@87: [slice(3, 5, None), slice(6, 9, None)] Chris@87: >>> a[:] = np.ma.masked Chris@87: >>> print np.ma.extras.flatnotmasked_edges(a) Chris@87: None Chris@87: Chris@87: """ Chris@87: m = getmask(a) Chris@87: if m is nomask: Chris@87: return slice(0, a.size, None) Chris@87: i = 0 Chris@87: result = [] Chris@87: for (k, g) in itertools.groupby(m.ravel()): Chris@87: n = len(list(g)) Chris@87: if not k: Chris@87: result.append(slice(i, i + n)) Chris@87: i += n Chris@87: return result or None Chris@87: Chris@87: def notmasked_contiguous(a, axis=None): Chris@87: """ Chris@87: Find contiguous unmasked data in a masked array along the given axis. Chris@87: Chris@87: Parameters Chris@87: ---------- Chris@87: a : array_like Chris@87: The input array. Chris@87: axis : int, optional Chris@87: Axis along which to perform the operation. Chris@87: If None (default), applies to a flattened version of the array. Chris@87: Chris@87: Returns Chris@87: ------- Chris@87: endpoints : list Chris@87: A list of slices (start and end indexes) of unmasked indexes Chris@87: in the array. Chris@87: Chris@87: See Also Chris@87: -------- Chris@87: flatnotmasked_edges, flatnotmasked_contiguous, notmasked_edges, Chris@87: clump_masked, clump_unmasked Chris@87: Chris@87: Notes Chris@87: ----- Chris@87: Only accepts 2-D arrays at most. Chris@87: Chris@87: Examples Chris@87: -------- Chris@87: >>> a = np.arange(9).reshape((3, 3)) Chris@87: >>> mask = np.zeros_like(a) Chris@87: >>> mask[1:, 1:] = 1 Chris@87: Chris@87: >>> ma = np.ma.array(a, mask=mask) Chris@87: >>> np.array(ma[~ma.mask]) Chris@87: array([0, 1, 2, 3, 6]) Chris@87: Chris@87: >>> np.ma.extras.notmasked_contiguous(ma) Chris@87: [slice(0, 4, None), slice(6, 7, None)] Chris@87: Chris@87: """ Chris@87: a = asarray(a) Chris@87: nd = a.ndim Chris@87: if nd > 2: Chris@87: raise NotImplementedError("Currently limited to atmost 2D array.") Chris@87: if axis is None or nd == 1: Chris@87: return flatnotmasked_contiguous(a) Chris@87: # Chris@87: result = [] Chris@87: # Chris@87: other = (axis + 1) % 2 Chris@87: idx = [0, 0] Chris@87: idx[axis] = slice(None, None) Chris@87: # Chris@87: for i in range(a.shape[other]): Chris@87: idx[other] = i Chris@87: result.append(flatnotmasked_contiguous(a[idx]) or None) Chris@87: return result Chris@87: Chris@87: Chris@87: def _ezclump(mask): Chris@87: """ Chris@87: Finds the clumps (groups of data with the same values) for a 1D bool array. Chris@87: Chris@87: Returns a series of slices. Chris@87: """ Chris@87: #def clump_masked(a): Chris@87: if mask.ndim > 1: Chris@87: mask = mask.ravel() Chris@87: idx = (mask[1:] ^ mask[:-1]).nonzero() Chris@87: idx = idx[0] + 1 Chris@87: slices = [slice(left, right) Chris@87: for (left, right) in zip(itertools.chain([0], idx), Chris@87: itertools.chain(idx, [len(mask)]),)] Chris@87: return slices Chris@87: Chris@87: Chris@87: def clump_unmasked(a): Chris@87: """ Chris@87: Return list of slices corresponding to the unmasked clumps of a 1-D array. Chris@87: (A "clump" is defined as a contiguous region of the array). Chris@87: Chris@87: Parameters Chris@87: ---------- Chris@87: a : ndarray Chris@87: A one-dimensional masked array. Chris@87: Chris@87: Returns Chris@87: ------- Chris@87: slices : list of slice Chris@87: The list of slices, one for each continuous region of unmasked Chris@87: elements in `a`. Chris@87: Chris@87: Notes Chris@87: ----- Chris@87: .. versionadded:: 1.4.0 Chris@87: Chris@87: See Also Chris@87: -------- Chris@87: flatnotmasked_edges, flatnotmasked_contiguous, notmasked_edges, Chris@87: notmasked_contiguous, clump_masked Chris@87: Chris@87: Examples Chris@87: -------- Chris@87: >>> a = np.ma.masked_array(np.arange(10)) Chris@87: >>> a[[0, 1, 2, 6, 8, 9]] = np.ma.masked Chris@87: >>> np.ma.extras.clump_unmasked(a) Chris@87: [slice(3, 6, None), slice(7, 8, None)] Chris@87: Chris@87: """ Chris@87: mask = getattr(a, '_mask', nomask) Chris@87: if mask is nomask: Chris@87: return [slice(0, a.size)] Chris@87: slices = _ezclump(mask) Chris@87: if a[0] is masked: Chris@87: result = slices[1::2] Chris@87: else: Chris@87: result = slices[::2] Chris@87: return result Chris@87: Chris@87: Chris@87: def clump_masked(a): Chris@87: """ Chris@87: Returns a list of slices corresponding to the masked clumps of a 1-D array. Chris@87: (A "clump" is defined as a contiguous region of the array). Chris@87: Chris@87: Parameters Chris@87: ---------- Chris@87: a : ndarray Chris@87: A one-dimensional masked array. Chris@87: Chris@87: Returns Chris@87: ------- Chris@87: slices : list of slice Chris@87: The list of slices, one for each continuous region of masked elements Chris@87: in `a`. Chris@87: Chris@87: Notes Chris@87: ----- Chris@87: .. versionadded:: 1.4.0 Chris@87: Chris@87: See Also Chris@87: -------- Chris@87: flatnotmasked_edges, flatnotmasked_contiguous, notmasked_edges, Chris@87: notmasked_contiguous, clump_unmasked Chris@87: Chris@87: Examples Chris@87: -------- Chris@87: >>> a = np.ma.masked_array(np.arange(10)) Chris@87: >>> a[[0, 1, 2, 6, 8, 9]] = np.ma.masked Chris@87: >>> np.ma.extras.clump_masked(a) Chris@87: [slice(0, 3, None), slice(6, 7, None), slice(8, 10, None)] Chris@87: Chris@87: """ Chris@87: mask = ma.getmask(a) Chris@87: if mask is nomask: Chris@87: return [] Chris@87: slices = _ezclump(mask) Chris@87: if len(slices): Chris@87: if a[0] is masked: Chris@87: slices = slices[::2] Chris@87: else: Chris@87: slices = slices[1::2] Chris@87: return slices Chris@87: Chris@87: Chris@87: Chris@87: #####-------------------------------------------------------------------------- Chris@87: #---- Polynomial fit --- Chris@87: #####-------------------------------------------------------------------------- Chris@87: Chris@87: def vander(x, n=None): Chris@87: """ Chris@87: Masked values in the input array result in rows of zeros. Chris@87: """ Chris@87: _vander = np.vander(x, n) Chris@87: m = getmask(x) Chris@87: if m is not nomask: Chris@87: _vander[m] = 0 Chris@87: return _vander Chris@87: vander.__doc__ = ma.doc_note(np.vander.__doc__, vander.__doc__) Chris@87: Chris@87: Chris@87: def polyfit(x, y, deg, rcond=None, full=False, w=None, cov=False): Chris@87: """ Chris@87: Any masked values in x is propagated in y, and vice-versa. Chris@87: """ Chris@87: x = asarray(x) Chris@87: y = asarray(y) Chris@87: Chris@87: m = getmask(x) Chris@87: if y.ndim == 1: Chris@87: m = mask_or(m, getmask(y)) Chris@87: elif y.ndim == 2: Chris@87: my = getmask(mask_rows(y)) Chris@87: if my is not nomask: Chris@87: m = mask_or(m, my[:, 0]) Chris@87: else: Chris@87: raise TypeError("Expected a 1D or 2D array for y!") Chris@87: Chris@87: if w is not None: Chris@87: w = asarray(w) Chris@87: if w.ndim != 1: Chris@87: raise TypeError("expected a 1-d array for weights") Chris@87: if w.shape[0] != y.shape[0] : Chris@87: raise TypeError("expected w and y to have the same length") Chris@87: m = mask_or(m, getmask(w)) Chris@87: Chris@87: if m is not nomask: Chris@87: if w is not None: Chris@87: w = ~m*w Chris@87: else: Chris@87: w = ~m Chris@87: Chris@87: return np.polyfit(x, y, deg, rcond, full, w, cov) Chris@87: Chris@87: polyfit.__doc__ = ma.doc_note(np.polyfit.__doc__, polyfit.__doc__) Chris@87: Chris@87: ################################################################################