codekingpro/portable-devtools
114k
1"""
2Masked arrays add-ons.
3
4A collection of utilities for `numpy.ma`.
5
6:author: Pierre Gerard-Marchant
7:contact: pierregm_at_uga_dot_edu
8
9"""
10__all__ = [
11 'apply_along_axis', 'apply_over_axes', 'atleast_1d', 'atleast_2d',
12 'atleast_3d', 'average', 'clump_masked', 'clump_unmasked', 'column_stack',
13 'compress_cols', 'compress_nd', 'compress_rowcols', 'compress_rows',
14 'count_masked', 'corrcoef', 'cov', 'diagflat', 'dot', 'dstack', 'ediff1d',
15 'flatnotmasked_contiguous', 'flatnotmasked_edges', 'hsplit', 'hstack',
16 'isin', 'in1d', 'intersect1d', 'mask_cols', 'mask_rowcols', 'mask_rows',
17 'masked_all', 'masked_all_like', 'median', 'mr_', 'ndenumerate',
18 'notmasked_contiguous', 'notmasked_edges', 'polyfit', 'row_stack',
19 'setdiff1d', 'setxor1d', 'stack', 'unique', 'union1d', 'vander', 'vstack',
20 ]
21
22import functools
23import itertools
24import warnings
25
26import numpy as np
27from numpy import array as nxarray, ndarray
28from numpy.lib._function_base_impl import _ureduce
29from numpy.lib._index_tricks_impl import AxisConcatenator
30from numpy.lib.array_utils import normalize_axis_index, normalize_axis_tuple
31
32from . import core as ma
33from .core import ( # noqa: F401
34 MAError,
35 MaskedArray,
36 add,
37 array,
38 asarray,
39 concatenate,
40 count,
41 dot,
42 filled,
43 get_masked_subclass,
44 getdata,
45 getmask,
46 getmaskarray,
47 make_mask_descr,
48 mask_or,
49 masked,
50 masked_array,
51 nomask,
52 ones,
53 sort,
54 zeros,
55)
56
57
58def issequence(seq):
59 """
60 Is seq a sequence (ndarray, list or tuple)?
61
62 """
63 return isinstance(seq, (ndarray, tuple, list))
64
65
66def count_masked(arr, axis=None):
67 """
68 Count the number of masked elements along the given axis.
69
70 Parameters
71 ----------
72 arr : array_like
73 An array with (possibly) masked elements.
74 axis : int, optional
75 Axis along which to count. If None (default), a flattened
76 version of the array is used.
77
78 Returns
79 -------
80 count : int, ndarray
81 The total number of masked elements (axis=None) or the number
82 of masked elements along each slice of the given axis.
83
84 See Also
85 --------
86 MaskedArray.count : Count non-masked elements.
87
88 Examples
89 --------
90 >>> import numpy as np
91 >>> a = np.arange(9).reshape((3,3))
92 >>> a = np.ma.array(a)
93 >>> a[1, 0] = np.ma.masked
94 >>> a[1, 2] = np.ma.masked
95 >>> a[2, 1] = np.ma.masked
96 >>> a
97 masked_array(
98 data=[[0, 1, 2],
99 [--, 4, --],
100 [6, --, 8]],
101 mask=[[False, False, False],
102 [ True, False, True],
103 [False, True, False]],
104 fill_value=999999)
105 >>> np.ma.count_masked(a)
106 3
107
108 When the `axis` keyword is used an array is returned.
109
110 >>> np.ma.count_masked(a, axis=0)
111 array([1, 1, 1])
112 >>> np.ma.count_masked(a, axis=1)
113 array([0, 2, 1])
114
115 """
116 m = getmaskarray(arr)
117 return m.sum(axis)
118
119
120def masked_all(shape, dtype=float):
121 """
122 Empty masked array with all elements masked.
123
124 Return an empty masked array of the given shape and dtype, where all the
125 data are masked.
126
127 Parameters
128 ----------
129 shape : int or tuple of ints
130 Shape of the required MaskedArray, e.g., ``(2, 3)`` or ``2``.
131 dtype : dtype, optional
132 Data type of the output.
133
134 Returns
135 -------
136 a : MaskedArray
137 A masked array with all data masked.
138
139 See Also
140 --------
141 masked_all_like : Empty masked array modelled on an existing array.
142
143 Notes
144 -----
145 Unlike other masked array creation functions (e.g. `numpy.ma.zeros`,
146 `numpy.ma.ones`, `numpy.ma.full`), `masked_all` does not initialize the
147 values of the array, and may therefore be marginally faster. However,
148 the values stored in the newly allocated array are arbitrary. For
149 reproducible behavior, be sure to set each element of the array before
150 reading.
151
152 Examples
153 --------
154 >>> import numpy as np
155 >>> np.ma.masked_all((3, 3))
156 masked_array(
157 data=[[--, --, --],
158 [--, --, --],
159 [--, --, --]],
160 mask=[[ True, True, True],
161 [ True, True, True],
162 [ True, True, True]],
163 fill_value=1e+20,
164 dtype=float64)
165
166 The `dtype` parameter defines the underlying data type.
167
168 >>> a = np.ma.masked_all((3, 3))
169 >>> a.dtype
170 dtype('float64')
171 >>> a = np.ma.masked_all((3, 3), dtype=np.int32)
172 >>> a.dtype
173 dtype('int32')
174
175 """
176 a = masked_array(np.empty(shape, dtype),
177 mask=np.ones(shape, make_mask_descr(dtype)))
178 return a
179
180
181def masked_all_like(arr):
182 """
183 Empty masked array with the properties of an existing array.
184
185 Return an empty masked array of the same shape and dtype as
186 the array `arr`, where all the data are masked.
187
188 Parameters
189 ----------
190 arr : ndarray
191 An array describing the shape and dtype of the required MaskedArray.
192
193 Returns
194 -------
195 a : MaskedArray
196 A masked array with all data masked.
197
198 Raises
199 ------
200 AttributeError
201 If `arr` doesn't have a shape attribute (i.e. not an ndarray)
202
203 See Also
204 --------
205 masked_all : Empty masked array with all elements masked.
206
207 Notes
208 -----
209 Unlike other masked array creation functions (e.g. `numpy.ma.zeros_like`,
210 `numpy.ma.ones_like`, `numpy.ma.full_like`), `masked_all_like` does not
211 initialize the values of the array, and may therefore be marginally
212 faster. However, the values stored in the newly allocated array are
213 arbitrary. For reproducible behavior, be sure to set each element of the
214 array before reading.
215
216 Examples
217 --------
218 >>> import numpy as np
219 >>> arr = np.zeros((2, 3), dtype=np.float32)
220 >>> arr
221 array([[0., 0., 0.],
222 [0., 0., 0.]], dtype=float32)
223 >>> np.ma.masked_all_like(arr)
224 masked_array(
225 data=[[--, --, --],
226 [--, --, --]],
227 mask=[[ True, True, True],
228 [ True, True, True]],
229 fill_value=np.float64(1e+20),
230 dtype=float32)
231
232 The dtype of the masked array matches the dtype of `arr`.
233
234 >>> arr.dtype
235 dtype('float32')
236 >>> np.ma.masked_all_like(arr).dtype
237 dtype('float32')
238
239 """
240 a = np.empty_like(arr).view(MaskedArray)
241 a._mask = np.ones(a.shape, dtype=make_mask_descr(a.dtype))
242 return a
243
244
245#####--------------------------------------------------------------------------
246#---- --- Standard functions ---
247#####--------------------------------------------------------------------------
248
249def _fromnxfunction_function(_fromnxfunction):
250 """
251 Decorator to wrap a "_fromnxfunction" function, wrapping a numpy function as a
252 masked array function, with proper docstring and name.
253
254 Parameters
255 ----------
256 _fromnxfunction : ({params}) -> ndarray, {params}) -> masked_array
257 Wrapper function that calls the wrapped numpy function
258
259 Returns
260 -------
261 decorator : (f: ({params}) -> ndarray) -> ({params}) -> masked_array
262 Function that accepts a numpy function and returns a masked array function
263
264 """
265 def decorator(npfunc, /):
266 def wrapper(*args, **kwargs):
267 return _fromnxfunction(npfunc, *args, **kwargs)
268
269 functools.update_wrapper(wrapper, npfunc, assigned=("__name__", "__qualname__"))
270 wrapper.__doc__ = ma.doc_note(
271 npfunc.__doc__,
272 "The function is applied to both the ``_data`` and the ``_mask``, if any.",
273 )
274 return wrapper
275
276 return decorator
277
278
279@_fromnxfunction_function
280def _fromnxfunction_single(npfunc, a, /, *args, **kwargs):
281 """
282 Wraps a NumPy function that can be called with a single array argument followed by
283 auxiliary args that are passed verbatim for both the data and mask calls.
284 """
285 return masked_array(
286 data=npfunc(np.asarray(a), *args, **kwargs),
287 mask=npfunc(getmaskarray(a), *args, **kwargs),
288 )
289
290
291@_fromnxfunction_function
292def _fromnxfunction_seq(npfunc, arys, /, *args, **kwargs):
293 """
294 Wraps a NumPy function that can be called with a single sequence of arrays followed
295 by auxiliary args that are passed verbatim for both the data and mask calls.
296 """
297 return masked_array(
298 data=npfunc(tuple(np.asarray(a) for a in arys), *args, **kwargs),
299 mask=npfunc(tuple(getmaskarray(a) for a in arys), *args, **kwargs),
300 )
301
302@_fromnxfunction_function
303def _fromnxfunction_allargs(npfunc, /, *arys, **kwargs):
304 """
305 Wraps a NumPy function that can be called with multiple array arguments.
306 All args are converted to arrays even if they are not so already.
307 This makes it possible to process scalars as 1-D arrays.
308 Only keyword arguments are passed through verbatim for the data and mask calls.
309 Arrays arguments are processed independently and the results are returned in a list.
310 If only one arg is present, the return value is just the processed array instead of
311 a list.
312 """
313 out = tuple(
314 masked_array(
315 data=npfunc(np.asarray(a), **kwargs),
316 mask=npfunc(getmaskarray(a), **kwargs),
317 )
318 for a in arys
319 )
320 return out[0] if len(out) == 1 else out
321
322
323atleast_1d = _fromnxfunction_allargs(np.atleast_1d)
324atleast_2d = _fromnxfunction_allargs(np.atleast_2d)
325atleast_3d = _fromnxfunction_allargs(np.atleast_3d)
326
327vstack = row_stack = _fromnxfunction_seq(np.vstack)
328hstack = _fromnxfunction_seq(np.hstack)
329column_stack = _fromnxfunction_seq(np.column_stack)
330dstack = _fromnxfunction_seq(np.dstack)
331stack = _fromnxfunction_seq(np.stack)
332
333hsplit = _fromnxfunction_single(np.hsplit)
334diagflat = _fromnxfunction_single(np.diagflat)
335
336
337#####--------------------------------------------------------------------------
338#----
339#####--------------------------------------------------------------------------
340def flatten_inplace(seq):
341 """Flatten a sequence in place."""
342 k = 0
343 while (k != len(seq)):
344 while hasattr(seq[k], '__iter__'):
345 seq[k:(k + 1)] = seq[k]
346 k += 1
347 return seq
348
349
350def apply_along_axis(func1d, axis, arr, *args, **kwargs):
351 """
352 (This docstring should be overwritten)
353 """
354 arr = array(arr, copy=False, subok=True)
355 nd = arr.ndim
356 axis = normalize_axis_index(axis, nd)
357 ind = [0] * (nd - 1)
358 i = np.zeros(nd, 'O')
359 indlist = list(range(nd))
360 indlist.remove(axis)
361 i[axis] = slice(None, None)
362 outshape = np.asarray(arr.shape).take(indlist)
363 i.put(indlist, ind)
364 res = func1d(arr[tuple(i.tolist())], *args, **kwargs)
365 # if res is a number, then we have a smaller output array
366 asscalar = np.isscalar(res)
367 if not asscalar:
368 try:
369 len(res)
370 except TypeError:
371 asscalar = True
372 # Note: we shouldn't set the dtype of the output from the first result
373 # so we force the type to object, and build a list of dtypes. We'll
374 # just take the largest, to avoid some downcasting
375 dtypes = []
376 if asscalar:
377 dtypes.append(np.asarray(res).dtype)
378 outarr = zeros(outshape, object)
379 outarr[tuple(ind)] = res
380 Ntot = np.prod(outshape)
381 k = 1
382 while k < Ntot:
383 # increment the index
384 ind[-1] += 1
385 n = -1
386 while (ind[n] >= outshape[n]) and (n > (1 - nd)):
387 ind[n - 1] += 1
388 ind[n] = 0
389 n -= 1
390 i.put(indlist, ind)
391 res = func1d(arr[tuple(i.tolist())], *args, **kwargs)
392 outarr[tuple(ind)] = res
393 dtypes.append(asarray(res).dtype)
394 k += 1
395 else:
396 res = array(res, copy=False, subok=True)
397 j = i.copy()
398 j[axis] = ([slice(None, None)] * res.ndim)
399 j.put(indlist, ind)
400 Ntot = np.prod(outshape)
401 holdshape = outshape
402 outshape = list(arr.shape)
403 outshape[axis] = res.shape
404 dtypes.append(asarray(res).dtype)
405 outshape = flatten_inplace(outshape)
406 outarr = zeros(outshape, object)
407 outarr[tuple(flatten_inplace(j.tolist()))] = res
408 k = 1
409 while k < Ntot:
410 # increment the index
411 ind[-1] += 1
412 n = -1
413 while (ind[n] >= holdshape[n]) and (n > (1 - nd)):
414 ind[n - 1] += 1
415 ind[n] = 0
416 n -= 1
417 i.put(indlist, ind)
418 j.put(indlist, ind)
419 res = func1d(arr[tuple(i.tolist())], *args, **kwargs)
420 outarr[tuple(flatten_inplace(j.tolist()))] = res
421 dtypes.append(asarray(res).dtype)
422 k += 1
423 max_dtypes = np.dtype(np.asarray(dtypes).max())
424 if not hasattr(arr, '_mask'):
425 result = np.asarray(outarr, dtype=max_dtypes)
426 else:
427 result = asarray(outarr, dtype=max_dtypes)
428 result.fill_value = ma.default_fill_value(result)
429 return result
430
431
432apply_along_axis.__doc__ = np.apply_along_axis.__doc__
433
434
435def apply_over_axes(func, a, axes):
436 """
437 (This docstring will be overwritten)
438 """
439 val = asarray(a)
440 N = a.ndim
441 if array(axes).ndim == 0:
442 axes = (axes,)
443 for axis in axes:
444 if axis < 0:
445 axis = N + axis
446 args = (val, axis)
447 res = func(*args)
448 if res.ndim == val.ndim:
449 val = res
450 else:
451 res = ma.expand_dims(res, axis)
452 if res.ndim == val.ndim:
453 val = res
454 else:
455 raise ValueError("function is not returning "
456 "an array of the correct shape")
457 return val
458
459
460if apply_over_axes.__doc__ is not None:
461 apply_over_axes.__doc__ = np.apply_over_axes.__doc__[
462 :np.apply_over_axes.__doc__.find('Notes')].rstrip() + \
463 """
464
465 Examples
466 --------
467 >>> import numpy as np
468 >>> a = np.ma.arange(24).reshape(2,3,4)
469 >>> a[:,0,1] = np.ma.masked
470 >>> a[:,1,:] = np.ma.masked
471 >>> a
472 masked_array(
473 data=[[[0, --, 2, 3],
474 [--, --, --, --],
475 [8, 9, 10, 11]],
476 [[12, --, 14, 15],
477 [--, --, --, --],
478 [20, 21, 22, 23]]],
479 mask=[[[False, True, False, False],
480 [ True, True, True, True],
481 [False, False, False, False]],
482 [[False, True, False, False],
483 [ True, True, True, True],
484 [False, False, False, False]]],
485 fill_value=999999)
486 >>> np.ma.apply_over_axes(np.ma.sum, a, [0,2])
487 masked_array(
488 data=[[[46],
489 [--],
490 [124]]],
491 mask=[[[False],
492 [ True],
493 [False]]],
494 fill_value=999999)
495
496 Tuple axis arguments to ufuncs are equivalent:
497
498 >>> np.ma.sum(a, axis=(0,2)).reshape((1,-1,1))
499 masked_array(
500 data=[[[46],
501 [--],
502 [124]]],
503 mask=[[[False],
504 [ True],
505 [False]]],
506 fill_value=999999)
507 """
508
509
510def average(a, axis=None, weights=None, returned=False, *,
511 keepdims=np._NoValue):
512 """
513 Return the weighted average of array over the given axis.
514
515 Parameters
516 ----------
517 a : array_like
518 Data to be averaged.
519 Masked entries are not taken into account in the computation.
520 axis : None or int or tuple of ints, optional
521 Axis or axes along which to average `a`. The default,
522 `axis=None`, will average over all of the elements of the input array.
523 If axis is a tuple of ints, averaging is performed on all of the axes
524 specified in the tuple instead of a single axis or all the axes as
525 before.
526 weights : array_like, optional
527 An array of weights associated with the values in `a`. Each value in
528 `a` contributes to the average according to its associated weight.
529 The array of weights must be the same shape as `a` if no axis is
530 specified, otherwise the weights must have dimensions and shape
531 consistent with `a` along the specified axis.
532 If `weights=None`, then all data in `a` are assumed to have a
533 weight equal to one.
534 The calculation is::
535
536 avg = sum(a * weights) / sum(weights)
537
538 where the sum is over all included elements.
539 The only constraint on the values of `weights` is that `sum(weights)`
540 must not be 0.
541 returned : bool, optional
542 Flag indicating whether a tuple ``(result, sum of weights)``
543 should be returned as output (True), or just the result (False).
544 Default is False.
545 keepdims : bool, optional
546 If this is set to True, the axes which are reduced are left
547 in the result as dimensions with size one. With this option,
548 the result will broadcast correctly against the original `a`.
549 *Note:* `keepdims` will not work with instances of `numpy.matrix`
550 or other classes whose methods do not support `keepdims`.
551
552 .. versionadded:: 1.23.0
553
554 Returns
555 -------
556 average, [sum_of_weights] : (tuple of) scalar or MaskedArray
557 The average along the specified axis. When returned is `True`,
558 return a tuple with the average as the first element and the sum
559 of the weights as the second element. The return type is `np.float64`
560 if `a` is of integer type and floats smaller than `float64`, or the
561 input data-type, otherwise. If returned, `sum_of_weights` is always
562 `float64`.
563
564 Raises
565 ------
566 ZeroDivisionError
567 When all weights along axis are zero. See `numpy.ma.average` for a
568 version robust to this type of error.
569 TypeError
570 When `weights` does not have the same shape as `a`, and `axis=None`.
571 ValueError
572 When `weights` does not have dimensions and shape consistent with `a`
573 along specified `axis`.
574
575 Examples
576 --------
577 >>> import numpy as np
578 >>> a = np.ma.array([1., 2., 3., 4.], mask=[False, False, True, True])
579 >>> np.ma.average(a, weights=[3, 1, 0, 0])
580 1.25
581
582 >>> x = np.ma.arange(6.).reshape(3, 2)
583 >>> x
584 masked_array(
585 data=[[0., 1.],
586 [2., 3.],
587 [4., 5.]],
588 mask=False,
589 fill_value=1e+20)
590 >>> data = np.arange(8).reshape((2, 2, 2))
591 >>> data
592 array([[[0, 1],
593 [2, 3]],
594 [[4, 5],
595 [6, 7]]])
596 >>> np.ma.average(data, axis=(0, 1), weights=[[1./4, 3./4], [1., 1./2]])
597 masked_array(data=[3.4, 4.4],
598 mask=[False, False],
599 fill_value=1e+20)
600 >>> np.ma.average(data, axis=0, weights=[[1./4, 3./4], [1., 1./2]])
601 Traceback (most recent call last):
602 ...
603 ValueError: Shape of weights must be consistent
604 with shape of a along specified axis.
605
606 >>> avg, sumweights = np.ma.average(x, axis=0, weights=[1, 2, 3],
607 ... returned=True)
608 >>> avg
609 masked_array(data=[2.6666666666666665, 3.6666666666666665],
610 mask=[False, False],
611 fill_value=1e+20)
612
613 With ``keepdims=True``, the following result has shape (3, 1).
614
615 >>> np.ma.average(x, axis=1, keepdims=True)
616 masked_array(
617 data=[[0.5],
618 [2.5],
619 [4.5]],
620 mask=False,
621 fill_value=1e+20)
622 """
623 a = asarray(a)
624 m = getmask(a)
625
626 if axis is not None:
627 axis = normalize_axis_tuple(axis, a.ndim, argname="axis")
628
629 if keepdims is np._NoValue:
630 # Don't pass on the keepdims argument if one wasn't given.
631 keepdims_kw = {}
632 else:
633 keepdims_kw = {'keepdims': keepdims}
634
635 if weights is None:
636 avg = a.mean(axis, **keepdims_kw)
637 scl = avg.dtype.type(a.count(axis))
638 else:
639 wgt = asarray(weights)
640
641 if issubclass(a.dtype.type, (np.integer, np.bool)):
642 result_dtype = np.result_type(a.dtype, wgt.dtype, 'f8')
643 else:
644 result_dtype = np.result_type(a.dtype, wgt.dtype)
645
646 # Sanity checks
647 if a.shape != wgt.shape:
648 if axis is None:
649 raise TypeError(
650 "Axis must be specified when shapes of a and weights "
651 "differ.")
652 if wgt.shape != tuple(a.shape[ax] for ax in axis):
653 raise ValueError(
654 "Shape of weights must be consistent with "
655 "shape of a along specified axis.")
656
657 # setup wgt to broadcast along axis
658 wgt = wgt.transpose(np.argsort(axis))
659 wgt = wgt.reshape(tuple((s if ax in axis else 1)
660 for ax, s in enumerate(a.shape)))
661
662 if m is not nomask:
663 wgt = wgt * (~a.mask)
664 wgt.mask |= a.mask
665
666 scl = wgt.sum(axis=axis, dtype=result_dtype, **keepdims_kw)
667 avg = np.multiply(a, wgt,
668 dtype=result_dtype).sum(axis, **keepdims_kw) / scl
669
670 if returned:
671 if scl.shape != avg.shape:
672 scl = np.broadcast_to(scl, avg.shape).copy()
673 return avg, scl
674 else:
675 return avg
676
677
678def median(a, axis=None, out=None, overwrite_input=False, keepdims=False):
679 """
680 Compute the median along the specified axis.
681
682 Returns the median of the array elements.
683
684 Parameters
685 ----------
686 a : array_like
687 Input array or object that can be converted to an array.
688 axis : int, optional
689 Axis along which the medians are computed. The default (None) is
690 to compute the median along a flattened version of the array.
691 out : ndarray, optional
692 Alternative output array in which to place the result. It must
693 have the same shape and buffer length as the expected output
694 but the type will be cast if necessary.
695 overwrite_input : bool, optional
696 If True, then allow use of memory of input array (a) for
697 calculations. The input array will be modified by the call to
698 median. This will save memory when you do not need to preserve
699 the contents of the input array. Treat the input as undefined,
700 but it will probably be fully or partially sorted. Default is
701 False. Note that, if `overwrite_input` is True, and the input
702 is not already an `ndarray`, an error will be raised.
703 keepdims : bool, optional
704 If this is set to True, the axes which are reduced are left
705 in the result as dimensions with size one. With this option,
706 the result will broadcast correctly against the input array.
707
708 Returns
709 -------
710 median : ndarray
711 A new array holding the result is returned unless out is
712 specified, in which case a reference to out is returned.
713 Return data-type is `float64` for integers and floats smaller than
714 `float64`, or the input data-type, otherwise.
715
716 See Also
717 --------
718 mean
719
720 Notes
721 -----
722 Given a vector ``V`` with ``N`` non masked values, the median of ``V``
723 is the middle value of a sorted copy of ``V`` (``Vs``) - i.e.
724 ``Vs[(N-1)/2]``, when ``N`` is odd, or ``{Vs[N/2 - 1] + Vs[N/2]}/2``
725 when ``N`` is even.
726
727 Examples
728 --------
729 >>> import numpy as np
730 >>> x = np.ma.array(np.arange(8), mask=[0]*4 + [1]*4)
731 >>> np.ma.median(x)
732 1.5
733
734 >>> x = np.ma.array(np.arange(10).reshape(2, 5), mask=[0]*6 + [1]*4)
735 >>> np.ma.median(x)
736 2.5
737 >>> np.ma.median(x, axis=-1, overwrite_input=True)
738 masked_array(data=[2.0, 5.0],
739 mask=[False, False],
740 fill_value=1e+20)
741
742 """
743 if not hasattr(a, 'mask'):
744 m = np.median(getdata(a, subok=True), axis=axis,
745 out=out, overwrite_input=overwrite_input,
746 keepdims=keepdims)
747 if isinstance(m, np.ndarray) and 1 <= m.ndim:
748 return masked_array(m, copy=False)
749 else:
750 return m
751
752 return _ureduce(a, func=_median, keepdims=keepdims, axis=axis, out=out,
753 overwrite_input=overwrite_input)
754
755
756def _median(a, axis=None, out=None, overwrite_input=False):
757 # when an unmasked NaN is present return it, so we need to sort the NaN
758 # values behind the mask
759 if np.issubdtype(a.dtype, np.inexact):
760 fill_value = np.inf
761 else:
762 fill_value = None
763 if overwrite_input:
764 if axis is None:
765 asorted = a.ravel()
766 asorted.sort(fill_value=fill_value)
767 else:
768 a.sort(axis=axis, fill_value=fill_value)
769 asorted = a
770 else:
771 asorted = sort(a, axis=axis, fill_value=fill_value)
772
773 if axis is None:
774 axis = 0
775 else:
776 axis = normalize_axis_index(axis, asorted.ndim)
777
778 if asorted.shape[axis] == 0:
779 # for empty axis integer indices fail so use slicing to get same result
780 # as median (which is mean of empty slice = nan)
781 indexer = [slice(None)] * asorted.ndim
782 indexer[axis] = slice(0, 0)
783 indexer = tuple(indexer)
784 return np.ma.mean(asorted[indexer], axis=axis, out=out)
785
786 if asorted.ndim == 1:
787 idx, odd = divmod(count(asorted), 2)
788 mid = asorted[idx + odd - 1:idx + 1]
789 if np.issubdtype(asorted.dtype, np.inexact) and asorted.size > 0:
790 # avoid inf / x = masked
791 s = mid.sum(out=out)
792 if not odd:
793 s = np.true_divide(s, 2., casting='safe', out=out)
794 s = np.lib._utils_impl._median_nancheck(asorted, s, axis)
795 else:
796 s = mid.mean(out=out)
797
798 # if result is masked either the input contained enough
799 # minimum_fill_value so that it would be the median or all values
800 # masked
801 if np.ma.is_masked(s) and not np.all(asorted.mask):
802 return np.ma.minimum_fill_value(asorted)
803 return s
804
805 counts = count(asorted, axis=axis, keepdims=True)
806 h = counts // 2
807
808 # duplicate high if odd number of elements so mean does nothing
809 odd = counts % 2 == 1
810 l = np.where(odd, h, h - 1)
811
812 lh = np.concatenate([l, h], axis=axis)
813
814 # get low and high median
815 low_high = np.take_along_axis(asorted, lh, axis=axis)
816
817 def replace_masked(s):
818 # Replace masked entries with minimum_full_value unless it all values
819 # are masked. This is required as the sort order of values equal or
820 # larger than the fill value is undefined and a valid value placed
821 # elsewhere, e.g. [4, --, inf].
822 if np.ma.is_masked(s):
823 rep = (~np.all(asorted.mask, axis=axis, keepdims=True)) & s.mask
824 s.data[rep] = np.ma.minimum_fill_value(asorted)
825 s.mask[rep] = False
826
827 replace_masked(low_high)
828
829 if np.issubdtype(asorted.dtype, np.inexact):
830 # avoid inf / x = masked
831 s = np.ma.sum(low_high, axis=axis, out=out)
832 np.true_divide(s.data, 2., casting='unsafe', out=s.data)
833
834 s = np.lib._utils_impl._median_nancheck(asorted, s, axis)
835 else:
836 s = np.ma.mean(low_high, axis=axis, out=out)
837
838 return s
839
840
841def compress_nd(x, axis=None):
842 """Suppress slices from multiple dimensions which contain masked values.
843
844 Parameters
845 ----------
846 x : array_like, MaskedArray
847 The array to operate on. If not a MaskedArray instance (or if no array
848 elements are masked), `x` is interpreted as a MaskedArray with `mask`
849 set to `nomask`.
850 axis : tuple of ints or int, optional
851 Which dimensions to suppress slices from can be configured with this
852 parameter.
853 - If axis is a tuple of ints, those are the axes to suppress slices from.
854 - If axis is an int, then that is the only axis to suppress slices from.
855 - If axis is None, all axis are selected.
856
857 Returns
858 -------
859 compress_array : ndarray
860 The compressed array.
861
862 Examples
863 --------
864 >>> import numpy as np
865 >>> arr = [[1, 2], [3, 4]]
866 >>> mask = [[0, 1], [0, 0]]
867 >>> x = np.ma.array(arr, mask=mask)
868 >>> np.ma.compress_nd(x, axis=0)
869 array([[3, 4]])
870 >>> np.ma.compress_nd(x, axis=1)
871 array([[1],
872 [3]])
873 >>> np.ma.compress_nd(x)
874 array([[3]])
875
876 """
877 x = asarray(x)
878 m = getmask(x)
879 # Set axis to tuple of ints
880 if axis is None:
881 axis = tuple(range(x.ndim))
882 else:
883 axis = normalize_axis_tuple(axis, x.ndim)
884
885 # Nothing is masked: return x
886 if m is nomask or not m.any():
887 return x._data
888 # All is masked: return empty
889 if m.all():
890 return nxarray([])
891 # Filter elements through boolean indexing
892 data = x._data
893 for ax in axis:
894 axes = tuple(list(range(ax)) + list(range(ax + 1, x.ndim)))
895 data = data[(slice(None),) * ax + (~m.any(axis=axes),)]
896 return data
897
898
899def compress_rowcols(x, axis=None):
900 """
901 Suppress the rows and/or columns of a 2-D array that contain
902 masked values.
903
904 The suppression behavior is selected with the `axis` parameter.
905
906 - If axis is None, both rows and columns are suppressed.
907 - If axis is 0, only rows are suppressed.
908 - If axis is 1 or -1, only columns are suppressed.
909
910 Parameters
911 ----------
912 x : array_like, MaskedArray
913 The array to operate on. If not a MaskedArray instance (or if no array
914 elements are masked), `x` is interpreted as a MaskedArray with
915 `mask` set to `nomask`. Must be a 2D array.
916 axis : int, optional
917 Axis along which to perform the operation. Default is None.
918
919 Returns
920 -------
921 compressed_array : ndarray
922 The compressed array.
923
924 Examples
925 --------
926 >>> import numpy as np
927 >>> x = np.ma.array(np.arange(9).reshape(3, 3), mask=[[1, 0, 0],
928 ... [1, 0, 0],
929 ... [0, 0, 0]])
930 >>> x
931 masked_array(
932 data=[[--, 1, 2],
933 [--, 4, 5],
934 [6, 7, 8]],
935 mask=[[ True, False, False],
936 [ True, False, False],
937 [False, False, False]],
938 fill_value=999999)
939
940 >>> np.ma.compress_rowcols(x)
941 array([[7, 8]])
942 >>> np.ma.compress_rowcols(x, 0)
943 array([[6, 7, 8]])
944 >>> np.ma.compress_rowcols(x, 1)
945 array([[1, 2],
946 [4, 5],
947 [7, 8]])
948
949 """
950 if asarray(x).ndim != 2:
951 raise NotImplementedError("compress_rowcols works for 2D arrays only.")
952 return compress_nd(x, axis=axis)
953
954
955def compress_rows(a):
956 """
957 Suppress whole rows of a 2-D array that contain masked values.
958
959 This is equivalent to ``np.ma.compress_rowcols(a, 0)``, see
960 `compress_rowcols` for details.
961
962 Parameters
963 ----------
964 x : array_like, MaskedArray
965 The array to operate on. If not a MaskedArray instance (or if no array
966 elements are masked), `x` is interpreted as a MaskedArray with
967 `mask` set to `nomask`. Must be a 2D array.
968
969 Returns
970 -------
971 compressed_array : ndarray
972 The compressed array.
973
974 See Also
975 --------
976 compress_rowcols
977
978 Examples
979 --------
980 >>> import numpy as np
981 >>> a = np.ma.array(np.arange(9).reshape(3, 3), mask=[[1, 0, 0],
982 ... [1, 0, 0],
983 ... [0, 0, 0]])
984 >>> np.ma.compress_rows(a)
985 array([[6, 7, 8]])
986
987 """
988 a = asarray(a)
989 if a.ndim != 2:
990 raise NotImplementedError("compress_rows works for 2D arrays only.")
991 return compress_rowcols(a, 0)
992
993
994def compress_cols(a):
995 """
996 Suppress whole columns of a 2-D array that contain masked values.
997
998 This is equivalent to ``np.ma.compress_rowcols(a, 1)``, see
999 `compress_rowcols` for details.
1000
1001 Parameters
1002 ----------
1003 x : array_like, MaskedArray
1004 The array to operate on. If not a MaskedArray instance (or if no array
1005 elements are masked), `x` is interpreted as a MaskedArray with
1006 `mask` set to `nomask`. Must be a 2D array.
1007
1008 Returns
1009 -------
1010 compressed_array : ndarray
1011 The compressed array.
1012
1013 See Also
1014 --------
1015 compress_rowcols
1016
1017 Examples
1018 --------
1019 >>> import numpy as np
1020 >>> a = np.ma.array(np.arange(9).reshape(3, 3), mask=[[1, 0, 0],
1021 ... [1, 0, 0],
1022 ... [0, 0, 0]])
1023 >>> np.ma.compress_cols(a)
1024 array([[1, 2],
1025 [4, 5],
1026 [7, 8]])
1027
1028 """
1029 a = asarray(a)
1030 if a.ndim != 2:
1031 raise NotImplementedError("compress_cols works for 2D arrays only.")
1032 return compress_rowcols(a, 1)
1033
1034
1035def mask_rowcols(a, axis=None):
1036 """
1037 Mask rows and/or columns of a 2D array that contain masked values.
1038
1039 Mask whole rows and/or columns of a 2D array that contain
1040 masked values. The masking behavior is selected using the
1041 `axis` parameter.
1042
1043 - If `axis` is None, rows *and* columns are masked.
1044 - If `axis` is 0, only rows are masked.
1045 - If `axis` is 1 or -1, only columns are masked.
1046
1047 Parameters
1048 ----------
1049 a : array_like, MaskedArray
1050 The array to mask. If not a MaskedArray instance (or if no array
1051 elements are masked), the result is a MaskedArray with `mask` set
1052 to `nomask` (False). Must be a 2D array.
1053 axis : int, optional
1054 Axis along which to perform the operation. If None, applies to a
1055 flattened version of the array.
1056
1057 Returns
1058 -------
1059 a : MaskedArray
1060 A modified version of the input array, masked depending on the value
1061 of the `axis` parameter.
1062
1063 Raises
1064 ------
1065 NotImplementedError
1066 If input array `a` is not 2D.
1067
1068 See Also
1069 --------
1070 mask_rows : Mask rows of a 2D array that contain masked values.
1071 mask_cols : Mask cols of a 2D array that contain masked values.
1072 masked_where : Mask where a condition is met.
1073
1074 Notes
1075 -----
1076 The input array's mask is modified by this function.
1077
1078 Examples
1079 --------
1080 >>> import numpy as np
1081 >>> a = np.zeros((3, 3), dtype=int)
1082 >>> a[1, 1] = 1
1083 >>> a
1084 array([[0, 0, 0],
1085 [0, 1, 0],
1086 [0, 0, 0]])
1087 >>> a = np.ma.masked_equal(a, 1)
1088 >>> a
1089 masked_array(
1090 data=[[0, 0, 0],
1091 [0, --, 0],
1092 [0, 0, 0]],
1093 mask=[[False, False, False],
1094 [False, True, False],
1095 [False, False, False]],
1096 fill_value=1)
1097 >>> np.ma.mask_rowcols(a)
1098 masked_array(
1099 data=[[0, --, 0],
1100 [--, --, --],
1101 [0, --, 0]],
1102 mask=[[False, True, False],
1103 [ True, True, True],
1104 [False, True, False]],
1105 fill_value=1)
1106
1107 """
1108 a = array(a, subok=False)
1109 if a.ndim != 2:
1110 raise NotImplementedError("mask_rowcols works for 2D arrays only.")
1111 m = getmask(a)
1112 # Nothing is masked: return a
1113 if m is nomask or not m.any():
1114 return a
1115 maskedval = m.nonzero()
1116 a._mask = a._mask.copy()
1117 if not axis:
1118 a[np.unique(maskedval[0])] = masked
1119 if axis in [None, 1, -1]:
1120 a[:, np.unique(maskedval[1])] = masked
1121 return a
1122
1123
1124def mask_rows(a, axis=np._NoValue):
1125 """
1126 Mask rows of a 2D array that contain masked values.
1127
1128 This function is a shortcut to ``mask_rowcols`` with `axis` equal to 0.
1129
1130 See Also
1131 --------
1132 mask_rowcols : Mask rows and/or columns of a 2D array.
1133 masked_where : Mask where a condition is met.
1134
1135 Examples
1136 --------
1137 >>> import numpy as np
1138 >>> a = np.zeros((3, 3), dtype=int)
1139 >>> a[1, 1] = 1
1140 >>> a
1141 array([[0, 0, 0],
1142 [0, 1, 0],
1143 [0, 0, 0]])
1144 >>> a = np.ma.masked_equal(a, 1)
1145 >>> a
1146 masked_array(
1147 data=[[0, 0, 0],
1148 [0, --, 0],
1149 [0, 0, 0]],
1150 mask=[[False, False, False],
1151 [False, True, False],
1152 [False, False, False]],
1153 fill_value=1)
1154
1155 >>> np.ma.mask_rows(a)
1156 masked_array(
1157 data=[[0, 0, 0],
1158 [--, --, --],
1159 [0, 0, 0]],
1160 mask=[[False, False, False],
1161 [ True, True, True],
1162 [False, False, False]],
1163 fill_value=1)
1164
1165 """
1166 if axis is not np._NoValue:
1167 # remove the axis argument when this deprecation expires
1168 # NumPy 1.18.0, 2019-11-28
1169 warnings.warn(
1170 "The axis argument has always been ignored, in future passing it "
1171 "will raise TypeError", DeprecationWarning, stacklevel=2)
1172 return mask_rowcols(a, 0)
1173
1174
1175def mask_cols(a, axis=np._NoValue):
1176 """
1177 Mask columns of a 2D array that contain masked values.
1178
1179 This function is a shortcut to ``mask_rowcols`` with `axis` equal to 1.
1180
1181 See Also
1182 --------
1183 mask_rowcols : Mask rows and/or columns of a 2D array.
1184 masked_where : Mask where a condition is met.
1185
1186 Examples
1187 --------
1188 >>> import numpy as np
1189 >>> a = np.zeros((3, 3), dtype=int)
1190 >>> a[1, 1] = 1
1191 >>> a
1192 array([[0, 0, 0],
1193 [0, 1, 0],
1194 [0, 0, 0]])
1195 >>> a = np.ma.masked_equal(a, 1)
1196 >>> a
1197 masked_array(
1198 data=[[0, 0, 0],
1199 [0, --, 0],
1200 [0, 0, 0]],