extras.py 66 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753175417551756175717581759176017611762176317641765176617671768176917701771177217731774177517761777177817791780178117821783178417851786178717881789179017911792179317941795179617971798179918001801180218031804180518061807180818091810181118121813181418151816181718181819182018211822182318241825182618271828182918301831183218331834183518361837183818391840184118421843184418451846184718481849185018511852185318541855185618571858185918601861186218631864186518661867186818691870187118721873187418751876187718781879188018811882188318841885188618871888188918901891189218931894189518961897189818991900190119021903190419051906190719081909191019111912191319141915191619171918191919201921192219231924192519261927192819291930193119321933193419351936193719381939194019411942194319441945194619471948194919501951195219531954195519561957195819591960196119621963196419651966196719681969197019711972197319741975197619771978197919801981198219831984198519861987198819891990199119921993199419951996199719981999200020012002200320042005200620072008200920102011201220132014201520162017201820192020202120222023202420252026202720282029203020312032203320342035203620372038203920402041204220432044204520462047204820492050205120522053205420552056205720582059206020612062206320642065206620672068206920702071207220732074207520762077207820792080208120822083208420852086208720882089209020912092209320942095209620972098209921002101210221032104210521062107210821092110211121122113211421152116211721182119212021212122212321242125212621272128212921302131213221332134213521362137213821392140214121422143214421452146214721482149215021512152215321542155215621572158215921602161216221632164216521662167216821692170217121722173217421752176217721782179218021812182218321842185218621872188218921902191219221932194219521962197219821992200220122022203220422052206220722082209221022112212221322142215221622172218221922202221222222232224222522262227222822292230223122322233223422352236223722382239224022412242224322442245224622472248224922502251225222532254225522562257225822592260226122622263226422652266
  1. """
  2. Masked arrays add-ons.
  3. A collection of utilities for `numpy.ma`.
  4. :author: Pierre Gerard-Marchant
  5. :contact: pierregm_at_uga_dot_edu
  6. """
  7. __all__ = [
  8. 'apply_along_axis', 'apply_over_axes', 'atleast_1d', 'atleast_2d',
  9. 'atleast_3d', 'average', 'clump_masked', 'clump_unmasked', 'column_stack',
  10. 'compress_cols', 'compress_nd', 'compress_rowcols', 'compress_rows',
  11. 'count_masked', 'corrcoef', 'cov', 'diagflat', 'dot', 'dstack', 'ediff1d',
  12. 'flatnotmasked_contiguous', 'flatnotmasked_edges', 'hsplit', 'hstack',
  13. 'isin', 'in1d', 'intersect1d', 'mask_cols', 'mask_rowcols', 'mask_rows',
  14. 'masked_all', 'masked_all_like', 'median', 'mr_', 'ndenumerate',
  15. 'notmasked_contiguous', 'notmasked_edges', 'polyfit', 'row_stack',
  16. 'setdiff1d', 'setxor1d', 'stack', 'unique', 'union1d', 'vander', 'vstack',
  17. ]
  18. import functools
  19. import itertools
  20. import warnings
  21. import numpy as np
  22. from numpy import array as nxarray, ndarray
  23. from numpy.lib._function_base_impl import _ureduce
  24. from numpy.lib._index_tricks_impl import AxisConcatenator
  25. from numpy.lib.array_utils import normalize_axis_index, normalize_axis_tuple
  26. from . import core as ma
  27. from .core import ( # noqa: F401
  28. MAError,
  29. MaskedArray,
  30. add,
  31. array,
  32. asarray,
  33. concatenate,
  34. count,
  35. dot,
  36. filled,
  37. get_masked_subclass,
  38. getdata,
  39. getmask,
  40. getmaskarray,
  41. make_mask_descr,
  42. mask_or,
  43. masked,
  44. masked_array,
  45. nomask,
  46. ones,
  47. sort,
  48. zeros,
  49. )
  50. def issequence(seq):
  51. """
  52. Is seq a sequence (ndarray, list or tuple)?
  53. """
  54. return isinstance(seq, (ndarray, tuple, list))
  55. def count_masked(arr, axis=None):
  56. """
  57. Count the number of masked elements along the given axis.
  58. Parameters
  59. ----------
  60. arr : array_like
  61. An array with (possibly) masked elements.
  62. axis : int, optional
  63. Axis along which to count. If None (default), a flattened
  64. version of the array is used.
  65. Returns
  66. -------
  67. count : int, ndarray
  68. The total number of masked elements (axis=None) or the number
  69. of masked elements along each slice of the given axis.
  70. See Also
  71. --------
  72. MaskedArray.count : Count non-masked elements.
  73. Examples
  74. --------
  75. >>> import numpy as np
  76. >>> a = np.arange(9).reshape((3,3))
  77. >>> a = np.ma.array(a)
  78. >>> a[1, 0] = np.ma.masked
  79. >>> a[1, 2] = np.ma.masked
  80. >>> a[2, 1] = np.ma.masked
  81. >>> a
  82. masked_array(
  83. data=[[0, 1, 2],
  84. [--, 4, --],
  85. [6, --, 8]],
  86. mask=[[False, False, False],
  87. [ True, False, True],
  88. [False, True, False]],
  89. fill_value=999999)
  90. >>> np.ma.count_masked(a)
  91. 3
  92. When the `axis` keyword is used an array is returned.
  93. >>> np.ma.count_masked(a, axis=0)
  94. array([1, 1, 1])
  95. >>> np.ma.count_masked(a, axis=1)
  96. array([0, 2, 1])
  97. """
  98. m = getmaskarray(arr)
  99. return m.sum(axis)
  100. def masked_all(shape, dtype=float):
  101. """
  102. Empty masked array with all elements masked.
  103. Return an empty masked array of the given shape and dtype, where all the
  104. data are masked.
  105. Parameters
  106. ----------
  107. shape : int or tuple of ints
  108. Shape of the required MaskedArray, e.g., ``(2, 3)`` or ``2``.
  109. dtype : dtype, optional
  110. Data type of the output.
  111. Returns
  112. -------
  113. a : MaskedArray
  114. A masked array with all data masked.
  115. See Also
  116. --------
  117. masked_all_like : Empty masked array modelled on an existing array.
  118. Notes
  119. -----
  120. Unlike other masked array creation functions (e.g. `numpy.ma.zeros`,
  121. `numpy.ma.ones`, `numpy.ma.full`), `masked_all` does not initialize the
  122. values of the array, and may therefore be marginally faster. However,
  123. the values stored in the newly allocated array are arbitrary. For
  124. reproducible behavior, be sure to set each element of the array before
  125. reading.
  126. Examples
  127. --------
  128. >>> import numpy as np
  129. >>> np.ma.masked_all((3, 3))
  130. masked_array(
  131. data=[[--, --, --],
  132. [--, --, --],
  133. [--, --, --]],
  134. mask=[[ True, True, True],
  135. [ True, True, True],
  136. [ True, True, True]],
  137. fill_value=1e+20,
  138. dtype=float64)
  139. The `dtype` parameter defines the underlying data type.
  140. >>> a = np.ma.masked_all((3, 3))
  141. >>> a.dtype
  142. dtype('float64')
  143. >>> a = np.ma.masked_all((3, 3), dtype=np.int32)
  144. >>> a.dtype
  145. dtype('int32')
  146. """
  147. a = masked_array(np.empty(shape, dtype),
  148. mask=np.ones(shape, make_mask_descr(dtype)))
  149. return a
  150. def masked_all_like(arr):
  151. """
  152. Empty masked array with the properties of an existing array.
  153. Return an empty masked array of the same shape and dtype as
  154. the array `arr`, where all the data are masked.
  155. Parameters
  156. ----------
  157. arr : ndarray
  158. An array describing the shape and dtype of the required MaskedArray.
  159. Returns
  160. -------
  161. a : MaskedArray
  162. A masked array with all data masked.
  163. Raises
  164. ------
  165. AttributeError
  166. If `arr` doesn't have a shape attribute (i.e. not an ndarray)
  167. See Also
  168. --------
  169. masked_all : Empty masked array with all elements masked.
  170. Notes
  171. -----
  172. Unlike other masked array creation functions (e.g. `numpy.ma.zeros_like`,
  173. `numpy.ma.ones_like`, `numpy.ma.full_like`), `masked_all_like` does not
  174. initialize the values of the array, and may therefore be marginally
  175. faster. However, the values stored in the newly allocated array are
  176. arbitrary. For reproducible behavior, be sure to set each element of the
  177. array before reading.
  178. Examples
  179. --------
  180. >>> import numpy as np
  181. >>> arr = np.zeros((2, 3), dtype=np.float32)
  182. >>> arr
  183. array([[0., 0., 0.],
  184. [0., 0., 0.]], dtype=float32)
  185. >>> np.ma.masked_all_like(arr)
  186. masked_array(
  187. data=[[--, --, --],
  188. [--, --, --]],
  189. mask=[[ True, True, True],
  190. [ True, True, True]],
  191. fill_value=np.float64(1e+20),
  192. dtype=float32)
  193. The dtype of the masked array matches the dtype of `arr`.
  194. >>> arr.dtype
  195. dtype('float32')
  196. >>> np.ma.masked_all_like(arr).dtype
  197. dtype('float32')
  198. """
  199. a = np.empty_like(arr).view(MaskedArray)
  200. a._mask = np.ones(a.shape, dtype=make_mask_descr(a.dtype))
  201. return a
  202. #####--------------------------------------------------------------------------
  203. #---- --- Standard functions ---
  204. #####--------------------------------------------------------------------------
  205. def _fromnxfunction_function(_fromnxfunction):
  206. """
  207. Decorator to wrap a "_fromnxfunction" function, wrapping a numpy function as a
  208. masked array function, with proper docstring and name.
  209. Parameters
  210. ----------
  211. _fromnxfunction : ({params}) -> ndarray, {params}) -> masked_array
  212. Wrapper function that calls the wrapped numpy function
  213. Returns
  214. -------
  215. decorator : (f: ({params}) -> ndarray) -> ({params}) -> masked_array
  216. Function that accepts a numpy function and returns a masked array function
  217. """
  218. def decorator(npfunc, /):
  219. def wrapper(*args, **kwargs):
  220. return _fromnxfunction(npfunc, *args, **kwargs)
  221. functools.update_wrapper(wrapper, npfunc, assigned=("__name__", "__qualname__"))
  222. wrapper.__doc__ = ma.doc_note(
  223. npfunc.__doc__,
  224. "The function is applied to both the ``_data`` and the ``_mask``, if any.",
  225. )
  226. return wrapper
  227. return decorator
  228. @_fromnxfunction_function
  229. def _fromnxfunction_single(npfunc, a, /, *args, **kwargs):
  230. """
  231. Wraps a NumPy function that can be called with a single array argument followed by
  232. auxiliary args that are passed verbatim for both the data and mask calls.
  233. """
  234. return masked_array(
  235. data=npfunc(np.asarray(a), *args, **kwargs),
  236. mask=npfunc(getmaskarray(a), *args, **kwargs),
  237. )
  238. @_fromnxfunction_function
  239. def _fromnxfunction_seq(npfunc, arys, /, *args, **kwargs):
  240. """
  241. Wraps a NumPy function that can be called with a single sequence of arrays followed
  242. by auxiliary args that are passed verbatim for both the data and mask calls.
  243. """
  244. return masked_array(
  245. data=npfunc(tuple(np.asarray(a) for a in arys), *args, **kwargs),
  246. mask=npfunc(tuple(getmaskarray(a) for a in arys), *args, **kwargs),
  247. )
  248. @_fromnxfunction_function
  249. def _fromnxfunction_allargs(npfunc, /, *arys, **kwargs):
  250. """
  251. Wraps a NumPy function that can be called with multiple array arguments.
  252. All args are converted to arrays even if they are not so already.
  253. This makes it possible to process scalars as 1-D arrays.
  254. Only keyword arguments are passed through verbatim for the data and mask calls.
  255. Arrays arguments are processed independently and the results are returned in a list.
  256. If only one arg is present, the return value is just the processed array instead of
  257. a list.
  258. """
  259. out = tuple(
  260. masked_array(
  261. data=npfunc(np.asarray(a), **kwargs),
  262. mask=npfunc(getmaskarray(a), **kwargs),
  263. )
  264. for a in arys
  265. )
  266. return out[0] if len(out) == 1 else out
  267. atleast_1d = _fromnxfunction_allargs(np.atleast_1d)
  268. atleast_2d = _fromnxfunction_allargs(np.atleast_2d)
  269. atleast_3d = _fromnxfunction_allargs(np.atleast_3d)
  270. vstack = row_stack = _fromnxfunction_seq(np.vstack)
  271. hstack = _fromnxfunction_seq(np.hstack)
  272. column_stack = _fromnxfunction_seq(np.column_stack)
  273. dstack = _fromnxfunction_seq(np.dstack)
  274. stack = _fromnxfunction_seq(np.stack)
  275. hsplit = _fromnxfunction_single(np.hsplit)
  276. diagflat = _fromnxfunction_single(np.diagflat)
  277. #####--------------------------------------------------------------------------
  278. #----
  279. #####--------------------------------------------------------------------------
  280. def flatten_inplace(seq):
  281. """Flatten a sequence in place."""
  282. k = 0
  283. while (k != len(seq)):
  284. while hasattr(seq[k], '__iter__'):
  285. seq[k:(k + 1)] = seq[k]
  286. k += 1
  287. return seq
  288. def apply_along_axis(func1d, axis, arr, *args, **kwargs):
  289. """
  290. (This docstring should be overwritten)
  291. """
  292. arr = array(arr, copy=False, subok=True)
  293. nd = arr.ndim
  294. axis = normalize_axis_index(axis, nd)
  295. ind = [0] * (nd - 1)
  296. i = np.zeros(nd, 'O')
  297. indlist = list(range(nd))
  298. indlist.remove(axis)
  299. i[axis] = slice(None, None)
  300. outshape = np.asarray(arr.shape).take(indlist)
  301. i.put(indlist, ind)
  302. res = func1d(arr[tuple(i.tolist())], *args, **kwargs)
  303. # if res is a number, then we have a smaller output array
  304. asscalar = np.isscalar(res)
  305. if not asscalar:
  306. try:
  307. len(res)
  308. except TypeError:
  309. asscalar = True
  310. # Note: we shouldn't set the dtype of the output from the first result
  311. # so we force the type to object, and build a list of dtypes. We'll
  312. # just take the largest, to avoid some downcasting
  313. dtypes = []
  314. if asscalar:
  315. dtypes.append(np.asarray(res).dtype)
  316. outarr = zeros(outshape, object)
  317. outarr[tuple(ind)] = res
  318. Ntot = np.prod(outshape)
  319. k = 1
  320. while k < Ntot:
  321. # increment the index
  322. ind[-1] += 1
  323. n = -1
  324. while (ind[n] >= outshape[n]) and (n > (1 - nd)):
  325. ind[n - 1] += 1
  326. ind[n] = 0
  327. n -= 1
  328. i.put(indlist, ind)
  329. res = func1d(arr[tuple(i.tolist())], *args, **kwargs)
  330. outarr[tuple(ind)] = res
  331. dtypes.append(asarray(res).dtype)
  332. k += 1
  333. else:
  334. res = array(res, copy=False, subok=True)
  335. j = i.copy()
  336. j[axis] = ([slice(None, None)] * res.ndim)
  337. j.put(indlist, ind)
  338. Ntot = np.prod(outshape)
  339. holdshape = outshape
  340. outshape = list(arr.shape)
  341. outshape[axis] = res.shape
  342. dtypes.append(asarray(res).dtype)
  343. outshape = flatten_inplace(outshape)
  344. outarr = zeros(outshape, object)
  345. outarr[tuple(flatten_inplace(j.tolist()))] = res
  346. k = 1
  347. while k < Ntot:
  348. # increment the index
  349. ind[-1] += 1
  350. n = -1
  351. while (ind[n] >= holdshape[n]) and (n > (1 - nd)):
  352. ind[n - 1] += 1
  353. ind[n] = 0
  354. n -= 1
  355. i.put(indlist, ind)
  356. j.put(indlist, ind)
  357. res = func1d(arr[tuple(i.tolist())], *args, **kwargs)
  358. outarr[tuple(flatten_inplace(j.tolist()))] = res
  359. dtypes.append(asarray(res).dtype)
  360. k += 1
  361. max_dtypes = np.dtype(np.asarray(dtypes).max())
  362. if not hasattr(arr, '_mask'):
  363. result = np.asarray(outarr, dtype=max_dtypes)
  364. else:
  365. result = asarray(outarr, dtype=max_dtypes)
  366. result.fill_value = ma.default_fill_value(result)
  367. return result
  368. apply_along_axis.__doc__ = np.apply_along_axis.__doc__
  369. def apply_over_axes(func, a, axes):
  370. """
  371. (This docstring will be overwritten)
  372. """
  373. val = asarray(a)
  374. N = a.ndim
  375. if array(axes).ndim == 0:
  376. axes = (axes,)
  377. for axis in axes:
  378. if axis < 0:
  379. axis = N + axis
  380. args = (val, axis)
  381. res = func(*args)
  382. if res.ndim == val.ndim:
  383. val = res
  384. else:
  385. res = ma.expand_dims(res, axis)
  386. if res.ndim == val.ndim:
  387. val = res
  388. else:
  389. raise ValueError("function is not returning "
  390. "an array of the correct shape")
  391. return val
  392. if apply_over_axes.__doc__ is not None:
  393. apply_over_axes.__doc__ = np.apply_over_axes.__doc__[
  394. :np.apply_over_axes.__doc__.find('Notes')].rstrip() + \
  395. """
  396. Examples
  397. --------
  398. >>> import numpy as np
  399. >>> a = np.ma.arange(24).reshape(2,3,4)
  400. >>> a[:,0,1] = np.ma.masked
  401. >>> a[:,1,:] = np.ma.masked
  402. >>> a
  403. masked_array(
  404. data=[[[0, --, 2, 3],
  405. [--, --, --, --],
  406. [8, 9, 10, 11]],
  407. [[12, --, 14, 15],
  408. [--, --, --, --],
  409. [20, 21, 22, 23]]],
  410. mask=[[[False, True, False, False],
  411. [ True, True, True, True],
  412. [False, False, False, False]],
  413. [[False, True, False, False],
  414. [ True, True, True, True],
  415. [False, False, False, False]]],
  416. fill_value=999999)
  417. >>> np.ma.apply_over_axes(np.ma.sum, a, [0,2])
  418. masked_array(
  419. data=[[[46],
  420. [--],
  421. [124]]],
  422. mask=[[[False],
  423. [ True],
  424. [False]]],
  425. fill_value=999999)
  426. Tuple axis arguments to ufuncs are equivalent:
  427. >>> np.ma.sum(a, axis=(0,2)).reshape((1,-1,1))
  428. masked_array(
  429. data=[[[46],
  430. [--],
  431. [124]]],
  432. mask=[[[False],
  433. [ True],
  434. [False]]],
  435. fill_value=999999)
  436. """
  437. def average(a, axis=None, weights=None, returned=False, *,
  438. keepdims=np._NoValue):
  439. """
  440. Return the weighted average of array over the given axis.
  441. Parameters
  442. ----------
  443. a : array_like
  444. Data to be averaged.
  445. Masked entries are not taken into account in the computation.
  446. axis : None or int or tuple of ints, optional
  447. Axis or axes along which to average `a`. The default,
  448. `axis=None`, will average over all of the elements of the input array.
  449. If axis is a tuple of ints, averaging is performed on all of the axes
  450. specified in the tuple instead of a single axis or all the axes as
  451. before.
  452. weights : array_like, optional
  453. An array of weights associated with the values in `a`. Each value in
  454. `a` contributes to the average according to its associated weight.
  455. The array of weights must be the same shape as `a` if no axis is
  456. specified, otherwise the weights must have dimensions and shape
  457. consistent with `a` along the specified axis.
  458. If `weights=None`, then all data in `a` are assumed to have a
  459. weight equal to one.
  460. The calculation is::
  461. avg = sum(a * weights) / sum(weights)
  462. where the sum is over all included elements.
  463. The only constraint on the values of `weights` is that `sum(weights)`
  464. must not be 0.
  465. returned : bool, optional
  466. Flag indicating whether a tuple ``(result, sum of weights)``
  467. should be returned as output (True), or just the result (False).
  468. Default is False.
  469. keepdims : bool, optional
  470. If this is set to True, the axes which are reduced are left
  471. in the result as dimensions with size one. With this option,
  472. the result will broadcast correctly against the original `a`.
  473. *Note:* `keepdims` will not work with instances of `numpy.matrix`
  474. or other classes whose methods do not support `keepdims`.
  475. .. versionadded:: 1.23.0
  476. Returns
  477. -------
  478. average, [sum_of_weights] : (tuple of) scalar or MaskedArray
  479. The average along the specified axis. When returned is `True`,
  480. return a tuple with the average as the first element and the sum
  481. of the weights as the second element. The return type is `np.float64`
  482. if `a` is of integer type and floats smaller than `float64`, or the
  483. input data-type, otherwise. If returned, `sum_of_weights` is always
  484. `float64`.
  485. Raises
  486. ------
  487. ZeroDivisionError
  488. When all weights along axis are zero. See `numpy.ma.average` for a
  489. version robust to this type of error.
  490. TypeError
  491. When `weights` does not have the same shape as `a`, and `axis=None`.
  492. ValueError
  493. When `weights` does not have dimensions and shape consistent with `a`
  494. along specified `axis`.
  495. Examples
  496. --------
  497. >>> import numpy as np
  498. >>> a = np.ma.array([1., 2., 3., 4.], mask=[False, False, True, True])
  499. >>> np.ma.average(a, weights=[3, 1, 0, 0])
  500. 1.25
  501. >>> x = np.ma.arange(6.).reshape(3, 2)
  502. >>> x
  503. masked_array(
  504. data=[[0., 1.],
  505. [2., 3.],
  506. [4., 5.]],
  507. mask=False,
  508. fill_value=1e+20)
  509. >>> data = np.arange(8).reshape((2, 2, 2))
  510. >>> data
  511. array([[[0, 1],
  512. [2, 3]],
  513. [[4, 5],
  514. [6, 7]]])
  515. >>> np.ma.average(data, axis=(0, 1), weights=[[1./4, 3./4], [1., 1./2]])
  516. masked_array(data=[3.4, 4.4],
  517. mask=[False, False],
  518. fill_value=1e+20)
  519. >>> np.ma.average(data, axis=0, weights=[[1./4, 3./4], [1., 1./2]])
  520. Traceback (most recent call last):
  521. ...
  522. ValueError: Shape of weights must be consistent
  523. with shape of a along specified axis.
  524. >>> avg, sumweights = np.ma.average(x, axis=0, weights=[1, 2, 3],
  525. ... returned=True)
  526. >>> avg
  527. masked_array(data=[2.6666666666666665, 3.6666666666666665],
  528. mask=[False, False],
  529. fill_value=1e+20)
  530. With ``keepdims=True``, the following result has shape (3, 1).
  531. >>> np.ma.average(x, axis=1, keepdims=True)
  532. masked_array(
  533. data=[[0.5],
  534. [2.5],
  535. [4.5]],
  536. mask=False,
  537. fill_value=1e+20)
  538. """
  539. a = asarray(a)
  540. m = getmask(a)
  541. if axis is not None:
  542. axis = normalize_axis_tuple(axis, a.ndim, argname="axis")
  543. if keepdims is np._NoValue:
  544. # Don't pass on the keepdims argument if one wasn't given.
  545. keepdims_kw = {}
  546. else:
  547. keepdims_kw = {'keepdims': keepdims}
  548. if weights is None:
  549. avg = a.mean(axis, **keepdims_kw)
  550. scl = avg.dtype.type(a.count(axis))
  551. else:
  552. wgt = asarray(weights)
  553. if issubclass(a.dtype.type, (np.integer, np.bool)):
  554. result_dtype = np.result_type(a.dtype, wgt.dtype, 'f8')
  555. else:
  556. result_dtype = np.result_type(a.dtype, wgt.dtype)
  557. # Sanity checks
  558. if a.shape != wgt.shape:
  559. if axis is None:
  560. raise TypeError(
  561. "Axis must be specified when shapes of a and weights "
  562. "differ.")
  563. if wgt.shape != tuple(a.shape[ax] for ax in axis):
  564. raise ValueError(
  565. "Shape of weights must be consistent with "
  566. "shape of a along specified axis.")
  567. # setup wgt to broadcast along axis
  568. wgt = wgt.transpose(np.argsort(axis))
  569. wgt = wgt.reshape(tuple((s if ax in axis else 1)
  570. for ax, s in enumerate(a.shape)))
  571. if m is not nomask:
  572. wgt = wgt * (~a.mask)
  573. wgt.mask |= a.mask
  574. scl = wgt.sum(axis=axis, dtype=result_dtype, **keepdims_kw)
  575. avg = np.multiply(a, wgt,
  576. dtype=result_dtype).sum(axis, **keepdims_kw) / scl
  577. if returned:
  578. if scl.shape != avg.shape:
  579. scl = np.broadcast_to(scl, avg.shape).copy()
  580. return avg, scl
  581. else:
  582. return avg
  583. def median(a, axis=None, out=None, overwrite_input=False, keepdims=False):
  584. """
  585. Compute the median along the specified axis.
  586. Returns the median of the array elements.
  587. Parameters
  588. ----------
  589. a : array_like
  590. Input array or object that can be converted to an array.
  591. axis : int, optional
  592. Axis along which the medians are computed. The default (None) is
  593. to compute the median along a flattened version of the array.
  594. out : ndarray, optional
  595. Alternative output array in which to place the result. It must
  596. have the same shape and buffer length as the expected output
  597. but the type will be cast if necessary.
  598. overwrite_input : bool, optional
  599. If True, then allow use of memory of input array (a) for
  600. calculations. The input array will be modified by the call to
  601. median. This will save memory when you do not need to preserve
  602. the contents of the input array. Treat the input as undefined,
  603. but it will probably be fully or partially sorted. Default is
  604. False. Note that, if `overwrite_input` is True, and the input
  605. is not already an `ndarray`, an error will be raised.
  606. keepdims : bool, optional
  607. If this is set to True, the axes which are reduced are left
  608. in the result as dimensions with size one. With this option,
  609. the result will broadcast correctly against the input array.
  610. Returns
  611. -------
  612. median : ndarray
  613. A new array holding the result is returned unless out is
  614. specified, in which case a reference to out is returned.
  615. Return data-type is `float64` for integers and floats smaller than
  616. `float64`, or the input data-type, otherwise.
  617. See Also
  618. --------
  619. mean
  620. Notes
  621. -----
  622. Given a vector ``V`` with ``N`` non masked values, the median of ``V``
  623. is the middle value of a sorted copy of ``V`` (``Vs``) - i.e.
  624. ``Vs[(N-1)/2]``, when ``N`` is odd, or ``{Vs[N/2 - 1] + Vs[N/2]}/2``
  625. when ``N`` is even.
  626. Examples
  627. --------
  628. >>> import numpy as np
  629. >>> x = np.ma.array(np.arange(8), mask=[0]*4 + [1]*4)
  630. >>> np.ma.median(x)
  631. 1.5
  632. >>> x = np.ma.array(np.arange(10).reshape(2, 5), mask=[0]*6 + [1]*4)
  633. >>> np.ma.median(x)
  634. 2.5
  635. >>> np.ma.median(x, axis=-1, overwrite_input=True)
  636. masked_array(data=[2.0, 5.0],
  637. mask=[False, False],
  638. fill_value=1e+20)
  639. """
  640. if not hasattr(a, 'mask'):
  641. m = np.median(getdata(a, subok=True), axis=axis,
  642. out=out, overwrite_input=overwrite_input,
  643. keepdims=keepdims)
  644. if isinstance(m, np.ndarray) and 1 <= m.ndim:
  645. return masked_array(m, copy=False)
  646. else:
  647. return m
  648. return _ureduce(a, func=_median, keepdims=keepdims, axis=axis, out=out,
  649. overwrite_input=overwrite_input)
  650. def _median(a, axis=None, out=None, overwrite_input=False):
  651. # when an unmasked NaN is present return it, so we need to sort the NaN
  652. # values behind the mask
  653. if np.issubdtype(a.dtype, np.inexact):
  654. fill_value = np.inf
  655. else:
  656. fill_value = None
  657. if overwrite_input:
  658. if axis is None:
  659. asorted = a.ravel()
  660. asorted.sort(fill_value=fill_value)
  661. else:
  662. a.sort(axis=axis, fill_value=fill_value)
  663. asorted = a
  664. else:
  665. asorted = sort(a, axis=axis, fill_value=fill_value)
  666. if axis is None:
  667. axis = 0
  668. else:
  669. axis = normalize_axis_index(axis, asorted.ndim)
  670. if asorted.shape[axis] == 0:
  671. # for empty axis integer indices fail so use slicing to get same result
  672. # as median (which is mean of empty slice = nan)
  673. indexer = [slice(None)] * asorted.ndim
  674. indexer[axis] = slice(0, 0)
  675. indexer = tuple(indexer)
  676. return np.ma.mean(asorted[indexer], axis=axis, out=out)
  677. if asorted.ndim == 1:
  678. idx, odd = divmod(count(asorted), 2)
  679. mid = asorted[idx + odd - 1:idx + 1]
  680. if np.issubdtype(asorted.dtype, np.inexact) and asorted.size > 0:
  681. # avoid inf / x = masked
  682. s = mid.sum(out=out)
  683. if not odd:
  684. s = np.true_divide(s, 2., casting='safe', out=out)
  685. s = np.lib._utils_impl._median_nancheck(asorted, s, axis)
  686. else:
  687. s = mid.mean(out=out)
  688. # if result is masked either the input contained enough
  689. # minimum_fill_value so that it would be the median or all values
  690. # masked
  691. if np.ma.is_masked(s) and not np.all(asorted.mask):
  692. return np.ma.minimum_fill_value(asorted)
  693. return s
  694. counts = count(asorted, axis=axis, keepdims=True)
  695. h = counts // 2
  696. # duplicate high if odd number of elements so mean does nothing
  697. odd = counts % 2 == 1
  698. l = np.where(odd, h, h - 1)
  699. lh = np.concatenate([l, h], axis=axis)
  700. # get low and high median
  701. low_high = np.take_along_axis(asorted, lh, axis=axis)
  702. def replace_masked(s):
  703. # Replace masked entries with minimum_full_value unless it all values
  704. # are masked. This is required as the sort order of values equal or
  705. # larger than the fill value is undefined and a valid value placed
  706. # elsewhere, e.g. [4, --, inf].
  707. if np.ma.is_masked(s):
  708. rep = (~np.all(asorted.mask, axis=axis, keepdims=True)) & s.mask
  709. s.data[rep] = np.ma.minimum_fill_value(asorted)
  710. s.mask[rep] = False
  711. replace_masked(low_high)
  712. if np.issubdtype(asorted.dtype, np.inexact):
  713. # avoid inf / x = masked
  714. s = np.ma.sum(low_high, axis=axis, out=out)
  715. np.true_divide(s.data, 2., casting='unsafe', out=s.data)
  716. s = np.lib._utils_impl._median_nancheck(asorted, s, axis)
  717. else:
  718. s = np.ma.mean(low_high, axis=axis, out=out)
  719. return s
  720. def compress_nd(x, axis=None):
  721. """Suppress slices from multiple dimensions which contain masked values.
  722. Parameters
  723. ----------
  724. x : array_like, MaskedArray
  725. The array to operate on. If not a MaskedArray instance (or if no array
  726. elements are masked), `x` is interpreted as a MaskedArray with `mask`
  727. set to `nomask`.
  728. axis : tuple of ints or int, optional
  729. Which dimensions to suppress slices from can be configured with this
  730. parameter.
  731. - If axis is a tuple of ints, those are the axes to suppress slices from.
  732. - If axis is an int, then that is the only axis to suppress slices from.
  733. - If axis is None, all axis are selected.
  734. Returns
  735. -------
  736. compress_array : ndarray
  737. The compressed array.
  738. Examples
  739. --------
  740. >>> import numpy as np
  741. >>> arr = [[1, 2], [3, 4]]
  742. >>> mask = [[0, 1], [0, 0]]
  743. >>> x = np.ma.array(arr, mask=mask)
  744. >>> np.ma.compress_nd(x, axis=0)
  745. array([[3, 4]])
  746. >>> np.ma.compress_nd(x, axis=1)
  747. array([[1],
  748. [3]])
  749. >>> np.ma.compress_nd(x)
  750. array([[3]])
  751. """
  752. x = asarray(x)
  753. m = getmask(x)
  754. # Set axis to tuple of ints
  755. if axis is None:
  756. axis = tuple(range(x.ndim))
  757. else:
  758. axis = normalize_axis_tuple(axis, x.ndim)
  759. # Nothing is masked: return x
  760. if m is nomask or not m.any():
  761. return x._data
  762. # All is masked: return empty
  763. if m.all():
  764. return nxarray([])
  765. # Filter elements through boolean indexing
  766. data = x._data
  767. for ax in axis:
  768. axes = tuple(list(range(ax)) + list(range(ax + 1, x.ndim)))
  769. data = data[(slice(None),) * ax + (~m.any(axis=axes),)]
  770. return data
  771. def compress_rowcols(x, axis=None):
  772. """
  773. Suppress the rows and/or columns of a 2-D array that contain
  774. masked values.
  775. The suppression behavior is selected with the `axis` parameter.
  776. - If axis is None, both rows and columns are suppressed.
  777. - If axis is 0, only rows are suppressed.
  778. - If axis is 1 or -1, only columns are suppressed.
  779. Parameters
  780. ----------
  781. x : array_like, MaskedArray
  782. The array to operate on. If not a MaskedArray instance (or if no array
  783. elements are masked), `x` is interpreted as a MaskedArray with
  784. `mask` set to `nomask`. Must be a 2D array.
  785. axis : int, optional
  786. Axis along which to perform the operation. Default is None.
  787. Returns
  788. -------
  789. compressed_array : ndarray
  790. The compressed array.
  791. Examples
  792. --------
  793. >>> import numpy as np
  794. >>> x = np.ma.array(np.arange(9).reshape(3, 3), mask=[[1, 0, 0],
  795. ... [1, 0, 0],
  796. ... [0, 0, 0]])
  797. >>> x
  798. masked_array(
  799. data=[[--, 1, 2],
  800. [--, 4, 5],
  801. [6, 7, 8]],
  802. mask=[[ True, False, False],
  803. [ True, False, False],
  804. [False, False, False]],
  805. fill_value=999999)
  806. >>> np.ma.compress_rowcols(x)
  807. array([[7, 8]])
  808. >>> np.ma.compress_rowcols(x, 0)
  809. array([[6, 7, 8]])
  810. >>> np.ma.compress_rowcols(x, 1)
  811. array([[1, 2],
  812. [4, 5],
  813. [7, 8]])
  814. """
  815. if asarray(x).ndim != 2:
  816. raise NotImplementedError("compress_rowcols works for 2D arrays only.")
  817. return compress_nd(x, axis=axis)
  818. def compress_rows(a):
  819. """
  820. Suppress whole rows of a 2-D array that contain masked values.
  821. This is equivalent to ``np.ma.compress_rowcols(a, 0)``, see
  822. `compress_rowcols` for details.
  823. Parameters
  824. ----------
  825. x : array_like, MaskedArray
  826. The array to operate on. If not a MaskedArray instance (or if no array
  827. elements are masked), `x` is interpreted as a MaskedArray with
  828. `mask` set to `nomask`. Must be a 2D array.
  829. Returns
  830. -------
  831. compressed_array : ndarray
  832. The compressed array.
  833. See Also
  834. --------
  835. compress_rowcols
  836. Examples
  837. --------
  838. >>> import numpy as np
  839. >>> a = np.ma.array(np.arange(9).reshape(3, 3), mask=[[1, 0, 0],
  840. ... [1, 0, 0],
  841. ... [0, 0, 0]])
  842. >>> np.ma.compress_rows(a)
  843. array([[6, 7, 8]])
  844. """
  845. a = asarray(a)
  846. if a.ndim != 2:
  847. raise NotImplementedError("compress_rows works for 2D arrays only.")
  848. return compress_rowcols(a, 0)
  849. def compress_cols(a):
  850. """
  851. Suppress whole columns of a 2-D array that contain masked values.
  852. This is equivalent to ``np.ma.compress_rowcols(a, 1)``, see
  853. `compress_rowcols` for details.
  854. Parameters
  855. ----------
  856. x : array_like, MaskedArray
  857. The array to operate on. If not a MaskedArray instance (or if no array
  858. elements are masked), `x` is interpreted as a MaskedArray with
  859. `mask` set to `nomask`. Must be a 2D array.
  860. Returns
  861. -------
  862. compressed_array : ndarray
  863. The compressed array.
  864. See Also
  865. --------
  866. compress_rowcols
  867. Examples
  868. --------
  869. >>> import numpy as np
  870. >>> a = np.ma.array(np.arange(9).reshape(3, 3), mask=[[1, 0, 0],
  871. ... [1, 0, 0],
  872. ... [0, 0, 0]])
  873. >>> np.ma.compress_cols(a)
  874. array([[1, 2],
  875. [4, 5],
  876. [7, 8]])
  877. """
  878. a = asarray(a)
  879. if a.ndim != 2:
  880. raise NotImplementedError("compress_cols works for 2D arrays only.")
  881. return compress_rowcols(a, 1)
  882. def mask_rowcols(a, axis=None):
  883. """
  884. Mask rows and/or columns of a 2D array that contain masked values.
  885. Mask whole rows and/or columns of a 2D array that contain
  886. masked values. The masking behavior is selected using the
  887. `axis` parameter.
  888. - If `axis` is None, rows *and* columns are masked.
  889. - If `axis` is 0, only rows are masked.
  890. - If `axis` is 1 or -1, only columns are masked.
  891. Parameters
  892. ----------
  893. a : array_like, MaskedArray
  894. The array to mask. If not a MaskedArray instance (or if no array
  895. elements are masked), the result is a MaskedArray with `mask` set
  896. to `nomask` (False). Must be a 2D array.
  897. axis : int, optional
  898. Axis along which to perform the operation. If None, applies to a
  899. flattened version of the array.
  900. Returns
  901. -------
  902. a : MaskedArray
  903. A modified version of the input array, masked depending on the value
  904. of the `axis` parameter.
  905. Raises
  906. ------
  907. NotImplementedError
  908. If input array `a` is not 2D.
  909. See Also
  910. --------
  911. mask_rows : Mask rows of a 2D array that contain masked values.
  912. mask_cols : Mask cols of a 2D array that contain masked values.
  913. masked_where : Mask where a condition is met.
  914. Notes
  915. -----
  916. The input array's mask is modified by this function.
  917. Examples
  918. --------
  919. >>> import numpy as np
  920. >>> a = np.zeros((3, 3), dtype=int)
  921. >>> a[1, 1] = 1
  922. >>> a
  923. array([[0, 0, 0],
  924. [0, 1, 0],
  925. [0, 0, 0]])
  926. >>> a = np.ma.masked_equal(a, 1)
  927. >>> a
  928. masked_array(
  929. data=[[0, 0, 0],
  930. [0, --, 0],
  931. [0, 0, 0]],
  932. mask=[[False, False, False],
  933. [False, True, False],
  934. [False, False, False]],
  935. fill_value=1)
  936. >>> np.ma.mask_rowcols(a)
  937. masked_array(
  938. data=[[0, --, 0],
  939. [--, --, --],
  940. [0, --, 0]],
  941. mask=[[False, True, False],
  942. [ True, True, True],
  943. [False, True, False]],
  944. fill_value=1)
  945. """
  946. a = array(a, subok=False)
  947. if a.ndim != 2:
  948. raise NotImplementedError("mask_rowcols works for 2D arrays only.")
  949. m = getmask(a)
  950. # Nothing is masked: return a
  951. if m is nomask or not m.any():
  952. return a
  953. maskedval = m.nonzero()
  954. a._mask = a._mask.copy()
  955. if not axis:
  956. a[np.unique(maskedval[0])] = masked
  957. if axis in [None, 1, -1]:
  958. a[:, np.unique(maskedval[1])] = masked
  959. return a
  960. def mask_rows(a, axis=np._NoValue):
  961. """
  962. Mask rows of a 2D array that contain masked values.
  963. This function is a shortcut to ``mask_rowcols`` with `axis` equal to 0.
  964. See Also
  965. --------
  966. mask_rowcols : Mask rows and/or columns of a 2D array.
  967. masked_where : Mask where a condition is met.
  968. Examples
  969. --------
  970. >>> import numpy as np
  971. >>> a = np.zeros((3, 3), dtype=int)
  972. >>> a[1, 1] = 1
  973. >>> a
  974. array([[0, 0, 0],
  975. [0, 1, 0],
  976. [0, 0, 0]])
  977. >>> a = np.ma.masked_equal(a, 1)
  978. >>> a
  979. masked_array(
  980. data=[[0, 0, 0],
  981. [0, --, 0],
  982. [0, 0, 0]],
  983. mask=[[False, False, False],
  984. [False, True, False],
  985. [False, False, False]],
  986. fill_value=1)
  987. >>> np.ma.mask_rows(a)
  988. masked_array(
  989. data=[[0, 0, 0],
  990. [--, --, --],
  991. [0, 0, 0]],
  992. mask=[[False, False, False],
  993. [ True, True, True],
  994. [False, False, False]],
  995. fill_value=1)
  996. """
  997. if axis is not np._NoValue:
  998. # remove the axis argument when this deprecation expires
  999. # NumPy 1.18.0, 2019-11-28
  1000. warnings.warn(
  1001. "The axis argument has always been ignored, in future passing it "
  1002. "will raise TypeError", DeprecationWarning, stacklevel=2)
  1003. return mask_rowcols(a, 0)
  1004. def mask_cols(a, axis=np._NoValue):
  1005. """
  1006. Mask columns of a 2D array that contain masked values.
  1007. This function is a shortcut to ``mask_rowcols`` with `axis` equal to 1.
  1008. See Also
  1009. --------
  1010. mask_rowcols : Mask rows and/or columns of a 2D array.
  1011. masked_where : Mask where a condition is met.
  1012. Examples
  1013. --------
  1014. >>> import numpy as np
  1015. >>> a = np.zeros((3, 3), dtype=int)
  1016. >>> a[1, 1] = 1
  1017. >>> a
  1018. array([[0, 0, 0],
  1019. [0, 1, 0],
  1020. [0, 0, 0]])
  1021. >>> a = np.ma.masked_equal(a, 1)
  1022. >>> a
  1023. masked_array(
  1024. data=[[0, 0, 0],
  1025. [0, --, 0],
  1026. [0, 0, 0]],
  1027. mask=[[False, False, False],
  1028. [False, True, False],
  1029. [False, False, False]],
  1030. fill_value=1)
  1031. >>> np.ma.mask_cols(a)
  1032. masked_array(
  1033. data=[[0, --, 0],
  1034. [0, --, 0],
  1035. [0, --, 0]],
  1036. mask=[[False, True, False],
  1037. [False, True, False],
  1038. [False, True, False]],
  1039. fill_value=1)
  1040. """
  1041. if axis is not np._NoValue:
  1042. # remove the axis argument when this deprecation expires
  1043. # NumPy 1.18.0, 2019-11-28
  1044. warnings.warn(
  1045. "The axis argument has always been ignored, in future passing it "
  1046. "will raise TypeError", DeprecationWarning, stacklevel=2)
  1047. return mask_rowcols(a, 1)
  1048. #####--------------------------------------------------------------------------
  1049. #---- --- arraysetops ---
  1050. #####--------------------------------------------------------------------------
  1051. def ediff1d(arr, to_end=None, to_begin=None):
  1052. """
  1053. Compute the differences between consecutive elements of an array.
  1054. This function is the equivalent of `numpy.ediff1d` that takes masked
  1055. values into account, see `numpy.ediff1d` for details.
  1056. See Also
  1057. --------
  1058. numpy.ediff1d : Equivalent function for ndarrays.
  1059. Examples
  1060. --------
  1061. >>> import numpy as np
  1062. >>> arr = np.ma.array([1, 2, 4, 7, 0])
  1063. >>> np.ma.ediff1d(arr)
  1064. masked_array(data=[ 1, 2, 3, -7],
  1065. mask=False,
  1066. fill_value=999999)
  1067. """
  1068. arr = ma.asanyarray(arr).flat
  1069. ed = arr[1:] - arr[:-1]
  1070. arrays = [ed]
  1071. #
  1072. if to_begin is not None:
  1073. arrays.insert(0, to_begin)
  1074. if to_end is not None:
  1075. arrays.append(to_end)
  1076. #
  1077. if len(arrays) != 1:
  1078. # We'll save ourselves a copy of a potentially large array in the common
  1079. # case where neither to_begin or to_end was given.
  1080. ed = hstack(arrays)
  1081. #
  1082. return ed
  1083. def unique(ar1, return_index=False, return_inverse=False):
  1084. """
  1085. Finds the unique elements of an array.
  1086. Masked values are considered the same element (masked). The output array
  1087. is always a masked array. See `numpy.unique` for more details.
  1088. See Also
  1089. --------
  1090. numpy.unique : Equivalent function for ndarrays.
  1091. Examples
  1092. --------
  1093. >>> import numpy as np
  1094. >>> a = [1, 2, 1000, 2, 3]
  1095. >>> mask = [0, 0, 1, 0, 0]
  1096. >>> masked_a = np.ma.masked_array(a, mask)
  1097. >>> masked_a
  1098. masked_array(data=[1, 2, --, 2, 3],
  1099. mask=[False, False, True, False, False],
  1100. fill_value=999999)
  1101. >>> np.ma.unique(masked_a)
  1102. masked_array(data=[1, 2, 3, --],
  1103. mask=[False, False, False, True],
  1104. fill_value=999999)
  1105. >>> np.ma.unique(masked_a, return_index=True)
  1106. (masked_array(data=[1, 2, 3, --],
  1107. mask=[False, False, False, True],
  1108. fill_value=999999), array([0, 1, 4, 2]))
  1109. >>> np.ma.unique(masked_a, return_inverse=True)
  1110. (masked_array(data=[1, 2, 3, --],
  1111. mask=[False, False, False, True],
  1112. fill_value=999999), array([0, 1, 3, 1, 2]))
  1113. >>> np.ma.unique(masked_a, return_index=True, return_inverse=True)
  1114. (masked_array(data=[1, 2, 3, --],
  1115. mask=[False, False, False, True],
  1116. fill_value=999999), array([0, 1, 4, 2]), array([0, 1, 3, 1, 2]))
  1117. """
  1118. output = np.unique(ar1,
  1119. return_index=return_index,
  1120. return_inverse=return_inverse)
  1121. if isinstance(output, tuple):
  1122. output = list(output)
  1123. output[0] = output[0].view(MaskedArray)
  1124. output = tuple(output)
  1125. else:
  1126. output = output.view(MaskedArray)
  1127. return output
  1128. def intersect1d(ar1, ar2, assume_unique=False):
  1129. """
  1130. Returns the unique elements common to both arrays.
  1131. Masked values are considered equal one to the other.
  1132. The output is always a masked array.
  1133. See `numpy.intersect1d` for more details.
  1134. See Also
  1135. --------
  1136. numpy.intersect1d : Equivalent function for ndarrays.
  1137. Examples
  1138. --------
  1139. >>> import numpy as np
  1140. >>> x = np.ma.array([1, 3, 3, 3], mask=[0, 0, 0, 1])
  1141. >>> y = np.ma.array([3, 1, 1, 1], mask=[0, 0, 0, 1])
  1142. >>> np.ma.intersect1d(x, y)
  1143. masked_array(data=[1, 3, --],
  1144. mask=[False, False, True],
  1145. fill_value=999999)
  1146. """
  1147. if assume_unique:
  1148. aux = ma.concatenate((ar1, ar2))
  1149. else:
  1150. # Might be faster than unique( intersect1d( ar1, ar2 ) )?
  1151. aux = ma.concatenate((unique(ar1), unique(ar2)))
  1152. aux.sort()
  1153. return aux[:-1][aux[1:] == aux[:-1]]
  1154. def setxor1d(ar1, ar2, assume_unique=False):
  1155. """
  1156. Set exclusive-or of 1-D arrays with unique elements.
  1157. The output is always a masked array. See `numpy.setxor1d` for more details.
  1158. See Also
  1159. --------
  1160. numpy.setxor1d : Equivalent function for ndarrays.
  1161. Examples
  1162. --------
  1163. >>> import numpy as np
  1164. >>> ar1 = np.ma.array([1, 2, 3, 2, 4])
  1165. >>> ar2 = np.ma.array([2, 3, 5, 7, 5])
  1166. >>> np.ma.setxor1d(ar1, ar2)
  1167. masked_array(data=[1, 4, 5, 7],
  1168. mask=False,
  1169. fill_value=999999)
  1170. """
  1171. if not assume_unique:
  1172. ar1 = unique(ar1)
  1173. ar2 = unique(ar2)
  1174. aux = ma.concatenate((ar1, ar2), axis=None)
  1175. if aux.size == 0:
  1176. return aux
  1177. aux.sort()
  1178. auxf = aux.filled()
  1179. # flag = ediff1d( aux, to_end = 1, to_begin = 1 ) == 0
  1180. flag = ma.concatenate(([True], (auxf[1:] != auxf[:-1]), [True]))
  1181. # flag2 = ediff1d( flag ) == 0
  1182. flag2 = (flag[1:] == flag[:-1])
  1183. return aux[flag2]
  1184. def in1d(ar1, ar2, assume_unique=False, invert=False):
  1185. """
  1186. Test whether each element of an array is also present in a second
  1187. array.
  1188. The output is always a masked array.
  1189. We recommend using :func:`isin` instead of `in1d` for new code.
  1190. See Also
  1191. --------
  1192. isin : Version of this function that preserves the shape of ar1.
  1193. Examples
  1194. --------
  1195. >>> import numpy as np
  1196. >>> ar1 = np.ma.array([0, 1, 2, 5, 0])
  1197. >>> ar2 = [0, 2]
  1198. >>> np.ma.in1d(ar1, ar2)
  1199. masked_array(data=[ True, False, True, False, True],
  1200. mask=False,
  1201. fill_value=True)
  1202. """
  1203. if not assume_unique:
  1204. ar1, rev_idx = unique(ar1, return_inverse=True)
  1205. ar2 = unique(ar2)
  1206. ar = ma.concatenate((ar1, ar2))
  1207. # We need this to be a stable sort, so always use 'mergesort'
  1208. # here. The values from the first array should always come before
  1209. # the values from the second array.
  1210. order = ar.argsort(kind='mergesort')
  1211. sar = ar[order]
  1212. if invert:
  1213. bool_ar = (sar[1:] != sar[:-1])
  1214. else:
  1215. bool_ar = (sar[1:] == sar[:-1])
  1216. flag = ma.concatenate((bool_ar, [invert]))
  1217. indx = order.argsort(kind='mergesort')[:len(ar1)]
  1218. if assume_unique:
  1219. return flag[indx]
  1220. else:
  1221. return flag[indx][rev_idx]
  1222. def isin(element, test_elements, assume_unique=False, invert=False):
  1223. """
  1224. Calculates `element in test_elements`, broadcasting over
  1225. `element` only.
  1226. The output is always a masked array of the same shape as `element`.
  1227. See `numpy.isin` for more details.
  1228. See Also
  1229. --------
  1230. in1d : Flattened version of this function.
  1231. numpy.isin : Equivalent function for ndarrays.
  1232. Examples
  1233. --------
  1234. >>> import numpy as np
  1235. >>> element = np.ma.array([1, 2, 3, 4, 5, 6])
  1236. >>> test_elements = [0, 2]
  1237. >>> np.ma.isin(element, test_elements)
  1238. masked_array(data=[False, True, False, False, False, False],
  1239. mask=False,
  1240. fill_value=True)
  1241. """
  1242. element = ma.asarray(element)
  1243. return in1d(element, test_elements, assume_unique=assume_unique,
  1244. invert=invert).reshape(element.shape)
  1245. def union1d(ar1, ar2):
  1246. """
  1247. Union of two arrays.
  1248. The output is always a masked array. See `numpy.union1d` for more details.
  1249. See Also
  1250. --------
  1251. numpy.union1d : Equivalent function for ndarrays.
  1252. Examples
  1253. --------
  1254. >>> import numpy as np
  1255. >>> ar1 = np.ma.array([1, 2, 3, 4])
  1256. >>> ar2 = np.ma.array([3, 4, 5, 6])
  1257. >>> np.ma.union1d(ar1, ar2)
  1258. masked_array(data=[1, 2, 3, 4, 5, 6],
  1259. mask=False,
  1260. fill_value=999999)
  1261. """
  1262. return unique(ma.concatenate((ar1, ar2), axis=None))
  1263. def setdiff1d(ar1, ar2, assume_unique=False):
  1264. """
  1265. Set difference of 1D arrays with unique elements.
  1266. The output is always a masked array. See `numpy.setdiff1d` for more
  1267. details.
  1268. See Also
  1269. --------
  1270. numpy.setdiff1d : Equivalent function for ndarrays.
  1271. Examples
  1272. --------
  1273. >>> import numpy as np
  1274. >>> x = np.ma.array([1, 2, 3, 4], mask=[0, 1, 0, 1])
  1275. >>> np.ma.setdiff1d(x, [1, 2])
  1276. masked_array(data=[3, --],
  1277. mask=[False, True],
  1278. fill_value=999999)
  1279. """
  1280. if assume_unique:
  1281. ar1 = ma.asarray(ar1).ravel()
  1282. else:
  1283. ar1 = unique(ar1)
  1284. ar2 = unique(ar2)
  1285. return ar1[in1d(ar1, ar2, assume_unique=True, invert=True)]
  1286. ###############################################################################
  1287. # Covariance #
  1288. ###############################################################################
  1289. def _covhelper(x, y=None, rowvar=True, allow_masked=True):
  1290. """
  1291. Private function for the computation of covariance and correlation
  1292. coefficients.
  1293. """
  1294. x = ma.array(x, ndmin=2, copy=True, dtype=float)
  1295. xmask = ma.getmaskarray(x)
  1296. # Quick exit if we can't process masked data
  1297. if not allow_masked and xmask.any():
  1298. raise ValueError("Cannot process masked data.")
  1299. #
  1300. if x.shape[0] == 1:
  1301. rowvar = True
  1302. # Make sure that rowvar is either 0 or 1
  1303. rowvar = int(bool(rowvar))
  1304. axis = 1 - rowvar
  1305. if rowvar:
  1306. tup = (slice(None), None)
  1307. else:
  1308. tup = (None, slice(None))
  1309. #
  1310. if y is None:
  1311. # Check if we can guarantee that the integers in the (N - ddof)
  1312. # normalisation can be accurately represented with single-precision
  1313. # before computing the dot product.
  1314. if x.shape[0] > 2 ** 24 or x.shape[1] > 2 ** 24:
  1315. xnm_dtype = np.float64
  1316. else:
  1317. xnm_dtype = np.float32
  1318. xnotmask = np.logical_not(xmask).astype(xnm_dtype)
  1319. else:
  1320. y = array(y, copy=False, ndmin=2, dtype=float)
  1321. ymask = ma.getmaskarray(y)
  1322. if not allow_masked and ymask.any():
  1323. raise ValueError("Cannot process masked data.")
  1324. if xmask.any() or ymask.any():
  1325. if y.shape == x.shape:
  1326. # Define some common mask
  1327. common_mask = np.logical_or(xmask, ymask)
  1328. if common_mask is not nomask:
  1329. xmask = x._mask = y._mask = ymask = common_mask
  1330. x._sharedmask = False
  1331. y._sharedmask = False
  1332. x = ma.concatenate((x, y), axis)
  1333. # Check if we can guarantee that the integers in the (N - ddof)
  1334. # normalisation can be accurately represented with single-precision
  1335. # before computing the dot product.
  1336. if x.shape[0] > 2 ** 24 or x.shape[1] > 2 ** 24:
  1337. xnm_dtype = np.float64
  1338. else:
  1339. xnm_dtype = np.float32
  1340. xnotmask = np.logical_not(np.concatenate((xmask, ymask), axis)).astype(
  1341. xnm_dtype
  1342. )
  1343. x -= x.mean(axis=rowvar)[tup]
  1344. return (x, xnotmask, rowvar)
  1345. def cov(x, y=None, rowvar=True, bias=False, allow_masked=True, ddof=None):
  1346. """
  1347. Estimate the covariance matrix.
  1348. Except for the handling of missing data this function does the same as
  1349. `numpy.cov`. For more details and examples, see `numpy.cov`.
  1350. By default, masked values are recognized as such. If `x` and `y` have the
  1351. same shape, a common mask is allocated: if ``x[i,j]`` is masked, then
  1352. ``y[i,j]`` will also be masked.
  1353. Setting `allow_masked` to False will raise an exception if values are
  1354. missing in either of the input arrays.
  1355. Parameters
  1356. ----------
  1357. x : array_like
  1358. A 1-D or 2-D array containing multiple variables and observations.
  1359. Each row of `x` represents a variable, and each column a single
  1360. observation of all those variables. Also see `rowvar` below.
  1361. y : array_like, optional
  1362. An additional set of variables and observations. `y` has the same
  1363. shape as `x`.
  1364. rowvar : bool, optional
  1365. If `rowvar` is True (default), then each row represents a
  1366. variable, with observations in the columns. Otherwise, the relationship
  1367. is transposed: each column represents a variable, while the rows
  1368. contain observations.
  1369. bias : bool, optional
  1370. Default normalization (False) is by ``(N-1)``, where ``N`` is the
  1371. number of observations given (unbiased estimate). If `bias` is True,
  1372. then normalization is by ``N``. This keyword can be overridden by
  1373. the keyword ``ddof`` in numpy versions >= 1.5.
  1374. allow_masked : bool, optional
  1375. If True, masked values are propagated pair-wise: if a value is masked
  1376. in `x`, the corresponding value is masked in `y`.
  1377. If False, raises a `ValueError` exception when some values are missing.
  1378. ddof : {None, int}, optional
  1379. If not ``None`` normalization is by ``(N - ddof)``, where ``N`` is
  1380. the number of observations; this overrides the value implied by
  1381. ``bias``. The default value is ``None``.
  1382. Raises
  1383. ------
  1384. ValueError
  1385. Raised if some values are missing and `allow_masked` is False.
  1386. See Also
  1387. --------
  1388. numpy.cov
  1389. Examples
  1390. --------
  1391. >>> import numpy as np
  1392. >>> x = np.ma.array([[0, 1], [1, 1]], mask=[0, 1, 0, 1])
  1393. >>> y = np.ma.array([[1, 0], [0, 1]], mask=[0, 0, 1, 1])
  1394. >>> np.ma.cov(x, y)
  1395. masked_array(
  1396. data=[[--, --, --, --],
  1397. [--, --, --, --],
  1398. [--, --, --, --],
  1399. [--, --, --, --]],
  1400. mask=[[ True, True, True, True],
  1401. [ True, True, True, True],
  1402. [ True, True, True, True],
  1403. [ True, True, True, True]],
  1404. fill_value=1e+20,
  1405. dtype=float64)
  1406. """
  1407. # Check inputs
  1408. if ddof is not None and ddof != int(ddof):
  1409. raise ValueError("ddof must be an integer")
  1410. # Set up ddof
  1411. if ddof is None:
  1412. if bias:
  1413. ddof = 0
  1414. else:
  1415. ddof = 1
  1416. (x, xnotmask, rowvar) = _covhelper(x, y, rowvar, allow_masked)
  1417. if not rowvar:
  1418. fact = np.dot(xnotmask.T, xnotmask) - ddof
  1419. mask = np.less_equal(fact, 0, dtype=bool)
  1420. with np.errstate(divide="ignore", invalid="ignore"):
  1421. data = np.dot(filled(x.T, 0), filled(x.conj(), 0)) / fact
  1422. result = ma.array(data, mask=mask).squeeze()
  1423. else:
  1424. fact = np.dot(xnotmask, xnotmask.T) - ddof
  1425. mask = np.less_equal(fact, 0, dtype=bool)
  1426. with np.errstate(divide="ignore", invalid="ignore"):
  1427. data = np.dot(filled(x, 0), filled(x.T.conj(), 0)) / fact
  1428. result = ma.array(data, mask=mask).squeeze()
  1429. return result
  1430. def corrcoef(x, y=None, rowvar=True, allow_masked=True,
  1431. ):
  1432. """
  1433. Return Pearson product-moment correlation coefficients.
  1434. Except for the handling of missing data this function does the same as
  1435. `numpy.corrcoef`. For more details and examples, see `numpy.corrcoef`.
  1436. Parameters
  1437. ----------
  1438. x : array_like
  1439. A 1-D or 2-D array containing multiple variables and observations.
  1440. Each row of `x` represents a variable, and each column a single
  1441. observation of all those variables. Also see `rowvar` below.
  1442. y : array_like, optional
  1443. An additional set of variables and observations. `y` has the same
  1444. shape as `x`.
  1445. rowvar : bool, optional
  1446. If `rowvar` is True (default), then each row represents a
  1447. variable, with observations in the columns. Otherwise, the relationship
  1448. is transposed: each column represents a variable, while the rows
  1449. contain observations.
  1450. allow_masked : bool, optional
  1451. If True, masked values are propagated pair-wise: if a value is masked
  1452. in `x`, the corresponding value is masked in `y`.
  1453. If False, raises an exception. Because `bias` is deprecated, this
  1454. argument needs to be treated as keyword only to avoid a warning.
  1455. See Also
  1456. --------
  1457. numpy.corrcoef : Equivalent function in top-level NumPy module.
  1458. cov : Estimate the covariance matrix.
  1459. Examples
  1460. --------
  1461. >>> import numpy as np
  1462. >>> x = np.ma.array([[0, 1], [1, 1]], mask=[0, 1, 0, 1])
  1463. >>> np.ma.corrcoef(x)
  1464. masked_array(
  1465. data=[[--, --],
  1466. [--, --]],
  1467. mask=[[ True, True],
  1468. [ True, True]],
  1469. fill_value=1e+20,
  1470. dtype=float64)
  1471. """
  1472. # Estimate the covariance matrix.
  1473. corr = cov(x, y, rowvar, allow_masked=allow_masked)
  1474. # The non-masked version returns a masked value for a scalar.
  1475. try:
  1476. std = ma.sqrt(ma.diagonal(corr))
  1477. except ValueError:
  1478. return ma.MaskedConstant()
  1479. corr /= ma.multiply.outer(std, std)
  1480. return corr
  1481. #####--------------------------------------------------------------------------
  1482. #---- --- Concatenation helpers ---
  1483. #####--------------------------------------------------------------------------
  1484. class MAxisConcatenator(AxisConcatenator):
  1485. """
  1486. Translate slice objects to concatenation along an axis.
  1487. For documentation on usage, see `mr_class`.
  1488. See Also
  1489. --------
  1490. mr_class
  1491. """
  1492. __slots__ = ()
  1493. concatenate = staticmethod(concatenate)
  1494. @classmethod
  1495. def makemat(cls, arr):
  1496. # There used to be a view as np.matrix here, but we may eventually
  1497. # deprecate that class. In preparation, we use the unmasked version
  1498. # to construct the matrix (with copy=False for backwards compatibility
  1499. # with the .view)
  1500. data = super().makemat(arr.data, copy=False)
  1501. return array(data, mask=arr.mask)
  1502. def __getitem__(self, key):
  1503. # matrix builder syntax, like 'a, b; c, d'
  1504. if isinstance(key, str):
  1505. raise MAError("Unavailable for masked array.")
  1506. return super().__getitem__(key)
  1507. class mr_class(MAxisConcatenator):
  1508. """
  1509. Translate slice objects to concatenation along the first axis.
  1510. This is the masked array version of `r_`.
  1511. See Also
  1512. --------
  1513. r_
  1514. Examples
  1515. --------
  1516. >>> import numpy as np
  1517. >>> np.ma.mr_[np.ma.array([1,2,3]), 0, 0, np.ma.array([4,5,6])]
  1518. masked_array(data=[1, 2, 3, ..., 4, 5, 6],
  1519. mask=False,
  1520. fill_value=999999)
  1521. """
  1522. __slots__ = ()
  1523. def __init__(self):
  1524. MAxisConcatenator.__init__(self, 0)
  1525. mr_ = mr_class()
  1526. #####--------------------------------------------------------------------------
  1527. #---- Find unmasked data ---
  1528. #####--------------------------------------------------------------------------
  1529. def ndenumerate(a, compressed=True):
  1530. """
  1531. Multidimensional index iterator.
  1532. Return an iterator yielding pairs of array coordinates and values,
  1533. skipping elements that are masked. With `compressed=False`,
  1534. `ma.masked` is yielded as the value of masked elements. This
  1535. behavior differs from that of `numpy.ndenumerate`, which yields the
  1536. value of the underlying data array.
  1537. Notes
  1538. -----
  1539. .. versionadded:: 1.23.0
  1540. Parameters
  1541. ----------
  1542. a : array_like
  1543. An array with (possibly) masked elements.
  1544. compressed : bool, optional
  1545. If True (default), masked elements are skipped.
  1546. See Also
  1547. --------
  1548. numpy.ndenumerate : Equivalent function ignoring any mask.
  1549. Examples
  1550. --------
  1551. >>> import numpy as np
  1552. >>> a = np.ma.arange(9).reshape((3, 3))
  1553. >>> a[1, 0] = np.ma.masked
  1554. >>> a[1, 2] = np.ma.masked
  1555. >>> a[2, 1] = np.ma.masked
  1556. >>> a
  1557. masked_array(
  1558. data=[[0, 1, 2],
  1559. [--, 4, --],
  1560. [6, --, 8]],
  1561. mask=[[False, False, False],
  1562. [ True, False, True],
  1563. [False, True, False]],
  1564. fill_value=999999)
  1565. >>> for index, x in np.ma.ndenumerate(a):
  1566. ... print(index, x)
  1567. (0, 0) 0
  1568. (0, 1) 1
  1569. (0, 2) 2
  1570. (1, 1) 4
  1571. (2, 0) 6
  1572. (2, 2) 8
  1573. >>> for index, x in np.ma.ndenumerate(a, compressed=False):
  1574. ... print(index, x)
  1575. (0, 0) 0
  1576. (0, 1) 1
  1577. (0, 2) 2
  1578. (1, 0) --
  1579. (1, 1) 4
  1580. (1, 2) --
  1581. (2, 0) 6
  1582. (2, 1) --
  1583. (2, 2) 8
  1584. """
  1585. for it, mask in zip(np.ndenumerate(a), getmaskarray(a).flat):
  1586. if not mask:
  1587. yield it
  1588. elif not compressed:
  1589. yield it[0], masked
  1590. def flatnotmasked_edges(a):
  1591. """
  1592. Find the indices of the first and last unmasked values.
  1593. Expects a 1-D `MaskedArray`, returns None if all values are masked.
  1594. Parameters
  1595. ----------
  1596. a : array_like
  1597. Input 1-D `MaskedArray`
  1598. Returns
  1599. -------
  1600. edges : ndarray or None
  1601. The indices of first and last non-masked value in the array.
  1602. Returns None if all values are masked.
  1603. See Also
  1604. --------
  1605. flatnotmasked_contiguous, notmasked_contiguous, notmasked_edges
  1606. clump_masked, clump_unmasked
  1607. Notes
  1608. -----
  1609. Only accepts 1-D arrays.
  1610. Examples
  1611. --------
  1612. >>> import numpy as np
  1613. >>> a = np.ma.arange(10)
  1614. >>> np.ma.flatnotmasked_edges(a)
  1615. array([0, 9])
  1616. >>> mask = (a < 3) | (a > 8) | (a == 5)
  1617. >>> a[mask] = np.ma.masked
  1618. >>> np.array(a[~a.mask])
  1619. array([3, 4, 6, 7, 8])
  1620. >>> np.ma.flatnotmasked_edges(a)
  1621. array([3, 8])
  1622. >>> a[:] = np.ma.masked
  1623. >>> print(np.ma.flatnotmasked_edges(a))
  1624. None
  1625. """
  1626. m = getmask(a)
  1627. if m is nomask or not np.any(m):
  1628. return np.array([0, a.size - 1])
  1629. unmasked = np.flatnonzero(~m)
  1630. if len(unmasked) > 0:
  1631. return unmasked[[0, -1]]
  1632. else:
  1633. return None
  1634. def notmasked_edges(a, axis=None):
  1635. """
  1636. Find the indices of the first and last unmasked values along an axis.
  1637. If all values are masked, return None. Otherwise, return a list
  1638. of two tuples, corresponding to the indices of the first and last
  1639. unmasked values respectively.
  1640. Parameters
  1641. ----------
  1642. a : array_like
  1643. The input array.
  1644. axis : int, optional
  1645. Axis along which to perform the operation.
  1646. If None (default), applies to a flattened version of the array.
  1647. Returns
  1648. -------
  1649. edges : ndarray or list
  1650. An array of start and end indexes if there are any masked data in
  1651. the array. If there are no masked data in the array, `edges` is a
  1652. list of the first and last index.
  1653. See Also
  1654. --------
  1655. flatnotmasked_contiguous, flatnotmasked_edges, notmasked_contiguous
  1656. clump_masked, clump_unmasked
  1657. Examples
  1658. --------
  1659. >>> import numpy as np
  1660. >>> a = np.arange(9).reshape((3, 3))
  1661. >>> m = np.zeros_like(a)
  1662. >>> m[1:, 1:] = 1
  1663. >>> am = np.ma.array(a, mask=m)
  1664. >>> np.array(am[~am.mask])
  1665. array([0, 1, 2, 3, 6])
  1666. >>> np.ma.notmasked_edges(am)
  1667. array([0, 6])
  1668. """
  1669. a = asarray(a)
  1670. if axis is None or a.ndim == 1:
  1671. return flatnotmasked_edges(a)
  1672. m = getmaskarray(a)
  1673. idx = array(np.indices(a.shape), mask=np.asarray([m] * a.ndim))
  1674. return [tuple(idx[i].min(axis).compressed() for i in range(a.ndim)),
  1675. tuple(idx[i].max(axis).compressed() for i in range(a.ndim)), ]
  1676. def flatnotmasked_contiguous(a):
  1677. """
  1678. Find contiguous unmasked data in a masked array.
  1679. Parameters
  1680. ----------
  1681. a : array_like
  1682. The input array.
  1683. Returns
  1684. -------
  1685. slice_list : list
  1686. A sorted sequence of `slice` objects (start index, end index).
  1687. See Also
  1688. --------
  1689. flatnotmasked_edges, notmasked_contiguous, notmasked_edges
  1690. clump_masked, clump_unmasked
  1691. Notes
  1692. -----
  1693. Only accepts 2-D arrays at most.
  1694. Examples
  1695. --------
  1696. >>> import numpy as np
  1697. >>> a = np.ma.arange(10)
  1698. >>> np.ma.flatnotmasked_contiguous(a)
  1699. [slice(0, 10, None)]
  1700. >>> mask = (a < 3) | (a > 8) | (a == 5)
  1701. >>> a[mask] = np.ma.masked
  1702. >>> np.array(a[~a.mask])
  1703. array([3, 4, 6, 7, 8])
  1704. >>> np.ma.flatnotmasked_contiguous(a)
  1705. [slice(3, 5, None), slice(6, 9, None)]
  1706. >>> a[:] = np.ma.masked
  1707. >>> np.ma.flatnotmasked_contiguous(a)
  1708. []
  1709. """
  1710. m = getmask(a)
  1711. if m is nomask:
  1712. return [slice(0, a.size)]
  1713. i = 0
  1714. result = []
  1715. for (k, g) in itertools.groupby(m.ravel()):
  1716. n = len(list(g))
  1717. if not k:
  1718. result.append(slice(i, i + n))
  1719. i += n
  1720. return result
  1721. def notmasked_contiguous(a, axis=None):
  1722. """
  1723. Find contiguous unmasked data in a masked array along the given axis.
  1724. Parameters
  1725. ----------
  1726. a : array_like
  1727. The input array.
  1728. axis : int, optional
  1729. Axis along which to perform the operation.
  1730. If None (default), applies to a flattened version of the array, and this
  1731. is the same as `flatnotmasked_contiguous`.
  1732. Returns
  1733. -------
  1734. endpoints : list
  1735. A list of slices (start and end indexes) of unmasked indexes
  1736. in the array.
  1737. If the input is 2d and axis is specified, the result is a list of lists.
  1738. See Also
  1739. --------
  1740. flatnotmasked_edges, flatnotmasked_contiguous, notmasked_edges
  1741. clump_masked, clump_unmasked
  1742. Notes
  1743. -----
  1744. Only accepts 2-D arrays at most.
  1745. Examples
  1746. --------
  1747. >>> import numpy as np
  1748. >>> a = np.arange(12).reshape((3, 4))
  1749. >>> mask = np.zeros_like(a)
  1750. >>> mask[1:, :-1] = 1; mask[0, 1] = 1; mask[-1, 0] = 0
  1751. >>> ma = np.ma.array(a, mask=mask)
  1752. >>> ma
  1753. masked_array(
  1754. data=[[0, --, 2, 3],
  1755. [--, --, --, 7],
  1756. [8, --, --, 11]],
  1757. mask=[[False, True, False, False],
  1758. [ True, True, True, False],
  1759. [False, True, True, False]],
  1760. fill_value=999999)
  1761. >>> np.array(ma[~ma.mask])
  1762. array([ 0, 2, 3, 7, 8, 11])
  1763. >>> np.ma.notmasked_contiguous(ma)
  1764. [slice(0, 1, None), slice(2, 4, None), slice(7, 9, None), slice(11, 12, None)]
  1765. >>> np.ma.notmasked_contiguous(ma, axis=0)
  1766. [[slice(0, 1, None), slice(2, 3, None)], [], [slice(0, 1, None)], [slice(0, 3, None)]]
  1767. >>> np.ma.notmasked_contiguous(ma, axis=1)
  1768. [[slice(0, 1, None), slice(2, 4, None)], [slice(3, 4, None)], [slice(0, 1, None), slice(3, 4, None)]]
  1769. """ # noqa: E501
  1770. a = asarray(a)
  1771. nd = a.ndim
  1772. if nd > 2:
  1773. raise NotImplementedError("Currently limited to at most 2D array.")
  1774. if axis is None or nd == 1:
  1775. return flatnotmasked_contiguous(a)
  1776. #
  1777. result = []
  1778. #
  1779. other = (axis + 1) % 2
  1780. idx = [0, 0]
  1781. idx[axis] = slice(None, None)
  1782. #
  1783. for i in range(a.shape[other]):
  1784. idx[other] = i
  1785. result.append(flatnotmasked_contiguous(a[tuple(idx)]))
  1786. return result
  1787. def _ezclump(mask):
  1788. """
  1789. Finds the clumps (groups of data with the same values) for a 1D bool array.
  1790. Returns a series of slices.
  1791. """
  1792. if mask.ndim > 1:
  1793. mask = mask.ravel()
  1794. idx = (mask[1:] ^ mask[:-1]).nonzero()
  1795. idx = idx[0] + 1
  1796. if mask[0]:
  1797. if len(idx) == 0:
  1798. return [slice(0, mask.size)]
  1799. r = [slice(0, idx[0])]
  1800. r.extend((slice(left, right)
  1801. for left, right in zip(idx[1:-1:2], idx[2::2])))
  1802. else:
  1803. if len(idx) == 0:
  1804. return []
  1805. r = [slice(left, right) for left, right in zip(idx[:-1:2], idx[1::2])]
  1806. if mask[-1]:
  1807. r.append(slice(idx[-1], mask.size))
  1808. return r
  1809. def clump_unmasked(a):
  1810. """
  1811. Return list of slices corresponding to the unmasked clumps of a 1-D array.
  1812. (A "clump" is defined as a contiguous region of the array).
  1813. Parameters
  1814. ----------
  1815. a : ndarray
  1816. A one-dimensional masked array.
  1817. Returns
  1818. -------
  1819. slices : list of slice
  1820. The list of slices, one for each continuous region of unmasked
  1821. elements in `a`.
  1822. See Also
  1823. --------
  1824. flatnotmasked_edges, flatnotmasked_contiguous, notmasked_edges
  1825. notmasked_contiguous, clump_masked
  1826. Examples
  1827. --------
  1828. >>> import numpy as np
  1829. >>> a = np.ma.masked_array(np.arange(10))
  1830. >>> a[[0, 1, 2, 6, 8, 9]] = np.ma.masked
  1831. >>> np.ma.clump_unmasked(a)
  1832. [slice(3, 6, None), slice(7, 8, None)]
  1833. """
  1834. mask = getattr(a, '_mask', nomask)
  1835. if mask is nomask:
  1836. return [slice(0, a.size)]
  1837. return _ezclump(~mask)
  1838. def clump_masked(a):
  1839. """
  1840. Returns a list of slices corresponding to the masked clumps of a 1-D array.
  1841. (A "clump" is defined as a contiguous region of the array).
  1842. Parameters
  1843. ----------
  1844. a : ndarray
  1845. A one-dimensional masked array.
  1846. Returns
  1847. -------
  1848. slices : list of slice
  1849. The list of slices, one for each continuous region of masked elements
  1850. in `a`.
  1851. See Also
  1852. --------
  1853. flatnotmasked_edges, flatnotmasked_contiguous, notmasked_edges
  1854. notmasked_contiguous, clump_unmasked
  1855. Examples
  1856. --------
  1857. >>> import numpy as np
  1858. >>> a = np.ma.masked_array(np.arange(10))
  1859. >>> a[[0, 1, 2, 6, 8, 9]] = np.ma.masked
  1860. >>> np.ma.clump_masked(a)
  1861. [slice(0, 3, None), slice(6, 7, None), slice(8, 10, None)]
  1862. """
  1863. mask = ma.getmask(a)
  1864. if mask is nomask:
  1865. return []
  1866. return _ezclump(mask)
  1867. ###############################################################################
  1868. # Polynomial fit #
  1869. ###############################################################################
  1870. def vander(x, n=None):
  1871. """
  1872. Masked values in the input array result in rows of zeros.
  1873. """
  1874. _vander = np.vander(x, n)
  1875. m = getmask(x)
  1876. if m is not nomask:
  1877. _vander[m] = 0
  1878. return _vander
  1879. vander.__doc__ = ma.doc_note(np.vander.__doc__, vander.__doc__)
  1880. def polyfit(x, y, deg, rcond=None, full=False, w=None, cov=False):
  1881. """
  1882. Any masked values in x is propagated in y, and vice-versa.
  1883. """
  1884. x = asarray(x)
  1885. y = asarray(y)
  1886. m = getmask(x)
  1887. if y.ndim == 1:
  1888. m = mask_or(m, getmask(y))
  1889. elif y.ndim == 2:
  1890. my = getmask(mask_rows(y))
  1891. if my is not nomask:
  1892. m = mask_or(m, my[:, 0])
  1893. else:
  1894. raise TypeError("Expected a 1D or 2D array for y!")
  1895. if w is not None:
  1896. w = asarray(w)
  1897. if w.ndim != 1:
  1898. raise TypeError("expected a 1-d array for weights")
  1899. if w.shape[0] != y.shape[0]:
  1900. raise TypeError("expected w and y to have the same length")
  1901. m = mask_or(m, getmask(w))
  1902. if m is not nomask:
  1903. not_m = ~m
  1904. if w is not None:
  1905. w = w[not_m]
  1906. return np.polyfit(x[not_m], y[not_m], deg, rcond, full, w, cov)
  1907. else:
  1908. return np.polyfit(x, y, deg, rcond, full, w, cov)
  1909. polyfit.__doc__ = ma.doc_note(np.polyfit.__doc__, polyfit.__doc__)