test_extras.py 65 KB

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