"""

Masked arrays add-ons.

A collection of utilities for `numpy.ma`.

:author: Pierre Gerard-Marchant

:contact: pierregm_at_uga_dot_edu

"""

__all__ = [

'apply_along_axis', 'apply_over_axes', 'atleast_1d', 'atleast_2d',

'atleast_3d', 'average', 'clump_masked', 'clump_unmasked', 'column_stack',

'compress_cols', 'compress_nd', 'compress_rowcols', 'compress_rows',

'count_masked', 'corrcoef', 'cov', 'diagflat', 'dot', 'dstack', 'ediff1d',

'flatnotmasked_contiguous', 'flatnotmasked_edges', 'hsplit', 'hstack',

'isin', 'in1d', 'intersect1d', 'mask_cols', 'mask_rowcols', 'mask_rows',

'masked_all', 'masked_all_like', 'median', 'mr_', 'ndenumerate',

'notmasked_contiguous', 'notmasked_edges', 'polyfit', 'row_stack',

'setdiff1d', 'setxor1d', 'stack', 'unique', 'union1d', 'vander', 'vstack',

]

import functools

import itertools

import warnings

import numpy as np

from numpy import array as nxarray, ndarray

from numpy.lib._function_base_impl import _ureduce

from numpy.lib._index_tricks_impl import AxisConcatenator

from numpy.lib.array_utils import normalize_axis_index, normalize_axis_tuple

from . import core as ma

from .core import ( # noqa: F401

MAError,

MaskedArray,

add,

array,

asarray,

concatenate,

count,

dot,

filled,

get_masked_subclass,

getdata,

getmask,

getmaskarray,

make_mask_descr,

mask_or,

masked,

masked_array,

nomask,

ones,

sort,

zeros,

)

def issequence(seq):

"""

Is seq a sequence (ndarray, list or tuple)?

"""

return isinstance(seq, (ndarray, tuple, list))

def count_masked(arr, axis=None):

"""

Count the number of masked elements along the given axis.

Parameters

----------

arr : array_like

An array with (possibly) masked elements.

axis : int, optional

Axis along which to count. If None (default), a flattened

version of the array is used.

Returns

-------

count : int, ndarray

The total number of masked elements (axis=None) or the number

of masked elements along each slice of the given axis.

See Also

--------

MaskedArray.count : Count non-masked elements.

Examples

--------

>>> import numpy as np

>>> a = np.arange(9).reshape((3,3))

>>> a = np.ma.array(a)

>>> a[1, 0] = np.ma.masked

>>> a[1, 2] = np.ma.masked

>>> a[2, 1] = np.ma.masked

>>> a

masked_array(

data=[[0, 1, 2],

[--, 4, --],

[6, --, 8]],

mask=[[False, False, False],

[ True, False, True],

[False, True, False]],

fill_value=999999)

>>> np.ma.count_masked(a)

3

When the `axis` keyword is used an array is returned.

>>> np.ma.count_masked(a, axis=0)

array([1, 1, 1])

>>> np.ma.count_masked(a, axis=1)

array([0, 2, 1])

"""

m = getmaskarray(arr)

return m.sum(axis)

def masked_all(shape, dtype=float):

"""

Empty masked array with all elements masked.

Return an empty masked array of the given shape and dtype, where all the

data are masked.

Parameters

----------

shape : int or tuple of ints

Shape of the required MaskedArray, e.g., ``(2, 3)`` or ``2``.

dtype : dtype, optional

Data type of the output.

Returns

-------

a : MaskedArray

A masked array with all data masked.

See Also

--------

masked_all_like : Empty masked array modelled on an existing array.

Notes

-----

Unlike other masked array creation functions (e.g. `numpy.ma.zeros`,

`numpy.ma.ones`, `numpy.ma.full`), `masked_all` does not initialize the

values of the array, and may therefore be marginally faster. However,

the values stored in the newly allocated array are arbitrary. For

reproducible behavior, be sure to set each element of the array before

reading.

Examples

--------

>>> import numpy as np

>>> np.ma.masked_all((3, 3))

masked_array(

data=[[--, --, --],

[--, --, --],

[--, --, --]],

mask=[[ True, True, True],

[ True, True, True],

[ True, True, True]],

fill_value=1e+20,

dtype=float64)

The `dtype` parameter defines the underlying data type.

>>> a = np.ma.masked_all((3, 3))

>>> a.dtype

dtype('float64')

>>> a = np.ma.masked_all((3, 3), dtype=np.int32)

>>> a.dtype

dtype('int32')

"""

a = masked_array(np.empty(shape, dtype),

mask=np.ones(shape, make_mask_descr(dtype)))

return a

def masked_all_like(arr):

"""

Empty masked array with the properties of an existing array.

Return an empty masked array of the same shape and dtype as

the array `arr`, where all the data are masked.

Parameters

----------

arr : ndarray

An array describing the shape and dtype of the required MaskedArray.

Returns

-------

a : MaskedArray

A masked array with all data masked.

Raises

------

AttributeError

If `arr` doesn't have a shape attribute (i.e. not an ndarray)

See Also

--------

masked_all : Empty masked array with all elements masked.

Notes

-----

Unlike other masked array creation functions (e.g. `numpy.ma.zeros_like`,

`numpy.ma.ones_like`, `numpy.ma.full_like`), `masked_all_like` does not

initialize the values of the array, and may therefore be marginally

faster. However, the values stored in the newly allocated array are

arbitrary. For reproducible behavior, be sure to set each element of the

array before reading.

Examples

--------

>>> import numpy as np

>>> arr = np.zeros((2, 3), dtype=np.float32)

>>> arr

array([[0., 0., 0.],

[0., 0., 0.]], dtype=float32)

>>> np.ma.masked_all_like(arr)

masked_array(

data=[[--, --, --],

[--, --, --]],

mask=[[ True, True, True],

[ True, True, True]],

fill_value=np.float64(1e+20),

dtype=float32)

The dtype of the masked array matches the dtype of `arr`.

>>> arr.dtype

dtype('float32')

>>> np.ma.masked_all_like(arr).dtype

dtype('float32')

"""

a = np.empty_like(arr).view(MaskedArray)

a._mask = np.ones(a.shape, dtype=make_mask_descr(a.dtype))

return a

#####--------------------------------------------------------------------------

#---- --- Standard functions ---

#####--------------------------------------------------------------------------

def _fromnxfunction_function(_fromnxfunction):

"""

Decorator to wrap a "_fromnxfunction" function, wrapping a numpy function as a

masked array function, with proper docstring and name.

Parameters

----------

_fromnxfunction : ({params}) -> ndarray, {params}) -> masked_array

Wrapper function that calls the wrapped numpy function

Returns

-------

decorator : (f: ({params}) -> ndarray) -> ({params}) -> masked_array

Function that accepts a numpy function and returns a masked array function

"""

def decorator(npfunc, /):

def wrapper(*args, **kwargs):

return _fromnxfunction(npfunc, *args, **kwargs)

functools.update_wrapper(wrapper, npfunc, assigned=("__name__", "__qualname__"))

wrapper.__doc__ = ma.doc_note(

npfunc.__doc__,

"The function is applied to both the ``_data`` and the ``_mask``, if any.",

)

return wrapper

return decorator

@_fromnxfunction_function

def _fromnxfunction_single(npfunc, a, /, *args, **kwargs):

"""

Wraps a NumPy function that can be called with a single array argument followed by

auxiliary args that are passed verbatim for both the data and mask calls.

"""

return masked_array(

data=npfunc(np.asarray(a), *args, **kwargs),

mask=npfunc(getmaskarray(a), *args, **kwargs),

)

@_fromnxfunction_function

def _fromnxfunction_seq(npfunc, arys, /, *args, **kwargs):

"""

Wraps a NumPy function that can be called with a single sequence of arrays followed

by auxiliary args that are passed verbatim for both the data and mask calls.

"""

return masked_array(

data=npfunc(tuple(np.asarray(a) for a in arys), *args, **kwargs),

mask=npfunc(tuple(getmaskarray(a) for a in arys), *args, **kwargs),

)

@_fromnxfunction_function

def _fromnxfunction_allargs(npfunc, /, *arys, **kwargs):

"""

Wraps a NumPy function that can be called with multiple array arguments.

All args are converted to arrays even if they are not so already.

This makes it possible to process scalars as 1-D arrays.

Only keyword arguments are passed through verbatim for the data and mask calls.

Arrays arguments are processed independently and the results are returned in a list.

If only one arg is present, the return value is just the processed array instead of

a list.

"""

out = tuple(

masked_array(

data=npfunc(np.asarray(a), **kwargs),

mask=npfunc(getmaskarray(a), **kwargs),

)

for a in arys

)

return out[0] if len(out) == 1 else out

atleast_1d = _fromnxfunction_allargs(np.atleast_1d)

atleast_2d = _fromnxfunction_allargs(np.atleast_2d)

atleast_3d = _fromnxfunction_allargs(np.atleast_3d)

vstack = row_stack = _fromnxfunction_seq(np.vstack)

hstack = _fromnxfunction_seq(np.hstack)

column_stack = _fromnxfunction_seq(np.column_stack)

dstack = _fromnxfunction_seq(np.dstack)

stack = _fromnxfunction_seq(np.stack)

hsplit = _fromnxfunction_single(np.hsplit)

diagflat = _fromnxfunction_single(np.diagflat)

#####--------------------------------------------------------------------------

#----

#####--------------------------------------------------------------------------

def flatten_inplace(seq):

"""Flatten a sequence in place."""

k = 0

while (k != len(seq)):

while hasattr(seq[k], '__iter__'):

seq[k:(k + 1)] = seq[k]

k += 1

return seq

def apply_along_axis(func1d, axis, arr, *args, **kwargs):

"""

(This docstring should be overwritten)

"""

arr = array(arr, copy=False, subok=True)

nd = arr.ndim

axis = normalize_axis_index(axis, nd)

ind = [0] * (nd - 1)

i = np.zeros(nd, 'O')

indlist = list(range(nd))

indlist.remove(axis)

i[axis] = slice(None, None)

outshape = np.asarray(arr.shape).take(indlist)

i.put(indlist, ind)

res = func1d(arr[tuple(i.tolist())], *args, **kwargs)

# if res is a number, then we have a smaller output array

asscalar = np.isscalar(res)

if not asscalar:

try:

len(res)

except TypeError:

asscalar = True

# Note: we shouldn't set the dtype of the output from the first result

# so we force the type to object, and build a list of dtypes. We'll

# just take the largest, to avoid some downcasting

dtypes = []

if asscalar:

dtypes.append(np.asarray(res).dtype)

outarr = zeros(outshape, object)

outarr[tuple(ind)] = res

Ntot = np.prod(outshape)

k = 1

while k < Ntot:

# increment the index

ind[-1] += 1

n = -1

while (ind[n] >= outshape[n]) and (n > (1 - nd)):

ind[n - 1] += 1

ind[n] = 0

n -= 1

i.put(indlist, ind)

res = func1d(arr[tuple(i.tolist())], *args, **kwargs)

outarr[tuple(ind)] = res

dtypes.append(asarray(res).dtype)

k += 1

else:

res = array(res, copy=False, subok=True)

j = i.copy()

j[axis] = ([slice(None, None)] * res.ndim)

j.put(indlist, ind)

Ntot = np.prod(outshape)

holdshape = outshape

outshape = list(arr.shape)

outshape[axis] = res.shape

dtypes.append(asarray(res).dtype)

outshape = flatten_inplace(outshape)

outarr = zeros(outshape, object)

outarr[tuple(flatten_inplace(j.tolist()))] = res

k = 1

while k < Ntot:

# increment the index

ind[-1] += 1

n = -1

while (ind[n] >= holdshape[n]) and (n > (1 - nd)):

ind[n - 1] += 1

ind[n] = 0

n -= 1

i.put(indlist, ind)

j.put(indlist, ind)

res = func1d(arr[tuple(i.tolist())], *args, **kwargs)

outarr[tuple(flatten_inplace(j.tolist()))] = res

dtypes.append(asarray(res).dtype)

k += 1

max_dtypes = np.dtype(np.asarray(dtypes).max())

if not hasattr(arr, '_mask'):

result = np.asarray(outarr, dtype=max_dtypes)

else:

result = asarray(outarr, dtype=max_dtypes)

result.fill_value = ma.default_fill_value(result)

return result

apply_along_axis.__doc__ = np.apply_along_axis.__doc__

def apply_over_axes(func, a, axes):

"""

(This docstring will be overwritten)

"""

val = asarray(a)

N = a.ndim

if array(axes).ndim == 0:

axes = (axes,)

for axis in axes:

if axis < 0:

axis = N + axis

args = (val, axis)

res = func(*args)

if res.ndim == val.ndim:

val = res

else:

res = ma.expand_dims(res, axis)

if res.ndim == val.ndim:

val = res

else:

raise ValueError("function is not returning "

"an array of the correct shape")

return val

if apply_over_axes.__doc__ is not None:

apply_over_axes.__doc__ = np.apply_over_axes.__doc__[

:np.apply_over_axes.__doc__.find('Notes')].rstrip() + \

"""

Examples

--------

>>> import numpy as np

>>> a = np.ma.arange(24).reshape(2,3,4)

>>> a[:,0,1] = np.ma.masked

>>> a[:,1,:] = np.ma.masked

>>> a

masked_array(

data=[[[0, --, 2, 3],

[--, --, --, --],

[8, 9, 10, 11]],

[[12, --, 14, 15],

[--, --, --, --],

[20, 21, 22, 23]]],

mask=[[[False, True, False, False],

[ True, True, True, True],

[False, False, False, False]],

[[False, True, False, False],

[ True, True, True, True],

[False, False, False, False]]],

fill_value=999999)

>>> np.ma.apply_over_axes(np.ma.sum, a, [0,2])

masked_array(

data=[[[46],

[--],

[124]]],

mask=[[[False],

[ True],

[False]]],

fill_value=999999)

Tuple axis arguments to ufuncs are equivalent:

>>> np.ma.sum(a, axis=(0,2)).reshape((1,-1,1))

masked_array(

data=[[[46],

[--],

[124]]],

mask=[[[False],

[ True],

[False]]],

fill_value=999999)

"""

def average(a, axis=None, weights=None, returned=False, *,

keepdims=np._NoValue):

"""

Return the weighted average of array over the given axis.

Parameters

----------

a : array_like

Data to be averaged.

Masked entries are not taken into account in the computation.

axis : None or int or tuple of ints, optional

Axis or axes along which to average `a`. The default,

`axis=None`, will average over all of the elements of the input array.

If axis is a tuple of ints, averaging is performed on all of the axes

specified in the tuple instead of a single axis or all the axes as

before.

weights : array_like, optional

An array of weights associated with the values in `a`. Each value in

`a` contributes to the average according to its associated weight.

The array of weights must be the same shape as `a` if no axis is

specified, otherwise the weights must have dimensions and shape

consistent with `a` along the specified axis.

If `weights=None`, then all data in `a` are assumed to have a

weight equal to one.

The calculation is::

avg = sum(a * weights) / sum(weights)

where the sum is over all included elements.

The only constraint on the values of `weights` is that `sum(weights)`

must not be 0.

returned : bool, optional

Flag indicating whether a tuple ``(result, sum of weights)``

should be returned as output (True), or just the result (False).

Default is False.

keepdims : bool, optional

If this is set to True, the axes which are reduced are left

in the result as dimensions with size one. With this option,

the result will broadcast correctly against the original `a`.

*Note:* `keepdims` will not work with instances of `numpy.matrix`

or other classes whose methods do not support `keepdims`.

.. versionadded:: 1.23.0

Returns

-------

average, [sum_of_weights] : (tuple of) scalar or MaskedArray

The average along the specified axis. When returned is `True`,

return a tuple with the average as the first element and the sum

of the weights as the second element. The return type is `np.float64`

if `a` is of integer type and floats smaller than `float64`, or the

input data-type, otherwise. If returned, `sum_of_weights` is always

`float64`.

Raises

------

ZeroDivisionError

When all weights along axis are zero. See `numpy.ma.average` for a

version robust to this type of error.

TypeError

When `weights` does not have the same shape as `a`, and `axis=None`.

ValueError

When `weights` does not have dimensions and shape consistent with `a`

along specified `axis`.

Examples

--------

>>> import numpy as np

>>> a = np.ma.array([1., 2., 3., 4.], mask=[False, False, True, True])

>>> np.ma.average(a, weights=[3, 1, 0, 0])

1.25

>>> x = np.ma.arange(6.).reshape(3, 2)

>>> x

masked_array(

data=[[0., 1.],

[2., 3.],

[4., 5.]],

mask=False,

fill_value=1e+20)

>>> data = np.arange(8).reshape((2, 2, 2))

>>> data

array([[[0, 1],

[2, 3]],

[[4, 5],

[6, 7]]])

>>> np.ma.average(data, axis=(0, 1), weights=[[1./4, 3./4], [1., 1./2]])

masked_array(data=[3.4, 4.4],

mask=[False, False],

fill_value=1e+20)

>>> np.ma.average(data, axis=0, weights=[[1./4, 3./4], [1., 1./2]])

Traceback (most recent call last):

...

ValueError: Shape of weights must be consistent

with shape of a along specified axis.

>>> avg, sumweights = np.ma.average(x, axis=0, weights=[1, 2, 3],

... returned=True)

>>> avg

masked_array(data=[2.6666666666666665, 3.6666666666666665],

mask=[False, False],

fill_value=1e+20)

With ``keepdims=True``, the following result has shape (3, 1).

>>> np.ma.average(x, axis=1, keepdims=True)

masked_array(

data=[[0.5],

[2.5],

[4.5]],

mask=False,

fill_value=1e+20)

"""

a = asarray(a)

m = getmask(a)

if axis is not None:

axis = normalize_axis_tuple(axis, a.ndim, argname="axis")

if keepdims is np._NoValue:

# Don't pass on the keepdims argument if one wasn't given.

keepdims_kw = {}

else:

keepdims_kw = {'keepdims': keepdims}

if weights is None:

avg = a.mean(axis, **keepdims_kw)

scl = avg.dtype.type(a.count(axis))

else:

wgt = asarray(weights)

if issubclass(a.dtype.type, (np.integer, np.bool)):

result_dtype = np.result_type(a.dtype, wgt.dtype, 'f8')

else:

result_dtype = np.result_type(a.dtype, wgt.dtype)

# Sanity checks

if a.shape != wgt.shape:

if axis is None:

raise TypeError(

"Axis must be specified when shapes of a and weights "

"differ.")

if wgt.shape != tuple(a.shape[ax] for ax in axis):

raise ValueError(

"Shape of weights must be consistent with "

"shape of a along specified axis.")

# setup wgt to broadcast along axis

wgt = wgt.transpose(np.argsort(axis))

wgt = wgt.reshape(tuple((s if ax in axis else 1)

for ax, s in enumerate(a.shape)))

if m is not nomask:

wgt = wgt * (~a.mask)

wgt.mask |= a.mask

scl = wgt.sum(axis=axis, dtype=result_dtype, **keepdims_kw)

avg = np.multiply(a, wgt,

dtype=result_dtype).sum(axis, **keepdims_kw) / scl

if returned:

if scl.shape != avg.shape:

scl = np.broadcast_to(scl, avg.shape).copy()

return avg, scl

else:

return avg

def median(a, axis=None, out=None, overwrite_input=False, keepdims=False):

"""

Compute the median along the specified axis.

Returns the median of the array elements.

Parameters

----------

a : array_like

Input array or object that can be converted to an array.

axis : int, optional

Axis along which the medians are computed. The default (None) is

to compute the median along a flattened version of the array.

out : ndarray, optional

Alternative output array in which to place the result. It must

have the same shape and buffer length as the expected output

but the type will be cast if necessary.

overwrite_input : bool, optional

If True, then allow use of memory of input array (a) for

calculations. The input array will be modified by the call to

median. This will save memory when you do not need to preserve

the contents of the input array. Treat the input as undefined,

but it will probably be fully or partially sorted. Default is

False. Note that, if `overwrite_input` is True, and the input

is not already an `ndarray`, an error will be raised.

keepdims : bool, optional

If this is set to True, the axes which are reduced are left

in the result as dimensions with size one. With this option,

the result will broadcast correctly against the input array.

Returns

-------

median : ndarray

A new array holding the result is returned unless out is

specified, in which case a reference to out is returned.

Return data-type is `float64` for integers and floats smaller than

`float64`, or the input data-type, otherwise.

See Also

--------

mean

Notes

-----

Given a vector ``V`` with ``N`` non masked values, the median of ``V``

is the middle value of a sorted copy of ``V`` (``Vs``) - i.e.

``Vs[(N-1)/2]``, when ``N`` is odd, or ``{Vs[N/2 - 1] + Vs[N/2]}/2``

when ``N`` is even.

Examples

--------

>>> import numpy as np

>>> x = np.ma.array(np.arange(8), mask=[0]*4 + [1]*4)

>>> np.ma.median(x)

1.5

>>> x = np.ma.array(np.arange(10).reshape(2, 5), mask=[0]*6 + [1]*4)

>>> np.ma.median(x)

2.5

>>> np.ma.median(x, axis=-1, overwrite_input=True)

masked_array(data=[2.0, 5.0],

mask=[False, False],

fill_value=1e+20)

"""

if not hasattr(a, 'mask'):

m = np.median(getdata(a, subok=True), axis=axis,

out=out, overwrite_input=overwrite_input,

keepdims=keepdims)

if isinstance(m, np.ndarray) and 1 <= m.ndim:

return masked_array(m, copy=False)

else:

return m

return _ureduce(a, func=_median, keepdims=keepdims, axis=axis, out=out,

overwrite_input=overwrite_input)

def _median(a, axis=None, out=None, overwrite_input=False):

# when an unmasked NaN is present return it, so we need to sort the NaN

# values behind the mask

if np.issubdtype(a.dtype, np.inexact):

fill_value = np.inf

else:

fill_value = None

if overwrite_input:

if axis is None:

asorted = a.ravel()

asorted.sort(fill_value=fill_value)

else:

a.sort(axis=axis, fill_value=fill_value)

asorted = a

else:

asorted = sort(a, axis=axis, fill_value=fill_value)

if axis is None:

axis = 0

else:

axis = normalize_axis_index(axis, asorted.ndim)

if asorted.shape[axis] == 0:

# for empty axis integer indices fail so use slicing to get same result

# as median (which is mean of empty slice = nan)

indexer = [slice(None)] * asorted.ndim

indexer[axis] = slice(0, 0)

indexer = tuple(indexer)

return np.ma.mean(asorted[indexer], axis=axis, out=out)

if asorted.ndim == 1:

idx, odd = divmod(count(asorted), 2)

mid = asorted[idx + odd - 1:idx + 1]

if np.issubdtype(asorted.dtype, np.inexact) and asorted.size > 0:

# avoid inf / x = masked

s = mid.sum(out=out)

if not odd:

s = np.true_divide(s, 2., casting='safe', out=out)

s = np.lib._utils_impl._median_nancheck(asorted, s, axis)

else:

s = mid.mean(out=out)

# if result is masked either the input contained enough

# minimum_fill_value so that it would be the median or all values

# masked

if np.ma.is_masked(s) and not np.all(asorted.mask):

return np.ma.minimum_fill_value(asorted)

return s

counts = count(asorted, axis=axis, keepdims=True)

h = counts // 2

# duplicate high if odd number of elements so mean does nothing

odd = counts % 2 == 1

l = np.where(odd, h, h - 1)

lh = np.concatenate([l, h], axis=axis)

# get low and high median

low_high = np.take_along_axis(asorted, lh, axis=axis)

def replace_masked(s):

# Replace masked entries with minimum_full_value unless it all values

# are masked. This is required as the sort order of values equal or

# larger than the fill value is undefined and a valid value placed

# elsewhere, e.g. [4, --, inf].

if np.ma.is_masked(s):

rep = (~np.all(asorted.mask, axis=axis, keepdims=True)) & s.mask

s.data[rep] = np.ma.minimum_fill_value(asorted)

s.mask[rep] = False

replace_masked(low_high)

if np.issubdtype(asorted.dtype, np.inexact):

# avoid inf / x = masked

s = np.ma.sum(low_high, axis=axis, out=out)

np.true_divide(s.data, 2., casting='unsafe', out=s.data)

s = np.lib._utils_impl._median_nancheck(asorted, s, axis)

else:

s = np.ma.mean(low_high, axis=axis, out=out)

return s

def compress_nd(x, axis=None):

"""Suppress slices from multiple dimensions which contain masked values.

Parameters

----------

x : array_like, MaskedArray

The array to operate on. If not a MaskedArray instance (or if no array

elements are masked), `x` is interpreted as a MaskedArray with `mask`

set to `nomask`.

axis : tuple of ints or int, optional

Which dimensions to suppress slices from can be configured with this

parameter.

- If axis is a tuple of ints, those are the axes to suppress slices from.

- If axis is an int, then that is the only axis to suppress slices from.

- If axis is None, all axis are selected.

Returns

-------

compress_array : ndarray

The compressed array.

Examples

--------

>>> import numpy as np

>>> arr = [[1, 2], [3, 4]]

>>> mask = [[0, 1], [0, 0]]

>>> x = np.ma.array(arr, mask=mask)

>>> np.ma.compress_nd(x, axis=0)

array([[3, 4]])

>>> np.ma.compress_nd(x, axis=1)

array([[1],

[3]])

>>> np.ma.compress_nd(x)

array([[3]])

"""

x = asarray(x)

m = getmask(x)

# Set axis to tuple of ints

if axis is None:

axis = tuple(range(x.ndim))

else:

axis = normalize_axis_tuple(axis, x.ndim)

# Nothing is masked: return x

if m is nomask or not m.any():

return x._data

# All is masked: return empty

if m.all():

return nxarray([])

# Filter elements through boolean indexing

data = x._data

for ax in axis:

axes = tuple(list(range(ax)) + list(range(ax + 1, x.ndim)))

data = data[(slice(None),) * ax + (~m.any(axis=axes),)]

return data

def compress_rowcols(x, axis=None):

"""

Suppress the rows and/or columns of a 2-D array that contain

masked values.

The suppression behavior is selected with the `axis` parameter.

- If axis is None, both rows and columns are suppressed.

- If axis is 0, only rows are suppressed.

- If axis is 1 or -1, only columns are suppressed.

Parameters

----------

x : array_like, MaskedArray

The array to operate on. If not a MaskedArray instance (or if no array

elements are masked), `x` is interpreted as a MaskedArray with

`mask` set to `nomask`. Must be a 2D array.

axis : int, optional

Axis along which to perform the operation. Default is None.

Returns

-------

compressed_array : ndarray

The compressed array.

Examples

--------

>>> import numpy as np

>>> x = np.ma.array(np.arange(9).reshape(3, 3), mask=[[1, 0, 0],

... [1, 0, 0],

... [0, 0, 0]])

>>> x

masked_array(

data=[[--, 1, 2],

[--, 4, 5],

[6, 7, 8]],

mask=[[ True, False, False],

[ True, False, False],

[False, False, False]],

fill_value=999999)

>>> np.ma.compress_rowcols(x)

array([[7, 8]])

>>> np.ma.compress_rowcols(x, 0)

array([[6, 7, 8]])

>>> np.ma.compress_rowcols(x, 1)

array([[1, 2],

[4, 5],

[7, 8]])

"""

if asarray(x).ndim != 2:

raise NotImplementedError("compress_rowcols works for 2D arrays only.")

return compress_nd(x, axis=axis)

def compress_rows(a):

"""

Suppress whole rows of a 2-D array that contain masked values.

This is equivalent to ``np.ma.compress_rowcols(a, 0)``, see

`compress_rowcols` for details.

Parameters

----------

x : array_like, MaskedArray

The array to operate on. If not a MaskedArray instance (or if no array

elements are masked), `x` is interpreted as a MaskedArray with

`mask` set to `nomask`. Must be a 2D array.

Returns

-------

compressed_array : ndarray

The compressed array.

See Also

--------

compress_rowcols

Examples

--------

>>> import numpy as np

>>> a = np.ma.array(np.arange(9).reshape(3, 3), mask=[[1, 0, 0],

... [1, 0, 0],

... [0, 0, 0]])

>>> np.ma.compress_rows(a)

array([[6, 7, 8]])

"""

a = asarray(a)

if a.ndim != 2:

raise NotImplementedError("compress_rows works for 2D arrays only.")

return compress_rowcols(a, 0)

def compress_cols(a):

"""

Suppress whole columns of a 2-D array that contain masked values.

This is equivalent to ``np.ma.compress_rowcols(a, 1)``, see

`compress_rowcols` for details.