test_extras.py 65 KB

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