test_extras.py 65 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620162116221623162416251626162716281629163016311632163316341635163616371638163916401641164216431644164516461647164816491650165116521653165416551656165716581659166016611662166316641665166616671668
  1. # pylint: disable-msg=W0611, W0612, W0511
  2. """Tests suite for MaskedArray.
  3. Adapted from the original test_ma by Pierre Gerard-Marchant
  4. :author: Pierre Gerard-Marchant
  5. :contact: pierregm_at_uga_dot_edu
  6. :version: $Id: test_extras.py 3473 2007-10-29 15:18:13Z jarrod.millman $
  7. """
  8. from __future__ import division, absolute_import, print_function
  9. import warnings
  10. import itertools
  11. import pytest
  12. import numpy as np
  13. from numpy.testing import (
  14. assert_warns, suppress_warnings
  15. )
  16. from numpy.ma.testutils import (
  17. assert_, assert_array_equal, assert_equal, assert_almost_equal
  18. )
  19. from numpy.ma.core import (
  20. array, arange, masked, MaskedArray, masked_array, getmaskarray, shape,
  21. nomask, ones, zeros, count
  22. )
  23. from numpy.ma.extras import (
  24. atleast_1d, atleast_2d, atleast_3d, mr_, dot, polyfit, cov, corrcoef,
  25. median, average, unique, setxor1d, setdiff1d, union1d, intersect1d, in1d,
  26. ediff1d, apply_over_axes, apply_along_axis, compress_nd, compress_rowcols,
  27. mask_rowcols, clump_masked, clump_unmasked, flatnotmasked_contiguous,
  28. notmasked_contiguous, notmasked_edges, masked_all, masked_all_like, isin,
  29. diagflat, stack, vstack
  30. )
  31. class TestGeneric(object):
  32. #
  33. def test_masked_all(self):
  34. # Tests masked_all
  35. # Standard dtype
  36. test = masked_all((2,), dtype=float)
  37. control = array([1, 1], mask=[1, 1], dtype=float)
  38. assert_equal(test, control)
  39. # Flexible dtype
  40. dt = np.dtype({'names': ['a', 'b'], 'formats': ['f', 'f']})
  41. test = masked_all((2,), dtype=dt)
  42. control = array([(0, 0), (0, 0)], mask=[(1, 1), (1, 1)], dtype=dt)
  43. assert_equal(test, control)
  44. test = masked_all((2, 2), dtype=dt)
  45. control = array([[(0, 0), (0, 0)], [(0, 0), (0, 0)]],
  46. mask=[[(1, 1), (1, 1)], [(1, 1), (1, 1)]],
  47. dtype=dt)
  48. assert_equal(test, control)
  49. # Nested dtype
  50. dt = np.dtype([('a', 'f'), ('b', [('ba', 'f'), ('bb', 'f')])])
  51. test = masked_all((2,), dtype=dt)
  52. control = array([(1, (1, 1)), (1, (1, 1))],
  53. mask=[(1, (1, 1)), (1, (1, 1))], dtype=dt)
  54. assert_equal(test, control)
  55. test = masked_all((2,), dtype=dt)
  56. control = array([(1, (1, 1)), (1, (1, 1))],
  57. mask=[(1, (1, 1)), (1, (1, 1))], dtype=dt)
  58. assert_equal(test, control)
  59. test = masked_all((1, 1), dtype=dt)
  60. control = array([[(1, (1, 1))]], mask=[[(1, (1, 1))]], dtype=dt)
  61. assert_equal(test, control)
  62. def test_masked_all_like(self):
  63. # Tests masked_all
  64. # Standard dtype
  65. base = array([1, 2], dtype=float)
  66. test = masked_all_like(base)
  67. control = array([1, 1], mask=[1, 1], dtype=float)
  68. assert_equal(test, control)
  69. # Flexible dtype
  70. dt = np.dtype({'names': ['a', 'b'], 'formats': ['f', 'f']})
  71. base = array([(0, 0), (0, 0)], mask=[(1, 1), (1, 1)], dtype=dt)
  72. test = masked_all_like(base)
  73. control = array([(10, 10), (10, 10)], mask=[(1, 1), (1, 1)], dtype=dt)
  74. assert_equal(test, control)
  75. # Nested dtype
  76. dt = np.dtype([('a', 'f'), ('b', [('ba', 'f'), ('bb', 'f')])])
  77. control = array([(1, (1, 1)), (1, (1, 1))],
  78. mask=[(1, (1, 1)), (1, (1, 1))], dtype=dt)
  79. test = masked_all_like(control)
  80. assert_equal(test, control)
  81. def check_clump(self, f):
  82. for i in range(1, 7):
  83. for j in range(2**i):
  84. k = np.arange(i, dtype=int)
  85. ja = np.full(i, j, dtype=int)
  86. a = masked_array(2**k)
  87. a.mask = (ja & (2**k)) != 0
  88. s = 0
  89. for sl in f(a):
  90. s += a.data[sl].sum()
  91. if f == clump_unmasked:
  92. assert_equal(a.compressed().sum(), s)
  93. else:
  94. a.mask = ~a.mask
  95. assert_equal(a.compressed().sum(), s)
  96. def test_clump_masked(self):
  97. # Test clump_masked
  98. a = masked_array(np.arange(10))
  99. a[[0, 1, 2, 6, 8, 9]] = masked
  100. #
  101. test = clump_masked(a)
  102. control = [slice(0, 3), slice(6, 7), slice(8, 10)]
  103. assert_equal(test, control)
  104. self.check_clump(clump_masked)
  105. def test_clump_unmasked(self):
  106. # Test clump_unmasked
  107. a = masked_array(np.arange(10))
  108. a[[0, 1, 2, 6, 8, 9]] = masked
  109. test = clump_unmasked(a)
  110. control = [slice(3, 6), slice(7, 8), ]
  111. assert_equal(test, control)
  112. self.check_clump(clump_unmasked)
  113. def test_flatnotmasked_contiguous(self):
  114. # Test flatnotmasked_contiguous
  115. a = arange(10)
  116. # No mask
  117. test = flatnotmasked_contiguous(a)
  118. assert_equal(test, [slice(0, a.size)])
  119. # mask of all false
  120. a.mask = np.zeros(10, dtype=bool)
  121. assert_equal(test, [slice(0, a.size)])
  122. # Some mask
  123. a[(a < 3) | (a > 8) | (a == 5)] = masked
  124. test = flatnotmasked_contiguous(a)
  125. assert_equal(test, [slice(3, 5), slice(6, 9)])
  126. #
  127. a[:] = masked
  128. test = flatnotmasked_contiguous(a)
  129. assert_equal(test, [])
  130. class TestAverage(object):
  131. # Several tests of average. Why so many ? Good point...
  132. def test_testAverage1(self):
  133. # Test of average.
  134. ott = array([0., 1., 2., 3.], mask=[True, False, False, False])
  135. assert_equal(2.0, average(ott, axis=0))
  136. assert_equal(2.0, average(ott, weights=[1., 1., 2., 1.]))
  137. result, wts = average(ott, weights=[1., 1., 2., 1.], returned=True)
  138. assert_equal(2.0, result)
  139. assert_(wts == 4.0)
  140. ott[:] = masked
  141. assert_equal(average(ott, axis=0).mask, [True])
  142. ott = array([0., 1., 2., 3.], mask=[True, False, False, False])
  143. ott = ott.reshape(2, 2)
  144. ott[:, 1] = masked
  145. assert_equal(average(ott, axis=0), [2.0, 0.0])
  146. assert_equal(average(ott, axis=1).mask[0], [True])
  147. assert_equal([2., 0.], average(ott, axis=0))
  148. result, wts = average(ott, axis=0, returned=True)
  149. assert_equal(wts, [1., 0.])
  150. def test_testAverage2(self):
  151. # More tests of average.
  152. w1 = [0, 1, 1, 1, 1, 0]
  153. w2 = [[0, 1, 1, 1, 1, 0], [1, 0, 0, 0, 0, 1]]
  154. x = arange(6, dtype=np.float_)
  155. assert_equal(average(x, axis=0), 2.5)
  156. assert_equal(average(x, axis=0, weights=w1), 2.5)
  157. y = array([arange(6, dtype=np.float_), 2.0 * arange(6)])
  158. assert_equal(average(y, None), np.add.reduce(np.arange(6)) * 3. / 12.)
  159. assert_equal(average(y, axis=0), np.arange(6) * 3. / 2.)
  160. assert_equal(average(y, axis=1),
  161. [average(x, axis=0), average(x, axis=0) * 2.0])
  162. assert_equal(average(y, None, weights=w2), 20. / 6.)
  163. assert_equal(average(y, axis=0, weights=w2),
  164. [0., 1., 2., 3., 4., 10.])
  165. assert_equal(average(y, axis=1),
  166. [average(x, axis=0), average(x, axis=0) * 2.0])
  167. m1 = zeros(6)
  168. m2 = [0, 0, 1, 1, 0, 0]
  169. m3 = [[0, 0, 1, 1, 0, 0], [0, 1, 1, 1, 1, 0]]
  170. m4 = ones(6)
  171. m5 = [0, 1, 1, 1, 1, 1]
  172. assert_equal(average(masked_array(x, m1), axis=0), 2.5)
  173. assert_equal(average(masked_array(x, m2), axis=0), 2.5)
  174. assert_equal(average(masked_array(x, m4), axis=0).mask, [True])
  175. assert_equal(average(masked_array(x, m5), axis=0), 0.0)
  176. assert_equal(count(average(masked_array(x, m4), axis=0)), 0)
  177. z = masked_array(y, m3)
  178. assert_equal(average(z, None), 20. / 6.)
  179. assert_equal(average(z, axis=0), [0., 1., 99., 99., 4.0, 7.5])
  180. assert_equal(average(z, axis=1), [2.5, 5.0])
  181. assert_equal(average(z, axis=0, weights=w2),
  182. [0., 1., 99., 99., 4.0, 10.0])
  183. def test_testAverage3(self):
  184. # Yet more tests of average!
  185. a = arange(6)
  186. b = arange(6) * 3
  187. r1, w1 = average([[a, b], [b, a]], axis=1, returned=True)
  188. assert_equal(shape(r1), shape(w1))
  189. assert_equal(r1.shape, w1.shape)
  190. r2, w2 = average(ones((2, 2, 3)), axis=0, weights=[3, 1], returned=True)
  191. assert_equal(shape(w2), shape(r2))
  192. r2, w2 = average(ones((2, 2, 3)), returned=True)
  193. assert_equal(shape(w2), shape(r2))
  194. r2, w2 = average(ones((2, 2, 3)), weights=ones((2, 2, 3)), returned=True)
  195. assert_equal(shape(w2), shape(r2))
  196. a2d = array([[1, 2], [0, 4]], float)
  197. a2dm = masked_array(a2d, [[False, False], [True, False]])
  198. a2da = average(a2d, axis=0)
  199. assert_equal(a2da, [0.5, 3.0])
  200. a2dma = average(a2dm, axis=0)
  201. assert_equal(a2dma, [1.0, 3.0])
  202. a2dma = average(a2dm, axis=None)
  203. assert_equal(a2dma, 7. / 3.)
  204. a2dma = average(a2dm, axis=1)
  205. assert_equal(a2dma, [1.5, 4.0])
  206. def test_onintegers_with_mask(self):
  207. # Test average on integers with mask
  208. a = average(array([1, 2]))
  209. assert_equal(a, 1.5)
  210. a = average(array([1, 2, 3, 4], mask=[False, False, True, True]))
  211. assert_equal(a, 1.5)
  212. def test_complex(self):
  213. # Test with complex data.
  214. # (Regression test for https://github.com/numpy/numpy/issues/2684)
  215. mask = np.array([[0, 0, 0, 1, 0],
  216. [0, 1, 0, 0, 0]], dtype=bool)
  217. a = masked_array([[0, 1+2j, 3+4j, 5+6j, 7+8j],
  218. [9j, 0+1j, 2+3j, 4+5j, 7+7j]],
  219. mask=mask)
  220. av = average(a)
  221. expected = np.average(a.compressed())
  222. assert_almost_equal(av.real, expected.real)
  223. assert_almost_equal(av.imag, expected.imag)
  224. av0 = average(a, axis=0)
  225. expected0 = average(a.real, axis=0) + average(a.imag, axis=0)*1j
  226. assert_almost_equal(av0.real, expected0.real)
  227. assert_almost_equal(av0.imag, expected0.imag)
  228. av1 = average(a, axis=1)
  229. expected1 = average(a.real, axis=1) + average(a.imag, axis=1)*1j
  230. assert_almost_equal(av1.real, expected1.real)
  231. assert_almost_equal(av1.imag, expected1.imag)
  232. # Test with the 'weights' argument.
  233. wts = np.array([[0.5, 1.0, 2.0, 1.0, 0.5],
  234. [1.0, 1.0, 1.0, 1.0, 1.0]])
  235. wav = average(a, weights=wts)
  236. expected = np.average(a.compressed(), weights=wts[~mask])
  237. assert_almost_equal(wav.real, expected.real)
  238. assert_almost_equal(wav.imag, expected.imag)
  239. wav0 = average(a, weights=wts, axis=0)
  240. expected0 = (average(a.real, weights=wts, axis=0) +
  241. average(a.imag, weights=wts, axis=0)*1j)
  242. assert_almost_equal(wav0.real, expected0.real)
  243. assert_almost_equal(wav0.imag, expected0.imag)
  244. wav1 = average(a, weights=wts, axis=1)
  245. expected1 = (average(a.real, weights=wts, axis=1) +
  246. average(a.imag, weights=wts, axis=1)*1j)
  247. assert_almost_equal(wav1.real, expected1.real)
  248. assert_almost_equal(wav1.imag, expected1.imag)
  249. class TestConcatenator(object):
  250. # Tests for mr_, the equivalent of r_ for masked arrays.
  251. def test_1d(self):
  252. # Tests mr_ on 1D arrays.
  253. assert_array_equal(mr_[1, 2, 3, 4, 5, 6], array([1, 2, 3, 4, 5, 6]))
  254. b = ones(5)
  255. m = [1, 0, 0, 0, 0]
  256. d = masked_array(b, mask=m)
  257. c = mr_[d, 0, 0, d]
  258. assert_(isinstance(c, MaskedArray))
  259. assert_array_equal(c, [1, 1, 1, 1, 1, 0, 0, 1, 1, 1, 1, 1])
  260. assert_array_equal(c.mask, mr_[m, 0, 0, m])
  261. def test_2d(self):
  262. # Tests mr_ on 2D arrays.
  263. a_1 = np.random.rand(5, 5)
  264. a_2 = np.random.rand(5, 5)
  265. m_1 = np.round_(np.random.rand(5, 5), 0)
  266. m_2 = np.round_(np.random.rand(5, 5), 0)
  267. b_1 = masked_array(a_1, mask=m_1)
  268. b_2 = masked_array(a_2, mask=m_2)
  269. # append columns
  270. d = mr_['1', b_1, b_2]
  271. assert_(d.shape == (5, 10))
  272. assert_array_equal(d[:, :5], b_1)
  273. assert_array_equal(d[:, 5:], b_2)
  274. assert_array_equal(d.mask, np.r_['1', m_1, m_2])
  275. d = mr_[b_1, b_2]
  276. assert_(d.shape == (10, 5))
  277. assert_array_equal(d[:5,:], b_1)
  278. assert_array_equal(d[5:,:], b_2)
  279. assert_array_equal(d.mask, np.r_[m_1, m_2])
  280. def test_masked_constant(self):
  281. actual = mr_[np.ma.masked, 1]
  282. assert_equal(actual.mask, [True, False])
  283. assert_equal(actual.data[1], 1)
  284. actual = mr_[[1, 2], np.ma.masked]
  285. assert_equal(actual.mask, [False, False, True])
  286. assert_equal(actual.data[:2], [1, 2])
  287. class TestNotMasked(object):
  288. # Tests notmasked_edges and notmasked_contiguous.
  289. def test_edges(self):
  290. # Tests unmasked_edges
  291. data = masked_array(np.arange(25).reshape(5, 5),
  292. mask=[[0, 0, 1, 0, 0],
  293. [0, 0, 0, 1, 1],
  294. [1, 1, 0, 0, 0],
  295. [0, 0, 0, 0, 0],
  296. [1, 1, 1, 0, 0]],)
  297. test = notmasked_edges(data, None)
  298. assert_equal(test, [0, 24])
  299. test = notmasked_edges(data, 0)
  300. assert_equal(test[0], [(0, 0, 1, 0, 0), (0, 1, 2, 3, 4)])
  301. assert_equal(test[1], [(3, 3, 3, 4, 4), (0, 1, 2, 3, 4)])
  302. test = notmasked_edges(data, 1)
  303. assert_equal(test[0], [(0, 1, 2, 3, 4), (0, 0, 2, 0, 3)])
  304. assert_equal(test[1], [(0, 1, 2, 3, 4), (4, 2, 4, 4, 4)])
  305. #
  306. test = notmasked_edges(data.data, None)
  307. assert_equal(test, [0, 24])
  308. test = notmasked_edges(data.data, 0)
  309. assert_equal(test[0], [(0, 0, 0, 0, 0), (0, 1, 2, 3, 4)])
  310. assert_equal(test[1], [(4, 4, 4, 4, 4), (0, 1, 2, 3, 4)])
  311. test = notmasked_edges(data.data, -1)
  312. assert_equal(test[0], [(0, 1, 2, 3, 4), (0, 0, 0, 0, 0)])
  313. assert_equal(test[1], [(0, 1, 2, 3, 4), (4, 4, 4, 4, 4)])
  314. #
  315. data[-2] = masked
  316. test = notmasked_edges(data, 0)
  317. assert_equal(test[0], [(0, 0, 1, 0, 0), (0, 1, 2, 3, 4)])
  318. assert_equal(test[1], [(1, 1, 2, 4, 4), (0, 1, 2, 3, 4)])
  319. test = notmasked_edges(data, -1)
  320. assert_equal(test[0], [(0, 1, 2, 4), (0, 0, 2, 3)])
  321. assert_equal(test[1], [(0, 1, 2, 4), (4, 2, 4, 4)])
  322. def test_contiguous(self):
  323. # Tests notmasked_contiguous
  324. a = masked_array(np.arange(24).reshape(3, 8),
  325. mask=[[0, 0, 0, 0, 1, 1, 1, 1],
  326. [1, 1, 1, 1, 1, 1, 1, 1],
  327. [0, 0, 0, 0, 0, 0, 1, 0]])
  328. tmp = notmasked_contiguous(a, None)
  329. assert_equal(tmp, [
  330. slice(0, 4, None),
  331. slice(16, 22, None),
  332. slice(23, 24, None)
  333. ])
  334. tmp = notmasked_contiguous(a, 0)
  335. assert_equal(tmp, [
  336. [slice(0, 1, None), slice(2, 3, None)],
  337. [slice(0, 1, None), slice(2, 3, None)],
  338. [slice(0, 1, None), slice(2, 3, None)],
  339. [slice(0, 1, None), slice(2, 3, None)],
  340. [slice(2, 3, None)],
  341. [slice(2, 3, None)],
  342. [],
  343. [slice(2, 3, None)]
  344. ])
  345. #
  346. tmp = notmasked_contiguous(a, 1)
  347. assert_equal(tmp, [
  348. [slice(0, 4, None)],
  349. [],
  350. [slice(0, 6, None), slice(7, 8, None)]
  351. ])
  352. class TestCompressFunctions(object):
  353. def test_compress_nd(self):
  354. # Tests compress_nd
  355. x = np.array(list(range(3*4*5))).reshape(3, 4, 5)
  356. m = np.zeros((3,4,5)).astype(bool)
  357. m[1,1,1] = True
  358. x = array(x, mask=m)
  359. # axis=None
  360. a = compress_nd(x)
  361. assert_equal(a, [[[ 0, 2, 3, 4],
  362. [10, 12, 13, 14],
  363. [15, 17, 18, 19]],
  364. [[40, 42, 43, 44],
  365. [50, 52, 53, 54],
  366. [55, 57, 58, 59]]])
  367. # axis=0
  368. a = compress_nd(x, 0)
  369. assert_equal(a, [[[ 0, 1, 2, 3, 4],
  370. [ 5, 6, 7, 8, 9],
  371. [10, 11, 12, 13, 14],
  372. [15, 16, 17, 18, 19]],
  373. [[40, 41, 42, 43, 44],
  374. [45, 46, 47, 48, 49],
  375. [50, 51, 52, 53, 54],
  376. [55, 56, 57, 58, 59]]])
  377. # axis=1
  378. a = compress_nd(x, 1)
  379. assert_equal(a, [[[ 0, 1, 2, 3, 4],
  380. [10, 11, 12, 13, 14],
  381. [15, 16, 17, 18, 19]],
  382. [[20, 21, 22, 23, 24],
  383. [30, 31, 32, 33, 34],
  384. [35, 36, 37, 38, 39]],
  385. [[40, 41, 42, 43, 44],
  386. [50, 51, 52, 53, 54],
  387. [55, 56, 57, 58, 59]]])
  388. a2 = compress_nd(x, (1,))
  389. a3 = compress_nd(x, -2)
  390. a4 = compress_nd(x, (-2,))
  391. assert_equal(a, a2)
  392. assert_equal(a, a3)
  393. assert_equal(a, a4)
  394. # axis=2
  395. a = compress_nd(x, 2)
  396. assert_equal(a, [[[ 0, 2, 3, 4],
  397. [ 5, 7, 8, 9],
  398. [10, 12, 13, 14],
  399. [15, 17, 18, 19]],
  400. [[20, 22, 23, 24],
  401. [25, 27, 28, 29],
  402. [30, 32, 33, 34],
  403. [35, 37, 38, 39]],
  404. [[40, 42, 43, 44],
  405. [45, 47, 48, 49],
  406. [50, 52, 53, 54],
  407. [55, 57, 58, 59]]])
  408. a2 = compress_nd(x, (2,))
  409. a3 = compress_nd(x, -1)
  410. a4 = compress_nd(x, (-1,))
  411. assert_equal(a, a2)
  412. assert_equal(a, a3)
  413. assert_equal(a, a4)
  414. # axis=(0, 1)
  415. a = compress_nd(x, (0, 1))
  416. assert_equal(a, [[[ 0, 1, 2, 3, 4],
  417. [10, 11, 12, 13, 14],
  418. [15, 16, 17, 18, 19]],
  419. [[40, 41, 42, 43, 44],
  420. [50, 51, 52, 53, 54],
  421. [55, 56, 57, 58, 59]]])
  422. a2 = compress_nd(x, (0, -2))
  423. assert_equal(a, a2)
  424. # axis=(1, 2)
  425. a = compress_nd(x, (1, 2))
  426. assert_equal(a, [[[ 0, 2, 3, 4],
  427. [10, 12, 13, 14],
  428. [15, 17, 18, 19]],
  429. [[20, 22, 23, 24],
  430. [30, 32, 33, 34],
  431. [35, 37, 38, 39]],
  432. [[40, 42, 43, 44],
  433. [50, 52, 53, 54],
  434. [55, 57, 58, 59]]])
  435. a2 = compress_nd(x, (-2, 2))
  436. a3 = compress_nd(x, (1, -1))
  437. a4 = compress_nd(x, (-2, -1))
  438. assert_equal(a, a2)
  439. assert_equal(a, a3)
  440. assert_equal(a, a4)
  441. # axis=(0, 2)
  442. a = compress_nd(x, (0, 2))
  443. assert_equal(a, [[[ 0, 2, 3, 4],
  444. [ 5, 7, 8, 9],
  445. [10, 12, 13, 14],
  446. [15, 17, 18, 19]],
  447. [[40, 42, 43, 44],
  448. [45, 47, 48, 49],
  449. [50, 52, 53, 54],
  450. [55, 57, 58, 59]]])
  451. a2 = compress_nd(x, (0, -1))
  452. assert_equal(a, a2)
  453. def test_compress_rowcols(self):
  454. # Tests compress_rowcols
  455. x = array(np.arange(9).reshape(3, 3),
  456. mask=[[1, 0, 0], [0, 0, 0], [0, 0, 0]])
  457. assert_equal(compress_rowcols(x), [[4, 5], [7, 8]])
  458. assert_equal(compress_rowcols(x, 0), [[3, 4, 5], [6, 7, 8]])
  459. assert_equal(compress_rowcols(x, 1), [[1, 2], [4, 5], [7, 8]])
  460. x = array(x._data, mask=[[0, 0, 0], [0, 1, 0], [0, 0, 0]])
  461. assert_equal(compress_rowcols(x), [[0, 2], [6, 8]])
  462. assert_equal(compress_rowcols(x, 0), [[0, 1, 2], [6, 7, 8]])
  463. assert_equal(compress_rowcols(x, 1), [[0, 2], [3, 5], [6, 8]])
  464. x = array(x._data, mask=[[1, 0, 0], [0, 1, 0], [0, 0, 0]])
  465. assert_equal(compress_rowcols(x), [[8]])
  466. assert_equal(compress_rowcols(x, 0), [[6, 7, 8]])
  467. assert_equal(compress_rowcols(x, 1,), [[2], [5], [8]])
  468. x = array(x._data, mask=[[1, 0, 0], [0, 1, 0], [0, 0, 1]])
  469. assert_equal(compress_rowcols(x).size, 0)
  470. assert_equal(compress_rowcols(x, 0).size, 0)
  471. assert_equal(compress_rowcols(x, 1).size, 0)
  472. def test_mask_rowcols(self):
  473. # Tests mask_rowcols.
  474. x = array(np.arange(9).reshape(3, 3),
  475. mask=[[1, 0, 0], [0, 0, 0], [0, 0, 0]])
  476. assert_equal(mask_rowcols(x).mask,
  477. [[1, 1, 1], [1, 0, 0], [1, 0, 0]])
  478. assert_equal(mask_rowcols(x, 0).mask,
  479. [[1, 1, 1], [0, 0, 0], [0, 0, 0]])
  480. assert_equal(mask_rowcols(x, 1).mask,
  481. [[1, 0, 0], [1, 0, 0], [1, 0, 0]])
  482. x = array(x._data, mask=[[0, 0, 0], [0, 1, 0], [0, 0, 0]])
  483. assert_equal(mask_rowcols(x).mask,
  484. [[0, 1, 0], [1, 1, 1], [0, 1, 0]])
  485. assert_equal(mask_rowcols(x, 0).mask,
  486. [[0, 0, 0], [1, 1, 1], [0, 0, 0]])
  487. assert_equal(mask_rowcols(x, 1).mask,
  488. [[0, 1, 0], [0, 1, 0], [0, 1, 0]])
  489. x = array(x._data, mask=[[1, 0, 0], [0, 1, 0], [0, 0, 0]])
  490. assert_equal(mask_rowcols(x).mask,
  491. [[1, 1, 1], [1, 1, 1], [1, 1, 0]])
  492. assert_equal(mask_rowcols(x, 0).mask,
  493. [[1, 1, 1], [1, 1, 1], [0, 0, 0]])
  494. assert_equal(mask_rowcols(x, 1,).mask,
  495. [[1, 1, 0], [1, 1, 0], [1, 1, 0]])
  496. x = array(x._data, mask=[[1, 0, 0], [0, 1, 0], [0, 0, 1]])
  497. assert_(mask_rowcols(x).all() is masked)
  498. assert_(mask_rowcols(x, 0).all() is masked)
  499. assert_(mask_rowcols(x, 1).all() is masked)
  500. assert_(mask_rowcols(x).mask.all())
  501. assert_(mask_rowcols(x, 0).mask.all())
  502. assert_(mask_rowcols(x, 1).mask.all())
  503. @pytest.mark.parametrize("axis", [None, 0, 1])
  504. @pytest.mark.parametrize(["func", "rowcols_axis"],
  505. [(np.ma.mask_rows, 0), (np.ma.mask_cols, 1)])
  506. def test_mask_row_cols_axis_deprecation(self, axis, func, rowcols_axis):
  507. # Test deprecation of the axis argument to `mask_rows` and `mask_cols`
  508. x = array(np.arange(9).reshape(3, 3),
  509. mask=[[1, 0, 0], [0, 0, 0], [0, 0, 0]])
  510. with assert_warns(DeprecationWarning):
  511. res = func(x, axis=axis)
  512. assert_equal(res, mask_rowcols(x, rowcols_axis))
  513. def test_dot(self):
  514. # Tests dot product
  515. n = np.arange(1, 7)
  516. #
  517. m = [1, 0, 0, 0, 0, 0]
  518. a = masked_array(n, mask=m).reshape(2, 3)
  519. b = masked_array(n, mask=m).reshape(3, 2)
  520. c = dot(a, b, strict=True)
  521. assert_equal(c.mask, [[1, 1], [1, 0]])
  522. c = dot(b, a, strict=True)
  523. assert_equal(c.mask, [[1, 1, 1], [1, 0, 0], [1, 0, 0]])
  524. c = dot(a, b, strict=False)
  525. assert_equal(c, np.dot(a.filled(0), b.filled(0)))
  526. c = dot(b, a, strict=False)
  527. assert_equal(c, np.dot(b.filled(0), a.filled(0)))
  528. #
  529. m = [0, 0, 0, 0, 0, 1]
  530. a = masked_array(n, mask=m).reshape(2, 3)
  531. b = masked_array(n, mask=m).reshape(3, 2)
  532. c = dot(a, b, strict=True)
  533. assert_equal(c.mask, [[0, 1], [1, 1]])
  534. c = dot(b, a, strict=True)
  535. assert_equal(c.mask, [[0, 0, 1], [0, 0, 1], [1, 1, 1]])
  536. c = dot(a, b, strict=False)
  537. assert_equal(c, np.dot(a.filled(0), b.filled(0)))
  538. assert_equal(c, dot(a, b))
  539. c = dot(b, a, strict=False)
  540. assert_equal(c, np.dot(b.filled(0), a.filled(0)))
  541. #
  542. m = [0, 0, 0, 0, 0, 0]
  543. a = masked_array(n, mask=m).reshape(2, 3)
  544. b = masked_array(n, mask=m).reshape(3, 2)
  545. c = dot(a, b)
  546. assert_equal(c.mask, nomask)
  547. c = dot(b, a)
  548. assert_equal(c.mask, nomask)
  549. #
  550. a = masked_array(n, mask=[1, 0, 0, 0, 0, 0]).reshape(2, 3)
  551. b = masked_array(n, mask=[0, 0, 0, 0, 0, 0]).reshape(3, 2)
  552. c = dot(a, b, strict=True)
  553. assert_equal(c.mask, [[1, 1], [0, 0]])
  554. c = dot(a, b, strict=False)
  555. assert_equal(c, np.dot(a.filled(0), b.filled(0)))
  556. c = dot(b, a, strict=True)
  557. assert_equal(c.mask, [[1, 0, 0], [1, 0, 0], [1, 0, 0]])
  558. c = dot(b, a, strict=False)
  559. assert_equal(c, np.dot(b.filled(0), a.filled(0)))
  560. #
  561. a = masked_array(n, mask=[0, 0, 0, 0, 0, 1]).reshape(2, 3)
  562. b = masked_array(n, mask=[0, 0, 0, 0, 0, 0]).reshape(3, 2)
  563. c = dot(a, b, strict=True)
  564. assert_equal(c.mask, [[0, 0], [1, 1]])
  565. c = dot(a, b)
  566. assert_equal(c, np.dot(a.filled(0), b.filled(0)))
  567. c = dot(b, a, strict=True)
  568. assert_equal(c.mask, [[0, 0, 1], [0, 0, 1], [0, 0, 1]])
  569. c = dot(b, a, strict=False)
  570. assert_equal(c, np.dot(b.filled(0), a.filled(0)))
  571. #
  572. a = masked_array(n, mask=[0, 0, 0, 0, 0, 1]).reshape(2, 3)
  573. b = masked_array(n, mask=[0, 0, 1, 0, 0, 0]).reshape(3, 2)
  574. c = dot(a, b, strict=True)
  575. assert_equal(c.mask, [[1, 0], [1, 1]])
  576. c = dot(a, b, strict=False)
  577. assert_equal(c, np.dot(a.filled(0), b.filled(0)))
  578. c = dot(b, a, strict=True)
  579. assert_equal(c.mask, [[0, 0, 1], [1, 1, 1], [0, 0, 1]])
  580. c = dot(b, a, strict=False)
  581. assert_equal(c, np.dot(b.filled(0), a.filled(0)))
  582. def test_dot_returns_maskedarray(self):
  583. # See gh-6611
  584. a = np.eye(3)
  585. b = array(a)
  586. assert_(type(dot(a, a)) is MaskedArray)
  587. assert_(type(dot(a, b)) is MaskedArray)
  588. assert_(type(dot(b, a)) is MaskedArray)
  589. assert_(type(dot(b, b)) is MaskedArray)
  590. def test_dot_out(self):
  591. a = array(np.eye(3))
  592. out = array(np.zeros((3, 3)))
  593. res = dot(a, a, out=out)
  594. assert_(res is out)
  595. assert_equal(a, res)
  596. class TestApplyAlongAxis(object):
  597. # Tests 2D functions
  598. def test_3d(self):
  599. a = arange(12.).reshape(2, 2, 3)
  600. def myfunc(b):
  601. return b[1]
  602. xa = apply_along_axis(myfunc, 2, a)
  603. assert_equal(xa, [[1, 4], [7, 10]])
  604. # Tests kwargs functions
  605. def test_3d_kwargs(self):
  606. a = arange(12).reshape(2, 2, 3)
  607. def myfunc(b, offset=0):
  608. return b[1+offset]
  609. xa = apply_along_axis(myfunc, 2, a, offset=1)
  610. assert_equal(xa, [[2, 5], [8, 11]])
  611. class TestApplyOverAxes(object):
  612. # Tests apply_over_axes
  613. def test_basic(self):
  614. a = arange(24).reshape(2, 3, 4)
  615. test = apply_over_axes(np.sum, a, [0, 2])
  616. ctrl = np.array([[[60], [92], [124]]])
  617. assert_equal(test, ctrl)
  618. a[(a % 2).astype(bool)] = masked
  619. test = apply_over_axes(np.sum, a, [0, 2])
  620. ctrl = np.array([[[28], [44], [60]]])
  621. assert_equal(test, ctrl)
  622. class TestMedian(object):
  623. def test_pytype(self):
  624. r = np.ma.median([[np.inf, np.inf], [np.inf, np.inf]], axis=-1)
  625. assert_equal(r, np.inf)
  626. def test_inf(self):
  627. # test that even which computes handles inf / x = masked
  628. r = np.ma.median(np.ma.masked_array([[np.inf, np.inf],
  629. [np.inf, np.inf]]), axis=-1)
  630. assert_equal(r, np.inf)
  631. r = np.ma.median(np.ma.masked_array([[np.inf, np.inf],
  632. [np.inf, np.inf]]), axis=None)
  633. assert_equal(r, np.inf)
  634. # all masked
  635. r = np.ma.median(np.ma.masked_array([[np.inf, np.inf],
  636. [np.inf, np.inf]], mask=True),
  637. axis=-1)
  638. assert_equal(r.mask, True)
  639. r = np.ma.median(np.ma.masked_array([[np.inf, np.inf],
  640. [np.inf, np.inf]], mask=True),
  641. axis=None)
  642. assert_equal(r.mask, True)
  643. def test_non_masked(self):
  644. x = np.arange(9)
  645. assert_equal(np.ma.median(x), 4.)
  646. assert_(type(np.ma.median(x)) is not MaskedArray)
  647. x = range(8)
  648. assert_equal(np.ma.median(x), 3.5)
  649. assert_(type(np.ma.median(x)) is not MaskedArray)
  650. x = 5
  651. assert_equal(np.ma.median(x), 5.)
  652. assert_(type(np.ma.median(x)) is not MaskedArray)
  653. # integer
  654. x = np.arange(9 * 8).reshape(9, 8)
  655. assert_equal(np.ma.median(x, axis=0), np.median(x, axis=0))
  656. assert_equal(np.ma.median(x, axis=1), np.median(x, axis=1))
  657. assert_(np.ma.median(x, axis=1) is not MaskedArray)
  658. # float
  659. x = np.arange(9 * 8.).reshape(9, 8)
  660. assert_equal(np.ma.median(x, axis=0), np.median(x, axis=0))
  661. assert_equal(np.ma.median(x, axis=1), np.median(x, axis=1))
  662. assert_(np.ma.median(x, axis=1) is not MaskedArray)
  663. def test_docstring_examples(self):
  664. "test the examples given in the docstring of ma.median"
  665. x = array(np.arange(8), mask=[0]*4 + [1]*4)
  666. assert_equal(np.ma.median(x), 1.5)
  667. assert_equal(np.ma.median(x).shape, (), "shape mismatch")
  668. assert_(type(np.ma.median(x)) is not MaskedArray)
  669. x = array(np.arange(10).reshape(2, 5), mask=[0]*6 + [1]*4)
  670. assert_equal(np.ma.median(x), 2.5)
  671. assert_equal(np.ma.median(x).shape, (), "shape mismatch")
  672. assert_(type(np.ma.median(x)) is not MaskedArray)
  673. ma_x = np.ma.median(x, axis=-1, overwrite_input=True)
  674. assert_equal(ma_x, [2., 5.])
  675. assert_equal(ma_x.shape, (2,), "shape mismatch")
  676. assert_(type(ma_x) is MaskedArray)
  677. def test_axis_argument_errors(self):
  678. msg = "mask = %s, ndim = %s, axis = %s, overwrite_input = %s"
  679. for ndmin in range(5):
  680. for mask in [False, True]:
  681. x = array(1, ndmin=ndmin, mask=mask)
  682. # Valid axis values should not raise exception
  683. args = itertools.product(range(-ndmin, ndmin), [False, True])
  684. for axis, over in args:
  685. try:
  686. np.ma.median(x, axis=axis, overwrite_input=over)
  687. except Exception:
  688. raise AssertionError(msg % (mask, ndmin, axis, over))
  689. # Invalid axis values should raise exception
  690. args = itertools.product([-(ndmin + 1), ndmin], [False, True])
  691. for axis, over in args:
  692. try:
  693. np.ma.median(x, axis=axis, overwrite_input=over)
  694. except np.AxisError:
  695. pass
  696. else:
  697. raise AssertionError(msg % (mask, ndmin, axis, over))
  698. def test_masked_0d(self):
  699. # Check values
  700. x = array(1, mask=False)
  701. assert_equal(np.ma.median(x), 1)
  702. x = array(1, mask=True)
  703. assert_equal(np.ma.median(x), np.ma.masked)
  704. def test_masked_1d(self):
  705. x = array(np.arange(5), mask=True)
  706. assert_equal(np.ma.median(x), np.ma.masked)
  707. assert_equal(np.ma.median(x).shape, (), "shape mismatch")
  708. assert_(type(np.ma.median(x)) is np.ma.core.MaskedConstant)
  709. x = array(np.arange(5), mask=False)
  710. assert_equal(np.ma.median(x), 2.)
  711. assert_equal(np.ma.median(x).shape, (), "shape mismatch")
  712. assert_(type(np.ma.median(x)) is not MaskedArray)
  713. x = array(np.arange(5), mask=[0,1,0,0,0])
  714. assert_equal(np.ma.median(x), 2.5)
  715. assert_equal(np.ma.median(x).shape, (), "shape mismatch")
  716. assert_(type(np.ma.median(x)) is not MaskedArray)
  717. x = array(np.arange(5), mask=[0,1,1,1,1])
  718. assert_equal(np.ma.median(x), 0.)
  719. assert_equal(np.ma.median(x).shape, (), "shape mismatch")
  720. assert_(type(np.ma.median(x)) is not MaskedArray)
  721. # integer
  722. x = array(np.arange(5), mask=[0,1,1,0,0])
  723. assert_equal(np.ma.median(x), 3.)
  724. assert_equal(np.ma.median(x).shape, (), "shape mismatch")
  725. assert_(type(np.ma.median(x)) is not MaskedArray)
  726. # float
  727. x = array(np.arange(5.), mask=[0,1,1,0,0])
  728. assert_equal(np.ma.median(x), 3.)
  729. assert_equal(np.ma.median(x).shape, (), "shape mismatch")
  730. assert_(type(np.ma.median(x)) is not MaskedArray)
  731. # integer
  732. x = array(np.arange(6), mask=[0,1,1,1,1,0])
  733. assert_equal(np.ma.median(x), 2.5)
  734. assert_equal(np.ma.median(x).shape, (), "shape mismatch")
  735. assert_(type(np.ma.median(x)) is not MaskedArray)
  736. # float
  737. x = array(np.arange(6.), mask=[0,1,1,1,1,0])
  738. assert_equal(np.ma.median(x), 2.5)
  739. assert_equal(np.ma.median(x).shape, (), "shape mismatch")
  740. assert_(type(np.ma.median(x)) is not MaskedArray)
  741. def test_1d_shape_consistency(self):
  742. assert_equal(np.ma.median(array([1,2,3],mask=[0,0,0])).shape,
  743. np.ma.median(array([1,2,3],mask=[0,1,0])).shape )
  744. def test_2d(self):
  745. # Tests median w/ 2D
  746. (n, p) = (101, 30)
  747. x = masked_array(np.linspace(-1., 1., n),)
  748. x[:10] = x[-10:] = masked
  749. z = masked_array(np.empty((n, p), dtype=float))
  750. z[:, 0] = x[:]
  751. idx = np.arange(len(x))
  752. for i in range(1, p):
  753. np.random.shuffle(idx)
  754. z[:, i] = x[idx]
  755. assert_equal(median(z[:, 0]), 0)
  756. assert_equal(median(z), 0)
  757. assert_equal(median(z, axis=0), np.zeros(p))
  758. assert_equal(median(z.T, axis=1), np.zeros(p))
  759. def test_2d_waxis(self):
  760. # Tests median w/ 2D arrays and different axis.
  761. x = masked_array(np.arange(30).reshape(10, 3))
  762. x[:3] = x[-3:] = masked
  763. assert_equal(median(x), 14.5)
  764. assert_(type(np.ma.median(x)) is not MaskedArray)
  765. assert_equal(median(x, axis=0), [13.5, 14.5, 15.5])
  766. assert_(type(np.ma.median(x, axis=0)) is MaskedArray)
  767. assert_equal(median(x, axis=1), [0, 0, 0, 10, 13, 16, 19, 0, 0, 0])
  768. assert_(type(np.ma.median(x, axis=1)) is MaskedArray)
  769. assert_equal(median(x, axis=1).mask, [1, 1, 1, 0, 0, 0, 0, 1, 1, 1])
  770. def test_3d(self):
  771. # Tests median w/ 3D
  772. x = np.ma.arange(24).reshape(3, 4, 2)
  773. x[x % 3 == 0] = masked
  774. assert_equal(median(x, 0), [[12, 9], [6, 15], [12, 9], [18, 15]])
  775. x.shape = (4, 3, 2)
  776. assert_equal(median(x, 0), [[99, 10], [11, 99], [13, 14]])
  777. x = np.ma.arange(24).reshape(4, 3, 2)
  778. x[x % 5 == 0] = masked
  779. assert_equal(median(x, 0), [[12, 10], [8, 9], [16, 17]])
  780. def test_neg_axis(self):
  781. x = masked_array(np.arange(30).reshape(10, 3))
  782. x[:3] = x[-3:] = masked
  783. assert_equal(median(x, axis=-1), median(x, axis=1))
  784. def test_out_1d(self):
  785. # integer float even odd
  786. for v in (30, 30., 31, 31.):
  787. x = masked_array(np.arange(v))
  788. x[:3] = x[-3:] = masked
  789. out = masked_array(np.ones(()))
  790. r = median(x, out=out)
  791. if v == 30:
  792. assert_equal(out, 14.5)
  793. else:
  794. assert_equal(out, 15.)
  795. assert_(r is out)
  796. assert_(type(r) is MaskedArray)
  797. def test_out(self):
  798. # integer float even odd
  799. for v in (40, 40., 30, 30.):
  800. x = masked_array(np.arange(v).reshape(10, -1))
  801. x[:3] = x[-3:] = masked
  802. out = masked_array(np.ones(10))
  803. r = median(x, axis=1, out=out)
  804. if v == 30:
  805. e = masked_array([0.]*3 + [10, 13, 16, 19] + [0.]*3,
  806. mask=[True] * 3 + [False] * 4 + [True] * 3)
  807. else:
  808. e = masked_array([0.]*3 + [13.5, 17.5, 21.5, 25.5] + [0.]*3,
  809. mask=[True]*3 + [False]*4 + [True]*3)
  810. assert_equal(r, e)
  811. assert_(r is out)
  812. assert_(type(r) is MaskedArray)
  813. def test_single_non_masked_value_on_axis(self):
  814. data = [[1., 0.],
  815. [0., 3.],
  816. [0., 0.]]
  817. masked_arr = np.ma.masked_equal(data, 0)
  818. expected = [1., 3.]
  819. assert_array_equal(np.ma.median(masked_arr, axis=0),
  820. expected)
  821. def test_nan(self):
  822. for mask in (False, np.zeros(6, dtype=bool)):
  823. dm = np.ma.array([[1, np.nan, 3], [1, 2, 3]])
  824. dm.mask = mask
  825. # scalar result
  826. r = np.ma.median(dm, axis=None)
  827. assert_(np.isscalar(r))
  828. assert_array_equal(r, np.nan)
  829. r = np.ma.median(dm.ravel(), axis=0)
  830. assert_(np.isscalar(r))
  831. assert_array_equal(r, np.nan)
  832. r = np.ma.median(dm, axis=0)
  833. assert_equal(type(r), MaskedArray)
  834. assert_array_equal(r, [1, np.nan, 3])
  835. r = np.ma.median(dm, axis=1)
  836. assert_equal(type(r), MaskedArray)
  837. assert_array_equal(r, [np.nan, 2])
  838. r = np.ma.median(dm, axis=-1)
  839. assert_equal(type(r), MaskedArray)
  840. assert_array_equal(r, [np.nan, 2])
  841. dm = np.ma.array([[1, np.nan, 3], [1, 2, 3]])
  842. dm[:, 2] = np.ma.masked
  843. assert_array_equal(np.ma.median(dm, axis=None), np.nan)
  844. assert_array_equal(np.ma.median(dm, axis=0), [1, np.nan, 3])
  845. assert_array_equal(np.ma.median(dm, axis=1), [np.nan, 1.5])
  846. def test_out_nan(self):
  847. o = np.ma.masked_array(np.zeros((4,)))
  848. d = np.ma.masked_array(np.ones((3, 4)))
  849. d[2, 1] = np.nan
  850. d[2, 2] = np.ma.masked
  851. assert_equal(np.ma.median(d, 0, out=o), o)
  852. o = np.ma.masked_array(np.zeros((3,)))
  853. assert_equal(np.ma.median(d, 1, out=o), o)
  854. o = np.ma.masked_array(np.zeros(()))
  855. assert_equal(np.ma.median(d, out=o), o)
  856. def test_nan_behavior(self):
  857. a = np.ma.masked_array(np.arange(24, dtype=float))
  858. a[::3] = np.ma.masked
  859. a[2] = np.nan
  860. assert_array_equal(np.ma.median(a), np.nan)
  861. assert_array_equal(np.ma.median(a, axis=0), np.nan)
  862. a = np.ma.masked_array(np.arange(24, dtype=float).reshape(2, 3, 4))
  863. a.mask = np.arange(a.size) % 2 == 1
  864. aorig = a.copy()
  865. a[1, 2, 3] = np.nan
  866. a[1, 1, 2] = np.nan
  867. # no axis
  868. assert_array_equal(np.ma.median(a), np.nan)
  869. assert_(np.isscalar(np.ma.median(a)))
  870. # axis0
  871. b = np.ma.median(aorig, axis=0)
  872. b[2, 3] = np.nan
  873. b[1, 2] = np.nan
  874. assert_equal(np.ma.median(a, 0), b)
  875. # axis1
  876. b = np.ma.median(aorig, axis=1)
  877. b[1, 3] = np.nan
  878. b[1, 2] = np.nan
  879. assert_equal(np.ma.median(a, 1), b)
  880. # axis02
  881. b = np.ma.median(aorig, axis=(0, 2))
  882. b[1] = np.nan
  883. b[2] = np.nan
  884. assert_equal(np.ma.median(a, (0, 2)), b)
  885. def test_ambigous_fill(self):
  886. # 255 is max value, used as filler for sort
  887. a = np.array([[3, 3, 255], [3, 3, 255]], dtype=np.uint8)
  888. a = np.ma.masked_array(a, mask=a == 3)
  889. assert_array_equal(np.ma.median(a, axis=1), 255)
  890. assert_array_equal(np.ma.median(a, axis=1).mask, False)
  891. assert_array_equal(np.ma.median(a, axis=0), a[0])
  892. assert_array_equal(np.ma.median(a), 255)
  893. def test_special(self):
  894. for inf in [np.inf, -np.inf]:
  895. a = np.array([[inf, np.nan], [np.nan, np.nan]])
  896. a = np.ma.masked_array(a, mask=np.isnan(a))
  897. assert_equal(np.ma.median(a, axis=0), [inf, np.nan])
  898. assert_equal(np.ma.median(a, axis=1), [inf, np.nan])
  899. assert_equal(np.ma.median(a), inf)
  900. a = np.array([[np.nan, np.nan, inf], [np.nan, np.nan, inf]])
  901. a = np.ma.masked_array(a, mask=np.isnan(a))
  902. assert_array_equal(np.ma.median(a, axis=1), inf)
  903. assert_array_equal(np.ma.median(a, axis=1).mask, False)
  904. assert_array_equal(np.ma.median(a, axis=0), a[0])
  905. assert_array_equal(np.ma.median(a), inf)
  906. # no mask
  907. a = np.array([[inf, inf], [inf, inf]])
  908. assert_equal(np.ma.median(a), inf)
  909. assert_equal(np.ma.median(a, axis=0), inf)
  910. assert_equal(np.ma.median(a, axis=1), inf)
  911. a = np.array([[inf, 7, -inf, -9],
  912. [-10, np.nan, np.nan, 5],
  913. [4, np.nan, np.nan, inf]],
  914. dtype=np.float32)
  915. a = np.ma.masked_array(a, mask=np.isnan(a))
  916. if inf > 0:
  917. assert_equal(np.ma.median(a, axis=0), [4., 7., -inf, 5.])
  918. assert_equal(np.ma.median(a), 4.5)
  919. else:
  920. assert_equal(np.ma.median(a, axis=0), [-10., 7., -inf, -9.])
  921. assert_equal(np.ma.median(a), -2.5)
  922. assert_equal(np.ma.median(a, axis=1), [-1., -2.5, inf])
  923. for i in range(0, 10):
  924. for j in range(1, 10):
  925. a = np.array([([np.nan] * i) + ([inf] * j)] * 2)
  926. a = np.ma.masked_array(a, mask=np.isnan(a))
  927. assert_equal(np.ma.median(a), inf)
  928. assert_equal(np.ma.median(a, axis=1), inf)
  929. assert_equal(np.ma.median(a, axis=0),
  930. ([np.nan] * i) + [inf] * j)
  931. def test_empty(self):
  932. # empty arrays
  933. a = np.ma.masked_array(np.array([], dtype=float))
  934. with suppress_warnings() as w:
  935. w.record(RuntimeWarning)
  936. assert_array_equal(np.ma.median(a), np.nan)
  937. assert_(w.log[0].category is RuntimeWarning)
  938. # multiple dimensions
  939. a = np.ma.masked_array(np.array([], dtype=float, ndmin=3))
  940. # no axis
  941. with suppress_warnings() as w:
  942. w.record(RuntimeWarning)
  943. warnings.filterwarnings('always', '', RuntimeWarning)
  944. assert_array_equal(np.ma.median(a), np.nan)
  945. assert_(w.log[0].category is RuntimeWarning)
  946. # axis 0 and 1
  947. b = np.ma.masked_array(np.array([], dtype=float, ndmin=2))
  948. assert_equal(np.ma.median(a, axis=0), b)
  949. assert_equal(np.ma.median(a, axis=1), b)
  950. # axis 2
  951. b = np.ma.masked_array(np.array(np.nan, dtype=float, ndmin=2))
  952. with warnings.catch_warnings(record=True) as w:
  953. warnings.filterwarnings('always', '', RuntimeWarning)
  954. assert_equal(np.ma.median(a, axis=2), b)
  955. assert_(w[0].category is RuntimeWarning)
  956. def test_object(self):
  957. o = np.ma.masked_array(np.arange(7.))
  958. assert_(type(np.ma.median(o.astype(object))), float)
  959. o[2] = np.nan
  960. assert_(type(np.ma.median(o.astype(object))), float)
  961. class TestCov(object):
  962. def setup(self):
  963. self.data = array(np.random.rand(12))
  964. def test_1d_without_missing(self):
  965. # Test cov on 1D variable w/o missing values
  966. x = self.data
  967. assert_almost_equal(np.cov(x), cov(x))
  968. assert_almost_equal(np.cov(x, rowvar=False), cov(x, rowvar=False))
  969. assert_almost_equal(np.cov(x, rowvar=False, bias=True),
  970. cov(x, rowvar=False, bias=True))
  971. def test_2d_without_missing(self):
  972. # Test cov on 1 2D variable w/o missing values
  973. x = self.data.reshape(3, 4)
  974. assert_almost_equal(np.cov(x), cov(x))
  975. assert_almost_equal(np.cov(x, rowvar=False), cov(x, rowvar=False))
  976. assert_almost_equal(np.cov(x, rowvar=False, bias=True),
  977. cov(x, rowvar=False, bias=True))
  978. def test_1d_with_missing(self):
  979. # Test cov 1 1D variable w/missing values
  980. x = self.data
  981. x[-1] = masked
  982. x -= x.mean()
  983. nx = x.compressed()
  984. assert_almost_equal(np.cov(nx), cov(x))
  985. assert_almost_equal(np.cov(nx, rowvar=False), cov(x, rowvar=False))
  986. assert_almost_equal(np.cov(nx, rowvar=False, bias=True),
  987. cov(x, rowvar=False, bias=True))
  988. #
  989. try:
  990. cov(x, allow_masked=False)
  991. except ValueError:
  992. pass
  993. #
  994. # 2 1D variables w/ missing values
  995. nx = x[1:-1]
  996. assert_almost_equal(np.cov(nx, nx[::-1]), cov(x, x[::-1]))
  997. assert_almost_equal(np.cov(nx, nx[::-1], rowvar=False),
  998. cov(x, x[::-1], rowvar=False))
  999. assert_almost_equal(np.cov(nx, nx[::-1], rowvar=False, bias=True),
  1000. cov(x, x[::-1], rowvar=False, bias=True))
  1001. def test_2d_with_missing(self):
  1002. # Test cov on 2D variable w/ missing value
  1003. x = self.data
  1004. x[-1] = masked
  1005. x = x.reshape(3, 4)
  1006. valid = np.logical_not(getmaskarray(x)).astype(int)
  1007. frac = np.dot(valid, valid.T)
  1008. xf = (x - x.mean(1)[:, None]).filled(0)
  1009. assert_almost_equal(cov(x),
  1010. np.cov(xf) * (x.shape[1] - 1) / (frac - 1.))
  1011. assert_almost_equal(cov(x, bias=True),
  1012. np.cov(xf, bias=True) * x.shape[1] / frac)
  1013. frac = np.dot(valid.T, valid)
  1014. xf = (x - x.mean(0)).filled(0)
  1015. assert_almost_equal(cov(x, rowvar=False),
  1016. (np.cov(xf, rowvar=False) *
  1017. (x.shape[0] - 1) / (frac - 1.)))
  1018. assert_almost_equal(cov(x, rowvar=False, bias=True),
  1019. (np.cov(xf, rowvar=False, bias=True) *
  1020. x.shape[0] / frac))
  1021. class TestCorrcoef(object):
  1022. def setup(self):
  1023. self.data = array(np.random.rand(12))
  1024. self.data2 = array(np.random.rand(12))
  1025. def test_ddof(self):
  1026. # ddof raises DeprecationWarning
  1027. x, y = self.data, self.data2
  1028. expected = np.corrcoef(x)
  1029. expected2 = np.corrcoef(x, y)
  1030. with suppress_warnings() as sup:
  1031. warnings.simplefilter("always")
  1032. assert_warns(DeprecationWarning, corrcoef, x, ddof=-1)
  1033. sup.filter(DeprecationWarning, "bias and ddof have no effect")
  1034. # ddof has no or negligible effect on the function
  1035. assert_almost_equal(np.corrcoef(x, ddof=0), corrcoef(x, ddof=0))
  1036. assert_almost_equal(corrcoef(x, ddof=-1), expected)
  1037. assert_almost_equal(corrcoef(x, y, ddof=-1), expected2)
  1038. assert_almost_equal(corrcoef(x, ddof=3), expected)
  1039. assert_almost_equal(corrcoef(x, y, ddof=3), expected2)
  1040. def test_bias(self):
  1041. x, y = self.data, self.data2
  1042. expected = np.corrcoef(x)
  1043. # bias raises DeprecationWarning
  1044. with suppress_warnings() as sup:
  1045. warnings.simplefilter("always")
  1046. assert_warns(DeprecationWarning, corrcoef, x, y, True, False)
  1047. assert_warns(DeprecationWarning, corrcoef, x, y, True, True)
  1048. assert_warns(DeprecationWarning, corrcoef, x, bias=False)
  1049. sup.filter(DeprecationWarning, "bias and ddof have no effect")
  1050. # bias has no or negligible effect on the function
  1051. assert_almost_equal(corrcoef(x, bias=1), expected)
  1052. def test_1d_without_missing(self):
  1053. # Test cov on 1D variable w/o missing values
  1054. x = self.data
  1055. assert_almost_equal(np.corrcoef(x), corrcoef(x))
  1056. assert_almost_equal(np.corrcoef(x, rowvar=False),
  1057. corrcoef(x, rowvar=False))
  1058. with suppress_warnings() as sup:
  1059. sup.filter(DeprecationWarning, "bias and ddof have no effect")
  1060. assert_almost_equal(np.corrcoef(x, rowvar=False, bias=True),
  1061. corrcoef(x, rowvar=False, bias=True))
  1062. def test_2d_without_missing(self):
  1063. # Test corrcoef on 1 2D variable w/o missing values
  1064. x = self.data.reshape(3, 4)
  1065. assert_almost_equal(np.corrcoef(x), corrcoef(x))
  1066. assert_almost_equal(np.corrcoef(x, rowvar=False),
  1067. corrcoef(x, rowvar=False))
  1068. with suppress_warnings() as sup:
  1069. sup.filter(DeprecationWarning, "bias and ddof have no effect")
  1070. assert_almost_equal(np.corrcoef(x, rowvar=False, bias=True),
  1071. corrcoef(x, rowvar=False, bias=True))
  1072. def test_1d_with_missing(self):
  1073. # Test corrcoef 1 1D variable w/missing values
  1074. x = self.data
  1075. x[-1] = masked
  1076. x -= x.mean()
  1077. nx = x.compressed()
  1078. assert_almost_equal(np.corrcoef(nx), corrcoef(x))
  1079. assert_almost_equal(np.corrcoef(nx, rowvar=False),
  1080. corrcoef(x, rowvar=False))
  1081. with suppress_warnings() as sup:
  1082. sup.filter(DeprecationWarning, "bias and ddof have no effect")
  1083. assert_almost_equal(np.corrcoef(nx, rowvar=False, bias=True),
  1084. corrcoef(x, rowvar=False, bias=True))
  1085. try:
  1086. corrcoef(x, allow_masked=False)
  1087. except ValueError:
  1088. pass
  1089. # 2 1D variables w/ missing values
  1090. nx = x[1:-1]
  1091. assert_almost_equal(np.corrcoef(nx, nx[::-1]), corrcoef(x, x[::-1]))
  1092. assert_almost_equal(np.corrcoef(nx, nx[::-1], rowvar=False),
  1093. corrcoef(x, x[::-1], rowvar=False))
  1094. with suppress_warnings() as sup:
  1095. sup.filter(DeprecationWarning, "bias and ddof have no effect")
  1096. # ddof and bias have no or negligible effect on the function
  1097. assert_almost_equal(np.corrcoef(nx, nx[::-1]),
  1098. corrcoef(x, x[::-1], bias=1))
  1099. assert_almost_equal(np.corrcoef(nx, nx[::-1]),
  1100. corrcoef(x, x[::-1], ddof=2))
  1101. def test_2d_with_missing(self):
  1102. # Test corrcoef on 2D variable w/ missing value
  1103. x = self.data
  1104. x[-1] = masked
  1105. x = x.reshape(3, 4)
  1106. test = corrcoef(x)
  1107. control = np.corrcoef(x)
  1108. assert_almost_equal(test[:-1, :-1], control[:-1, :-1])
  1109. with suppress_warnings() as sup:
  1110. sup.filter(DeprecationWarning, "bias and ddof have no effect")
  1111. # ddof and bias have no or negligible effect on the function
  1112. assert_almost_equal(corrcoef(x, ddof=-2)[:-1, :-1],
  1113. control[:-1, :-1])
  1114. assert_almost_equal(corrcoef(x, ddof=3)[:-1, :-1],
  1115. control[:-1, :-1])
  1116. assert_almost_equal(corrcoef(x, bias=1)[:-1, :-1],
  1117. control[:-1, :-1])
  1118. class TestPolynomial(object):
  1119. #
  1120. def test_polyfit(self):
  1121. # Tests polyfit
  1122. # On ndarrays
  1123. x = np.random.rand(10)
  1124. y = np.random.rand(20).reshape(-1, 2)
  1125. assert_almost_equal(polyfit(x, y, 3), np.polyfit(x, y, 3))
  1126. # ON 1D maskedarrays
  1127. x = x.view(MaskedArray)
  1128. x[0] = masked
  1129. y = y.view(MaskedArray)
  1130. y[0, 0] = y[-1, -1] = masked
  1131. #
  1132. (C, R, K, S, D) = polyfit(x, y[:, 0], 3, full=True)
  1133. (c, r, k, s, d) = np.polyfit(x[1:], y[1:, 0].compressed(), 3,
  1134. full=True)
  1135. for (a, a_) in zip((C, R, K, S, D), (c, r, k, s, d)):
  1136. assert_almost_equal(a, a_)
  1137. #
  1138. (C, R, K, S, D) = polyfit(x, y[:, -1], 3, full=True)
  1139. (c, r, k, s, d) = np.polyfit(x[1:-1], y[1:-1, -1], 3, full=True)
  1140. for (a, a_) in zip((C, R, K, S, D), (c, r, k, s, d)):
  1141. assert_almost_equal(a, a_)
  1142. #
  1143. (C, R, K, S, D) = polyfit(x, y, 3, full=True)
  1144. (c, r, k, s, d) = np.polyfit(x[1:-1], y[1:-1,:], 3, full=True)
  1145. for (a, a_) in zip((C, R, K, S, D), (c, r, k, s, d)):
  1146. assert_almost_equal(a, a_)
  1147. #
  1148. w = np.random.rand(10) + 1
  1149. wo = w.copy()
  1150. xs = x[1:-1]
  1151. ys = y[1:-1]
  1152. ws = w[1:-1]
  1153. (C, R, K, S, D) = polyfit(x, y, 3, full=True, w=w)
  1154. (c, r, k, s, d) = np.polyfit(xs, ys, 3, full=True, w=ws)
  1155. assert_equal(w, wo)
  1156. for (a, a_) in zip((C, R, K, S, D), (c, r, k, s, d)):
  1157. assert_almost_equal(a, a_)
  1158. def test_polyfit_with_masked_NaNs(self):
  1159. x = np.random.rand(10)
  1160. y = np.random.rand(20).reshape(-1, 2)
  1161. x[0] = np.nan
  1162. y[-1,-1] = np.nan
  1163. x = x.view(MaskedArray)
  1164. y = y.view(MaskedArray)
  1165. x[0] = masked
  1166. y[-1,-1] = masked
  1167. (C, R, K, S, D) = polyfit(x, y, 3, full=True)
  1168. (c, r, k, s, d) = np.polyfit(x[1:-1], y[1:-1,:], 3, full=True)
  1169. for (a, a_) in zip((C, R, K, S, D), (c, r, k, s, d)):
  1170. assert_almost_equal(a, a_)
  1171. class TestArraySetOps(object):
  1172. def test_unique_onlist(self):
  1173. # Test unique on list
  1174. data = [1, 1, 1, 2, 2, 3]
  1175. test = unique(data, return_index=True, return_inverse=True)
  1176. assert_(isinstance(test[0], MaskedArray))
  1177. assert_equal(test[0], masked_array([1, 2, 3], mask=[0, 0, 0]))
  1178. assert_equal(test[1], [0, 3, 5])
  1179. assert_equal(test[2], [0, 0, 0, 1, 1, 2])
  1180. def test_unique_onmaskedarray(self):
  1181. # Test unique on masked data w/use_mask=True
  1182. data = masked_array([1, 1, 1, 2, 2, 3], mask=[0, 0, 1, 0, 1, 0])
  1183. test = unique(data, return_index=True, return_inverse=True)
  1184. assert_equal(test[0], masked_array([1, 2, 3, -1], mask=[0, 0, 0, 1]))
  1185. assert_equal(test[1], [0, 3, 5, 2])
  1186. assert_equal(test[2], [0, 0, 3, 1, 3, 2])
  1187. #
  1188. data.fill_value = 3
  1189. data = masked_array(data=[1, 1, 1, 2, 2, 3],
  1190. mask=[0, 0, 1, 0, 1, 0], fill_value=3)
  1191. test = unique(data, return_index=True, return_inverse=True)
  1192. assert_equal(test[0], masked_array([1, 2, 3, -1], mask=[0, 0, 0, 1]))
  1193. assert_equal(test[1], [0, 3, 5, 2])
  1194. assert_equal(test[2], [0, 0, 3, 1, 3, 2])
  1195. def test_unique_allmasked(self):
  1196. # Test all masked
  1197. data = masked_array([1, 1, 1], mask=True)
  1198. test = unique(data, return_index=True, return_inverse=True)
  1199. assert_equal(test[0], masked_array([1, ], mask=[True]))
  1200. assert_equal(test[1], [0])
  1201. assert_equal(test[2], [0, 0, 0])
  1202. #
  1203. # Test masked
  1204. data = masked
  1205. test = unique(data, return_index=True, return_inverse=True)
  1206. assert_equal(test[0], masked_array(masked))
  1207. assert_equal(test[1], [0])
  1208. assert_equal(test[2], [0])
  1209. def test_ediff1d(self):
  1210. # Tests mediff1d
  1211. x = masked_array(np.arange(5), mask=[1, 0, 0, 0, 1])
  1212. control = array([1, 1, 1, 4], mask=[1, 0, 0, 1])
  1213. test = ediff1d(x)
  1214. assert_equal(test, control)
  1215. assert_equal(test.filled(0), control.filled(0))
  1216. assert_equal(test.mask, control.mask)
  1217. def test_ediff1d_tobegin(self):
  1218. # Test ediff1d w/ to_begin
  1219. x = masked_array(np.arange(5), mask=[1, 0, 0, 0, 1])
  1220. test = ediff1d(x, to_begin=masked)
  1221. control = array([0, 1, 1, 1, 4], mask=[1, 1, 0, 0, 1])
  1222. assert_equal(test, control)
  1223. assert_equal(test.filled(0), control.filled(0))
  1224. assert_equal(test.mask, control.mask)
  1225. #
  1226. test = ediff1d(x, to_begin=[1, 2, 3])
  1227. control = array([1, 2, 3, 1, 1, 1, 4], mask=[0, 0, 0, 1, 0, 0, 1])
  1228. assert_equal(test, control)
  1229. assert_equal(test.filled(0), control.filled(0))
  1230. assert_equal(test.mask, control.mask)
  1231. def test_ediff1d_toend(self):
  1232. # Test ediff1d w/ to_end
  1233. x = masked_array(np.arange(5), mask=[1, 0, 0, 0, 1])
  1234. test = ediff1d(x, to_end=masked)
  1235. control = array([1, 1, 1, 4, 0], mask=[1, 0, 0, 1, 1])
  1236. assert_equal(test, control)
  1237. assert_equal(test.filled(0), control.filled(0))
  1238. assert_equal(test.mask, control.mask)
  1239. #
  1240. test = ediff1d(x, to_end=[1, 2, 3])
  1241. control = array([1, 1, 1, 4, 1, 2, 3], mask=[1, 0, 0, 1, 0, 0, 0])
  1242. assert_equal(test, control)
  1243. assert_equal(test.filled(0), control.filled(0))
  1244. assert_equal(test.mask, control.mask)
  1245. def test_ediff1d_tobegin_toend(self):
  1246. # Test ediff1d w/ to_begin and to_end
  1247. x = masked_array(np.arange(5), mask=[1, 0, 0, 0, 1])
  1248. test = ediff1d(x, to_end=masked, to_begin=masked)
  1249. control = array([0, 1, 1, 1, 4, 0], mask=[1, 1, 0, 0, 1, 1])
  1250. assert_equal(test, control)
  1251. assert_equal(test.filled(0), control.filled(0))
  1252. assert_equal(test.mask, control.mask)
  1253. #
  1254. test = ediff1d(x, to_end=[1, 2, 3], to_begin=masked)
  1255. control = array([0, 1, 1, 1, 4, 1, 2, 3],
  1256. mask=[1, 1, 0, 0, 1, 0, 0, 0])
  1257. assert_equal(test, control)
  1258. assert_equal(test.filled(0), control.filled(0))
  1259. assert_equal(test.mask, control.mask)
  1260. def test_ediff1d_ndarray(self):
  1261. # Test ediff1d w/ a ndarray
  1262. x = np.arange(5)
  1263. test = ediff1d(x)
  1264. control = array([1, 1, 1, 1], mask=[0, 0, 0, 0])
  1265. assert_equal(test, control)
  1266. assert_(isinstance(test, MaskedArray))
  1267. assert_equal(test.filled(0), control.filled(0))
  1268. assert_equal(test.mask, control.mask)
  1269. #
  1270. test = ediff1d(x, to_end=masked, to_begin=masked)
  1271. control = array([0, 1, 1, 1, 1, 0], mask=[1, 0, 0, 0, 0, 1])
  1272. assert_(isinstance(test, MaskedArray))
  1273. assert_equal(test.filled(0), control.filled(0))
  1274. assert_equal(test.mask, control.mask)
  1275. def test_intersect1d(self):
  1276. # Test intersect1d
  1277. x = array([1, 3, 3, 3], mask=[0, 0, 0, 1])
  1278. y = array([3, 1, 1, 1], mask=[0, 0, 0, 1])
  1279. test = intersect1d(x, y)
  1280. control = array([1, 3, -1], mask=[0, 0, 1])
  1281. assert_equal(test, control)
  1282. def test_setxor1d(self):
  1283. # Test setxor1d
  1284. a = array([1, 2, 5, 7, -1], mask=[0, 0, 0, 0, 1])
  1285. b = array([1, 2, 3, 4, 5, -1], mask=[0, 0, 0, 0, 0, 1])
  1286. test = setxor1d(a, b)
  1287. assert_equal(test, array([3, 4, 7]))
  1288. #
  1289. a = array([1, 2, 5, 7, -1], mask=[0, 0, 0, 0, 1])
  1290. b = [1, 2, 3, 4, 5]
  1291. test = setxor1d(a, b)
  1292. assert_equal(test, array([3, 4, 7, -1], mask=[0, 0, 0, 1]))
  1293. #
  1294. a = array([1, 2, 3])
  1295. b = array([6, 5, 4])
  1296. test = setxor1d(a, b)
  1297. assert_(isinstance(test, MaskedArray))
  1298. assert_equal(test, [1, 2, 3, 4, 5, 6])
  1299. #
  1300. a = array([1, 8, 2, 3], mask=[0, 1, 0, 0])
  1301. b = array([6, 5, 4, 8], mask=[0, 0, 0, 1])
  1302. test = setxor1d(a, b)
  1303. assert_(isinstance(test, MaskedArray))
  1304. assert_equal(test, [1, 2, 3, 4, 5, 6])
  1305. #
  1306. assert_array_equal([], setxor1d([], []))
  1307. def test_isin(self):
  1308. # the tests for in1d cover most of isin's behavior
  1309. # if in1d is removed, would need to change those tests to test
  1310. # isin instead.
  1311. a = np.arange(24).reshape([2, 3, 4])
  1312. mask = np.zeros([2, 3, 4])
  1313. mask[1, 2, 0] = 1
  1314. a = array(a, mask=mask)
  1315. b = array(data=[0, 10, 20, 30, 1, 3, 11, 22, 33],
  1316. mask=[0, 1, 0, 1, 0, 1, 0, 1, 0])
  1317. ec = zeros((2, 3, 4), dtype=bool)
  1318. ec[0, 0, 0] = True
  1319. ec[0, 0, 1] = True
  1320. ec[0, 2, 3] = True
  1321. c = isin(a, b)
  1322. assert_(isinstance(c, MaskedArray))
  1323. assert_array_equal(c, ec)
  1324. #compare results of np.isin to ma.isin
  1325. d = np.isin(a, b[~b.mask]) & ~a.mask
  1326. assert_array_equal(c, d)
  1327. def test_in1d(self):
  1328. # Test in1d
  1329. a = array([1, 2, 5, 7, -1], mask=[0, 0, 0, 0, 1])
  1330. b = array([1, 2, 3, 4, 5, -1], mask=[0, 0, 0, 0, 0, 1])
  1331. test = in1d(a, b)
  1332. assert_equal(test, [True, True, True, False, True])
  1333. #
  1334. a = array([5, 5, 2, 1, -1], mask=[0, 0, 0, 0, 1])
  1335. b = array([1, 5, -1], mask=[0, 0, 1])
  1336. test = in1d(a, b)
  1337. assert_equal(test, [True, True, False, True, True])
  1338. #
  1339. assert_array_equal([], in1d([], []))
  1340. def test_in1d_invert(self):
  1341. # Test in1d's invert parameter
  1342. a = array([1, 2, 5, 7, -1], mask=[0, 0, 0, 0, 1])
  1343. b = array([1, 2, 3, 4, 5, -1], mask=[0, 0, 0, 0, 0, 1])
  1344. assert_equal(np.invert(in1d(a, b)), in1d(a, b, invert=True))
  1345. a = array([5, 5, 2, 1, -1], mask=[0, 0, 0, 0, 1])
  1346. b = array([1, 5, -1], mask=[0, 0, 1])
  1347. assert_equal(np.invert(in1d(a, b)), in1d(a, b, invert=True))
  1348. assert_array_equal([], in1d([], [], invert=True))
  1349. def test_union1d(self):
  1350. # Test union1d
  1351. a = array([1, 2, 5, 7, 5, -1], mask=[0, 0, 0, 0, 0, 1])
  1352. b = array([1, 2, 3, 4, 5, -1], mask=[0, 0, 0, 0, 0, 1])
  1353. test = union1d(a, b)
  1354. control = array([1, 2, 3, 4, 5, 7, -1], mask=[0, 0, 0, 0, 0, 0, 1])
  1355. assert_equal(test, control)
  1356. # Tests gh-10340, arguments to union1d should be
  1357. # flattened if they are not already 1D
  1358. x = array([[0, 1, 2], [3, 4, 5]], mask=[[0, 0, 0], [0, 0, 1]])
  1359. y = array([0, 1, 2, 3, 4], mask=[0, 0, 0, 0, 1])
  1360. ez = array([0, 1, 2, 3, 4, 5], mask=[0, 0, 0, 0, 0, 1])
  1361. z = union1d(x, y)
  1362. assert_equal(z, ez)
  1363. #
  1364. assert_array_equal([], union1d([], []))
  1365. def test_setdiff1d(self):
  1366. # Test setdiff1d
  1367. a = array([6, 5, 4, 7, 7, 1, 2, 1], mask=[0, 0, 0, 0, 0, 0, 0, 1])
  1368. b = array([2, 4, 3, 3, 2, 1, 5])
  1369. test = setdiff1d(a, b)
  1370. assert_equal(test, array([6, 7, -1], mask=[0, 0, 1]))
  1371. #
  1372. a = arange(10)
  1373. b = arange(8)
  1374. assert_equal(setdiff1d(a, b), array([8, 9]))
  1375. a = array([], np.uint32, mask=[])
  1376. assert_equal(setdiff1d(a, []).dtype, np.uint32)
  1377. def test_setdiff1d_char_array(self):
  1378. # Test setdiff1d_charray
  1379. a = np.array(['a', 'b', 'c'])
  1380. b = np.array(['a', 'b', 's'])
  1381. assert_array_equal(setdiff1d(a, b), np.array(['c']))
  1382. class TestShapeBase(object):
  1383. def test_atleast_2d(self):
  1384. # Test atleast_2d
  1385. a = masked_array([0, 1, 2], mask=[0, 1, 0])
  1386. b = atleast_2d(a)
  1387. assert_equal(b.shape, (1, 3))
  1388. assert_equal(b.mask.shape, b.data.shape)
  1389. assert_equal(a.shape, (3,))
  1390. assert_equal(a.mask.shape, a.data.shape)
  1391. assert_equal(b.mask.shape, b.data.shape)
  1392. def test_shape_scalar(self):
  1393. # the atleast and diagflat function should work with scalars
  1394. # GitHub issue #3367
  1395. # Additionally, the atleast functions should accept multiple scalars
  1396. # correctly
  1397. b = atleast_1d(1.0)
  1398. assert_equal(b.shape, (1,))
  1399. assert_equal(b.mask.shape, b.shape)
  1400. assert_equal(b.data.shape, b.shape)
  1401. b = atleast_1d(1.0, 2.0)
  1402. for a in b:
  1403. assert_equal(a.shape, (1,))
  1404. assert_equal(a.mask.shape, a.shape)
  1405. assert_equal(a.data.shape, a.shape)
  1406. b = atleast_2d(1.0)
  1407. assert_equal(b.shape, (1, 1))
  1408. assert_equal(b.mask.shape, b.shape)
  1409. assert_equal(b.data.shape, b.shape)
  1410. b = atleast_2d(1.0, 2.0)
  1411. for a in b:
  1412. assert_equal(a.shape, (1, 1))
  1413. assert_equal(a.mask.shape, a.shape)
  1414. assert_equal(a.data.shape, a.shape)
  1415. b = atleast_3d(1.0)
  1416. assert_equal(b.shape, (1, 1, 1))
  1417. assert_equal(b.mask.shape, b.shape)
  1418. assert_equal(b.data.shape, b.shape)
  1419. b = atleast_3d(1.0, 2.0)
  1420. for a in b:
  1421. assert_equal(a.shape, (1, 1, 1))
  1422. assert_equal(a.mask.shape, a.shape)
  1423. assert_equal(a.data.shape, a.shape)
  1424. b = diagflat(1.0)
  1425. assert_equal(b.shape, (1, 1))
  1426. assert_equal(b.mask.shape, b.data.shape)
  1427. class TestStack(object):
  1428. def test_stack_1d(self):
  1429. a = masked_array([0, 1, 2], mask=[0, 1, 0])
  1430. b = masked_array([9, 8, 7], mask=[1, 0, 0])
  1431. c = stack([a, b], axis=0)
  1432. assert_equal(c.shape, (2, 3))
  1433. assert_array_equal(a.mask, c[0].mask)
  1434. assert_array_equal(b.mask, c[1].mask)
  1435. d = vstack([a, b])
  1436. assert_array_equal(c.data, d.data)
  1437. assert_array_equal(c.mask, d.mask)
  1438. c = stack([a, b], axis=1)
  1439. assert_equal(c.shape, (3, 2))
  1440. assert_array_equal(a.mask, c[:, 0].mask)
  1441. assert_array_equal(b.mask, c[:, 1].mask)
  1442. def test_stack_masks(self):
  1443. a = masked_array([0, 1, 2], mask=True)
  1444. b = masked_array([9, 8, 7], mask=False)
  1445. c = stack([a, b], axis=0)
  1446. assert_equal(c.shape, (2, 3))
  1447. assert_array_equal(a.mask, c[0].mask)
  1448. assert_array_equal(b.mask, c[1].mask)
  1449. d = vstack([a, b])
  1450. assert_array_equal(c.data, d.data)
  1451. assert_array_equal(c.mask, d.mask)
  1452. c = stack([a, b], axis=1)
  1453. assert_equal(c.shape, (3, 2))
  1454. assert_array_equal(a.mask, c[:, 0].mask)
  1455. assert_array_equal(b.mask, c[:, 1].mask)
  1456. def test_stack_nd(self):
  1457. # 2D
  1458. shp = (3, 2)
  1459. d1 = np.random.randint(0, 10, shp)
  1460. d2 = np.random.randint(0, 10, shp)
  1461. m1 = np.random.randint(0, 2, shp).astype(bool)
  1462. m2 = np.random.randint(0, 2, shp).astype(bool)
  1463. a1 = masked_array(d1, mask=m1)
  1464. a2 = masked_array(d2, mask=m2)
  1465. c = stack([a1, a2], axis=0)
  1466. c_shp = (2,) + shp
  1467. assert_equal(c.shape, c_shp)
  1468. assert_array_equal(a1.mask, c[0].mask)
  1469. assert_array_equal(a2.mask, c[1].mask)
  1470. c = stack([a1, a2], axis=-1)
  1471. c_shp = shp + (2,)
  1472. assert_equal(c.shape, c_shp)
  1473. assert_array_equal(a1.mask, c[..., 0].mask)
  1474. assert_array_equal(a2.mask, c[..., 1].mask)
  1475. # 4D
  1476. shp = (3, 2, 4, 5,)
  1477. d1 = np.random.randint(0, 10, shp)
  1478. d2 = np.random.randint(0, 10, shp)
  1479. m1 = np.random.randint(0, 2, shp).astype(bool)
  1480. m2 = np.random.randint(0, 2, shp).astype(bool)
  1481. a1 = masked_array(d1, mask=m1)
  1482. a2 = masked_array(d2, mask=m2)
  1483. c = stack([a1, a2], axis=0)
  1484. c_shp = (2,) + shp
  1485. assert_equal(c.shape, c_shp)
  1486. assert_array_equal(a1.mask, c[0].mask)
  1487. assert_array_equal(a2.mask, c[1].mask)
  1488. c = stack([a1, a2], axis=-1)
  1489. c_shp = shp + (2,)
  1490. assert_equal(c.shape, c_shp)
  1491. assert_array_equal(a1.mask, c[..., 0].mask)
  1492. assert_array_equal(a2.mask, c[..., 1].mask)