codekingpro/portable-devtools
114k
1"""Tests suite for MaskedArray.
2Adapted from the original test_ma by Pierre Gerard-Marchant
3
4:author: Pierre Gerard-Marchant
5:contact: pierregm_at_uga_dot_edu
6
7"""
8import inspect
9import itertools
10
11import pytest
12
13import numpy as np
14from numpy._core.numeric import normalize_axis_tuple
15from numpy.ma.core import (
16 MaskedArray,
17 arange,
18 array,
19 count,
20 getmaskarray,
21 masked,
22 masked_array,
23 nomask,
24 ones,
25 shape,
26 zeros,
27)
28from numpy.ma.extras import (
29 _covhelper,
30 apply_along_axis,
31 apply_over_axes,
32 atleast_1d,
33 atleast_2d,
34 atleast_3d,
35 average,
36 clump_masked,
37 clump_unmasked,
38 compress_nd,
39 compress_rowcols,
40 corrcoef,
41 cov,
42 diagflat,
43 dot,
44 ediff1d,
45 flatnotmasked_contiguous,
46 in1d,
47 intersect1d,
48 isin,
49 mask_rowcols,
50 masked_all,
51 masked_all_like,
52 median,
53 mr_,
54 ndenumerate,
55 notmasked_contiguous,
56 notmasked_edges,
57 polyfit,
58 setdiff1d,
59 setxor1d,
60 stack,
61 union1d,
62 unique,
63 vstack,
64)
65from numpy.ma.testutils import (
66 assert_,
67 assert_almost_equal,
68 assert_array_equal,
69 assert_equal,
70)
71
72
73class TestGeneric:
74 #
75 def test_masked_all(self):
76 # Tests masked_all
77 # Standard dtype
78 test = masked_all((2,), dtype=float)
79 control = array([1, 1], mask=[1, 1], dtype=float)
80 assert_equal(test, control)
81 # Flexible dtype
82 dt = np.dtype({'names': ['a', 'b'], 'formats': ['f', 'f']})
83 test = masked_all((2,), dtype=dt)
84 control = array([(0, 0), (0, 0)], mask=[(1, 1), (1, 1)], dtype=dt)
85 assert_equal(test, control)
86 test = masked_all((2, 2), dtype=dt)
87 control = array([[(0, 0), (0, 0)], [(0, 0), (0, 0)]],
88 mask=[[(1, 1), (1, 1)], [(1, 1), (1, 1)]],
89 dtype=dt)
90 assert_equal(test, control)
91 # Nested dtype
92 dt = np.dtype([('a', 'f'), ('b', [('ba', 'f'), ('bb', 'f')])])
93 test = masked_all((2,), dtype=dt)
94 control = array([(1, (1, 1)), (1, (1, 1))],
95 mask=[(1, (1, 1)), (1, (1, 1))], dtype=dt)
96 assert_equal(test, control)
97 test = masked_all((2,), dtype=dt)
98 control = array([(1, (1, 1)), (1, (1, 1))],
99 mask=[(1, (1, 1)), (1, (1, 1))], dtype=dt)
100 assert_equal(test, control)
101 test = masked_all((1, 1), dtype=dt)
102 control = array([[(1, (1, 1))]], mask=[[(1, (1, 1))]], dtype=dt)
103 assert_equal(test, control)
104
105 def test_masked_all_with_object_nested(self):
106 # Test masked_all works with nested array with dtype of an 'object'
107 # refers to issue #15895
108 my_dtype = np.dtype([('b', ([('c', object)], (1,)))])
109 masked_arr = np.ma.masked_all((1,), my_dtype)
110
111 assert_equal(type(masked_arr['b']), np.ma.core.MaskedArray)
112 assert_equal(type(masked_arr['b']['c']), np.ma.core.MaskedArray)
113 assert_equal(len(masked_arr['b']['c']), 1)
114 assert_equal(masked_arr['b']['c'].shape, (1, 1))
115 assert_equal(masked_arr['b']['c']._fill_value.shape, ())
116
117 def test_masked_all_with_object(self):
118 # same as above except that the array is not nested
119 my_dtype = np.dtype([('b', (object, (1,)))])
120 masked_arr = np.ma.masked_all((1,), my_dtype)
121
122 assert_equal(type(masked_arr['b']), np.ma.core.MaskedArray)
123 assert_equal(len(masked_arr['b']), 1)
124 assert_equal(masked_arr['b'].shape, (1, 1))
125 assert_equal(masked_arr['b']._fill_value.shape, ())
126
127 def test_masked_all_like(self):
128 # Tests masked_all
129 # Standard dtype
130 base = array([1, 2], dtype=float)
131 test = masked_all_like(base)
132 control = array([1, 1], mask=[1, 1], dtype=float)
133 assert_equal(test, control)
134 # Flexible dtype
135 dt = np.dtype({'names': ['a', 'b'], 'formats': ['f', 'f']})
136 base = array([(0, 0), (0, 0)], mask=[(1, 1), (1, 1)], dtype=dt)
137 test = masked_all_like(base)
138 control = array([(10, 10), (10, 10)], mask=[(1, 1), (1, 1)], dtype=dt)
139 assert_equal(test, control)
140 # Nested dtype
141 dt = np.dtype([('a', 'f'), ('b', [('ba', 'f'), ('bb', 'f')])])
142 control = array([(1, (1, 1)), (1, (1, 1))],
143 mask=[(1, (1, 1)), (1, (1, 1))], dtype=dt)
144 test = masked_all_like(control)
145 assert_equal(test, control)
146
147 def check_clump(self, f):
148 for i in range(1, 7):
149 for j in range(2**i):
150 k = np.arange(i, dtype=int)
151 ja = np.full(i, j, dtype=int)
152 a = masked_array(2**k)
153 a.mask = (ja & (2**k)) != 0
154 s = 0
155 for sl in f(a):
156 s += a.data[sl].sum()
157 if f == clump_unmasked:
158 assert_equal(a.compressed().sum(), s)
159 else:
160 a.mask = ~a.mask
161 assert_equal(a.compressed().sum(), s)
162
163 def test_clump_masked(self):
164 # Test clump_masked
165 a = masked_array(np.arange(10))
166 a[[0, 1, 2, 6, 8, 9]] = masked
167 #
168 test = clump_masked(a)
169 control = [slice(0, 3), slice(6, 7), slice(8, 10)]
170 assert_equal(test, control)
171
172 self.check_clump(clump_masked)
173
174 def test_clump_unmasked(self):
175 # Test clump_unmasked
176 a = masked_array(np.arange(10))
177 a[[0, 1, 2, 6, 8, 9]] = masked
178 test = clump_unmasked(a)
179 control = [slice(3, 6), slice(7, 8), ]
180 assert_equal(test, control)
181
182 self.check_clump(clump_unmasked)
183
184 def test_flatnotmasked_contiguous(self):
185 # Test flatnotmasked_contiguous
186 a = arange(10)
187 # No mask
188 test = flatnotmasked_contiguous(a)
189 assert_equal(test, [slice(0, a.size)])
190 # mask of all false
191 a.mask = np.zeros(10, dtype=bool)
192 assert_equal(test, [slice(0, a.size)])
193 # Some mask
194 a[(a < 3) | (a > 8) | (a == 5)] = masked
195 test = flatnotmasked_contiguous(a)
196 assert_equal(test, [slice(3, 5), slice(6, 9)])
197 #
198 a[:] = masked
199 test = flatnotmasked_contiguous(a)
200 assert_equal(test, [])
201
202
203class TestAverage:
204 # Several tests of average. Why so many ? Good point...
205 def test_testAverage1(self):
206 # Test of average.
207 ott = array([0., 1., 2., 3.], mask=[True, False, False, False])
208 assert_equal(2.0, average(ott, axis=0))
209 assert_equal(2.0, average(ott, weights=[1., 1., 2., 1.]))
210 result, wts = average(ott, weights=[1., 1., 2., 1.], returned=True)
211 assert_equal(2.0, result)
212 assert_(wts == 4.0)
213 ott[:] = masked
214 assert_equal(average(ott, axis=0).mask, [True])
215 ott = array([0., 1., 2., 3.], mask=[True, False, False, False])
216 ott = ott.reshape(2, 2)
217 ott[:, 1] = masked
218 assert_equal(average(ott, axis=0), [2.0, 0.0])
219 assert_equal(average(ott, axis=1).mask[0], [True])
220 assert_equal([2., 0.], average(ott, axis=0))
221 result, wts = average(ott, axis=0, returned=True)
222 assert_equal(wts, [1., 0.])
223
224 def test_testAverage2(self):
225 # More tests of average.
226 w1 = [0, 1, 1, 1, 1, 0]
227 w2 = [[0, 1, 1, 1, 1, 0], [1, 0, 0, 0, 0, 1]]
228 x = arange(6, dtype=np.float64)
229 assert_equal(average(x, axis=0), 2.5)
230 assert_equal(average(x, axis=0, weights=w1), 2.5)
231 y = array([arange(6, dtype=np.float64), 2.0 * arange(6)])
232 assert_equal(average(y, None), np.add.reduce(np.arange(6)) * 3. / 12.)
233 assert_equal(average(y, axis=0), np.arange(6) * 3. / 2.)
234 assert_equal(average(y, axis=1),
235 [average(x, axis=0), average(x, axis=0) * 2.0])
236 assert_equal(average(y, None, weights=w2), 20. / 6.)
237 assert_equal(average(y, axis=0, weights=w2),
238 [0., 1., 2., 3., 4., 10.])
239 assert_equal(average(y, axis=1),
240 [average(x, axis=0), average(x, axis=0) * 2.0])
241 m1 = zeros(6)
242 m2 = [0, 0, 1, 1, 0, 0]
243 m3 = [[0, 0, 1, 1, 0, 0], [0, 1, 1, 1, 1, 0]]
244 m4 = ones(6)
245 m5 = [0, 1, 1, 1, 1, 1]
246 assert_equal(average(masked_array(x, m1), axis=0), 2.5)
247 assert_equal(average(masked_array(x, m2), axis=0), 2.5)
248 assert_equal(average(masked_array(x, m4), axis=0).mask, [True])
249 assert_equal(average(masked_array(x, m5), axis=0), 0.0)
250 assert_equal(count(average(masked_array(x, m4), axis=0)), 0)
251 z = masked_array(y, m3)
252 assert_equal(average(z, None), 20. / 6.)
253 assert_equal(average(z, axis=0), [0., 1., 99., 99., 4.0, 7.5])
254 assert_equal(average(z, axis=1), [2.5, 5.0])
255 assert_equal(average(z, axis=0, weights=w2),
256 [0., 1., 99., 99., 4.0, 10.0])
257
258 def test_testAverage3(self):
259 # Yet more tests of average!
260 a = arange(6)
261 b = arange(6) * 3
262 r1, w1 = average([[a, b], [b, a]], axis=1, returned=True)
263 assert_equal(shape(r1), shape(w1))
264 assert_equal(r1.shape, w1.shape)
265 r2, w2 = average(ones((2, 2, 3)), axis=0, weights=[3, 1], returned=True)
266 assert_equal(shape(w2), shape(r2))
267 r2, w2 = average(ones((2, 2, 3)), returned=True)
268 assert_equal(shape(w2), shape(r2))
269 r2, w2 = average(ones((2, 2, 3)), weights=ones((2, 2, 3)), returned=True)
270 assert_equal(shape(w2), shape(r2))
271 a2d = array([[1, 2], [0, 4]], float)
272 a2dm = masked_array(a2d, [[False, False], [True, False]])
273 a2da = average(a2d, axis=0)
274 assert_equal(a2da, [0.5, 3.0])
275 a2dma = average(a2dm, axis=0)
276 assert_equal(a2dma, [1.0, 3.0])
277 a2dma = average(a2dm, axis=None)
278 assert_equal(a2dma, 7. / 3.)
279 a2dma = average(a2dm, axis=1)
280 assert_equal(a2dma, [1.5, 4.0])
281
282 def test_testAverage4(self):
283 # Test that `keepdims` works with average
284 x = np.array([2, 3, 4]).reshape(3, 1)
285 b = np.ma.array(x, mask=[[False], [False], [True]])
286 w = np.array([4, 5, 6]).reshape(3, 1)
287 actual = average(b, weights=w, axis=1, keepdims=True)
288 desired = masked_array([[2.], [3.], [4.]], [[False], [False], [True]])
289 assert_equal(actual, desired)
290
291 def test_weight_and_input_dims_different(self):
292 # this test mirrors a test for np.average()
293 # in lib/test/test_function_base.py
294 y = np.arange(12).reshape(2, 2, 3)
295 w = np.array([0., 0., 1., .5, .5, 0., 0., .5, .5, 1., 0., 0.])\
296 .reshape(2, 2, 3)
297
298 m = np.full((2, 2, 3), False)
299 yma = np.ma.array(y, mask=m)
300 subw0 = w[:, :, 0]
301
302 actual = average(yma, axis=(0, 1), weights=subw0)
303 desired = masked_array([7., 8., 9.], mask=[False, False, False])
304 assert_almost_equal(actual, desired)
305
306 m = np.full((2, 2, 3), False)
307 m[:, :, 0] = True
308 m[0, 0, 1] = True
309 yma = np.ma.array(y, mask=m)
310 actual = average(yma, axis=(0, 1), weights=subw0)
311 desired = masked_array(
312 [np.nan, 8., 9.],
313 mask=[True, False, False])
314 assert_almost_equal(actual, desired)
315
316 m = np.full((2, 2, 3), False)
317 yma = np.ma.array(y, mask=m)
318
319 subw1 = w[1, :, :]
320 actual = average(yma, axis=(1, 2), weights=subw1)
321 desired = masked_array([2.25, 8.25], mask=[False, False])
322 assert_almost_equal(actual, desired)
323
324 # here the weights have the wrong shape for the specified axes
325 with pytest.raises(
326 ValueError,
327 match="Shape of weights must be consistent with "
328 "shape of a along specified axis"):
329 average(yma, axis=(0, 1, 2), weights=subw0)
330
331 with pytest.raises(
332 ValueError,
333 match="Shape of weights must be consistent with "
334 "shape of a along specified axis"):
335 average(yma, axis=(0, 1), weights=subw1)
336
337 # swapping the axes should be same as transposing weights
338 actual = average(yma, axis=(1, 0), weights=subw0)
339 desired = average(yma, axis=(0, 1), weights=subw0.T)
340 assert_almost_equal(actual, desired)
341
342 def test_onintegers_with_mask(self):
343 # Test average on integers with mask
344 a = average(array([1, 2]))
345 assert_equal(a, 1.5)
346 a = average(array([1, 2, 3, 4], mask=[False, False, True, True]))
347 assert_equal(a, 1.5)
348
349 def test_complex(self):
350 # Test with complex data.
351 # (Regression test for https://github.com/numpy/numpy/issues/2684)
352 mask = np.array([[0, 0, 0, 1, 0],
353 [0, 1, 0, 0, 0]], dtype=bool)
354 a = masked_array([[0, 1 + 2j, 3 + 4j, 5 + 6j, 7 + 8j],
355 [9j, 0 + 1j, 2 + 3j, 4 + 5j, 7 + 7j]],
356 mask=mask)
357
358 av = average(a)
359 expected = np.average(a.compressed())
360 assert_almost_equal(av.real, expected.real)
361 assert_almost_equal(av.imag, expected.imag)
362
363 av0 = average(a, axis=0)
364 expected0 = average(a.real, axis=0) + average(a.imag, axis=0) * 1j
365 assert_almost_equal(av0.real, expected0.real)
366 assert_almost_equal(av0.imag, expected0.imag)
367
368 av1 = average(a, axis=1)
369 expected1 = average(a.real, axis=1) + average(a.imag, axis=1) * 1j
370 assert_almost_equal(av1.real, expected1.real)
371 assert_almost_equal(av1.imag, expected1.imag)
372
373 # Test with the 'weights' argument.
374 wts = np.array([[0.5, 1.0, 2.0, 1.0, 0.5],
375 [1.0, 1.0, 1.0, 1.0, 1.0]])
376 wav = average(a, weights=wts)
377 expected = np.average(a.compressed(), weights=wts[~mask])
378 assert_almost_equal(wav.real, expected.real)
379 assert_almost_equal(wav.imag, expected.imag)
380
381 wav0 = average(a, weights=wts, axis=0)
382 expected0 = (average(a.real, weights=wts, axis=0) +
383 average(a.imag, weights=wts, axis=0) * 1j)
384 assert_almost_equal(wav0.real, expected0.real)
385 assert_almost_equal(wav0.imag, expected0.imag)
386
387 wav1 = average(a, weights=wts, axis=1)
388 expected1 = (average(a.real, weights=wts, axis=1) +
389 average(a.imag, weights=wts, axis=1) * 1j)
390 assert_almost_equal(wav1.real, expected1.real)
391 assert_almost_equal(wav1.imag, expected1.imag)
392
393 @pytest.mark.parametrize(
394 'x, axis, expected_avg, weights, expected_wavg, expected_wsum',
395 [([1, 2, 3], None, [2.0], [3, 4, 1], [1.75], [8.0]),
396 ([[1, 2, 5], [1, 6, 11]], 0, [[1.0, 4.0, 8.0]],
397 [1, 3], [[1.0, 5.0, 9.5]], [[4, 4, 4]])],
398 )
399 def test_basic_keepdims(self, x, axis, expected_avg,
400 weights, expected_wavg, expected_wsum):
401 avg = np.ma.average(x, axis=axis, keepdims=True)
402 assert avg.shape == np.shape(expected_avg)
403 assert_array_equal(avg, expected_avg)
404
405 wavg = np.ma.average(x, axis=axis, weights=weights, keepdims=True)
406 assert wavg.shape == np.shape(expected_wavg)
407 assert_array_equal(wavg, expected_wavg)
408
409 wavg, wsum = np.ma.average(x, axis=axis, weights=weights,
410 returned=True, keepdims=True)
411 assert wavg.shape == np.shape(expected_wavg)
412 assert_array_equal(wavg, expected_wavg)
413 assert wsum.shape == np.shape(expected_wsum)
414 assert_array_equal(wsum, expected_wsum)
415
416 def test_masked_weights(self):
417 # Test with masked weights.
418 # (Regression test for https://github.com/numpy/numpy/issues/10438)
419 a = np.ma.array(np.arange(9).reshape(3, 3),
420 mask=[[1, 0, 0], [1, 0, 0], [0, 0, 0]])
421 weights_unmasked = masked_array([5, 28, 31], mask=False)
422 weights_masked = masked_array([5, 28, 31], mask=[1, 0, 0])
423
424 avg_unmasked = average(a, axis=0,
425 weights=weights_unmasked, returned=False)
426 expected_unmasked = np.array([6.0, 5.21875, 6.21875])
427 assert_almost_equal(avg_unmasked, expected_unmasked)
428
429 avg_masked = average(a, axis=0, weights=weights_masked, returned=False)
430 expected_masked = np.array([6.0, 5.576271186440678, 6.576271186440678])
431 assert_almost_equal(avg_masked, expected_masked)
432
433 # weights should be masked if needed
434 # depending on the array mask. This is to avoid summing
435 # masked nan or other values that are not cancelled by a zero
436 a = np.ma.array([1.0, 2.0, 3.0, 4.0],
437 mask=[False, False, True, True])
438 avg_unmasked = average(a, weights=[1, 1, 1, np.nan])
439
440 assert_almost_equal(avg_unmasked, 1.5)
441
442 a = np.ma.array([
443 [1.0, 2.0, 3.0, 4.0],
444 [5.0, 6.0, 7.0, 8.0],
445 [9.0, 1.0, 2.0, 3.0],
446 ], mask=[
447 [False, True, True, False],
448 [True, False, True, True],
449 [True, False, True, False],
450 ])
451
452 avg_masked = np.ma.average(a, weights=[1, np.nan, 1], axis=0)
453 avg_expected = np.ma.array([1.0, np.nan, np.nan, 3.5],
454 mask=[False, True, True, False])
455
456 assert_almost_equal(avg_masked, avg_expected)
457 assert_equal(avg_masked.mask, avg_expected.mask)
458
459
460class TestConcatenator:
461 # Tests for mr_, the equivalent of r_ for masked arrays.
462
463 def test_1d(self):
464 # Tests mr_ on 1D arrays.
465 assert_array_equal(mr_[1, 2, 3, 4, 5, 6], array([1, 2, 3, 4, 5, 6]))
466 b = ones(5)
467 m = [1, 0, 0, 0, 0]
468 d = masked_array(b, mask=m)
469 c = mr_[d, 0, 0, d]
470 assert_(isinstance(c, MaskedArray))
471 assert_array_equal(c, [1, 1, 1, 1, 1, 0, 0, 1, 1, 1, 1, 1])
472 assert_array_equal(c.mask, mr_[m, 0, 0, m])
473
474 def test_2d(self):
475 # Tests mr_ on 2D arrays.
476 a_1 = np.random.rand(5, 5)
477 a_2 = np.random.rand(5, 5)
478 m_1 = np.round(np.random.rand(5, 5), 0)
479 m_2 = np.round(np.random.rand(5, 5), 0)
480 b_1 = masked_array(a_1, mask=m_1)
481 b_2 = masked_array(a_2, mask=m_2)
482 # append columns
483 d = mr_['1', b_1, b_2]
484 assert_(d.shape == (5, 10))
485 assert_array_equal(d[:, :5], b_1)
486 assert_array_equal(d[:, 5:], b_2)
487 assert_array_equal(d.mask, np.r_['1', m_1, m_2])
488 d = mr_[b_1, b_2]
489 assert_(d.shape == (10, 5))
490 assert_array_equal(d[:5, :], b_1)
491 assert_array_equal(d[5:, :], b_2)
492 assert_array_equal(d.mask, np.r_[m_1, m_2])
493
494 def test_masked_constant(self):
495 actual = mr_[np.ma.masked, 1]
496 assert_equal(actual.mask, [True, False])
497 assert_equal(actual.data[1], 1)
498
499 actual = mr_[[1, 2], np.ma.masked]
500 assert_equal(actual.mask, [False, False, True])
501 assert_equal(actual.data[:2], [1, 2])
502
503
504class TestNotMasked:
505 # Tests notmasked_edges and notmasked_contiguous.
506
507 def test_edges(self):
508 # Tests unmasked_edges
509 data = masked_array(np.arange(25).reshape(5, 5),
510 mask=[[0, 0, 1, 0, 0],
511 [0, 0, 0, 1, 1],
512 [1, 1, 0, 0, 0],
513 [0, 0, 0, 0, 0],
514 [1, 1, 1, 0, 0]],)
515 test = notmasked_edges(data, None)
516 assert_equal(test, [0, 24])
517 test = notmasked_edges(data, 0)
518 assert_equal(test[0], [(0, 0, 1, 0, 0), (0, 1, 2, 3, 4)])
519 assert_equal(test[1], [(3, 3, 3, 4, 4), (0, 1, 2, 3, 4)])
520 test = notmasked_edges(data, 1)
521 assert_equal(test[0], [(0, 1, 2, 3, 4), (0, 0, 2, 0, 3)])
522 assert_equal(test[1], [(0, 1, 2, 3, 4), (4, 2, 4, 4, 4)])
523 #
524 test = notmasked_edges(data.data, None)
525 assert_equal(test, [0, 24])
526 test = notmasked_edges(data.data, 0)
527 assert_equal(test[0], [(0, 0, 0, 0, 0), (0, 1, 2, 3, 4)])
528 assert_equal(test[1], [(4, 4, 4, 4, 4), (0, 1, 2, 3, 4)])
529 test = notmasked_edges(data.data, -1)
530 assert_equal(test[0], [(0, 1, 2, 3, 4), (0, 0, 0, 0, 0)])
531 assert_equal(test[1], [(0, 1, 2, 3, 4), (4, 4, 4, 4, 4)])
532 #
533 data[-2] = masked
534 test = notmasked_edges(data, 0)
535 assert_equal(test[0], [(0, 0, 1, 0, 0), (0, 1, 2, 3, 4)])
536 assert_equal(test[1], [(1, 1, 2, 4, 4), (0, 1, 2, 3, 4)])
537 test = notmasked_edges(data, -1)
538 assert_equal(test[0], [(0, 1, 2, 4), (0, 0, 2, 3)])
539 assert_equal(test[1], [(0, 1, 2, 4), (4, 2, 4, 4)])
540
541 def test_contiguous(self):
542 # Tests notmasked_contiguous
543 a = masked_array(np.arange(24).reshape(3, 8),
544 mask=[[0, 0, 0, 0, 1, 1, 1, 1],
545 [1, 1, 1, 1, 1, 1, 1, 1],
546 [0, 0, 0, 0, 0, 0, 1, 0]])
547 tmp = notmasked_contiguous(a, None)
548 assert_equal(tmp, [
549 slice(0, 4, None),
550 slice(16, 22, None),
551 slice(23, 24, None)
552 ])
553
554 tmp = notmasked_contiguous(a, 0)
555 assert_equal(tmp, [
556 [slice(0, 1, None), slice(2, 3, None)],
557 [slice(0, 1, None), slice(2, 3, None)],
558 [slice(0, 1, None), slice(2, 3, None)],
559 [slice(0, 1, None), slice(2, 3, None)],
560 [slice(2, 3, None)],
561 [slice(2, 3, None)],
562 [],
563 [slice(2, 3, None)]
564 ])
565 #
566 tmp = notmasked_contiguous(a, 1)
567 assert_equal(tmp, [
568 [slice(0, 4, None)],
569 [],
570 [slice(0, 6, None), slice(7, 8, None)]
571 ])
572
573
574class TestCompressFunctions:
575
576 def test_compress_nd(self):
577 # Tests compress_nd
578 x = np.array(list(range(3 * 4 * 5))).reshape(3, 4, 5)
579 m = np.zeros((3, 4, 5)).astype(bool)
580 m[1, 1, 1] = True
581 x = array(x, mask=m)
582
583 # axis=None
584 a = compress_nd(x)
585 assert_equal(a, [[[ 0, 2, 3, 4],
586 [10, 12, 13, 14],
587 [15, 17, 18, 19]],
588 [[40, 42, 43, 44],
589 [50, 52, 53, 54],
590 [55, 57, 58, 59]]])
591
592 # axis=0
593 a = compress_nd(x, 0)
594 assert_equal(a, [[[ 0, 1, 2, 3, 4],
595 [ 5, 6, 7, 8, 9],
596 [10, 11, 12, 13, 14],
597 [15, 16, 17, 18, 19]],
598 [[40, 41, 42, 43, 44],
599 [45, 46, 47, 48, 49],
600 [50, 51, 52, 53, 54],
601 [55, 56, 57, 58, 59]]])
602
603 # axis=1
604 a = compress_nd(x, 1)
605 assert_equal(a, [[[ 0, 1, 2, 3, 4],
606 [10, 11, 12, 13, 14],
607 [15, 16, 17, 18, 19]],
608 [[20, 21, 22, 23, 24],
609 [30, 31, 32, 33, 34],
610 [35, 36, 37, 38, 39]],
611 [[40, 41, 42, 43, 44],
612 [50, 51, 52, 53, 54],
613 [55, 56, 57, 58, 59]]])
614
615 a2 = compress_nd(x, (1,))
616 a3 = compress_nd(x, -2)
617 a4 = compress_nd(x, (-2,))
618 assert_equal(a, a2)
619 assert_equal(a, a3)
620 assert_equal(a, a4)
621
622 # axis=2
623 a = compress_nd(x, 2)
624 assert_equal(a, [[[ 0, 2, 3, 4],
625 [ 5, 7, 8, 9],
626 [10, 12, 13, 14],
627 [15, 17, 18, 19]],
628 [[20, 22, 23, 24],
629 [25, 27, 28, 29],
630 [30, 32, 33, 34],
631 [35, 37, 38, 39]],
632 [[40, 42, 43, 44],
633 [45, 47, 48, 49],
634 [50, 52, 53, 54],
635 [55, 57, 58, 59]]])
636
637 a2 = compress_nd(x, (2,))
638 a3 = compress_nd(x, -1)
639 a4 = compress_nd(x, (-1,))
640 assert_equal(a, a2)
641 assert_equal(a, a3)
642 assert_equal(a, a4)
643
644 # axis=(0, 1)
645 a = compress_nd(x, (0, 1))
646 assert_equal(a, [[[ 0, 1, 2, 3, 4],
647 [10, 11, 12, 13, 14],
648 [15, 16, 17, 18, 19]],
649 [[40, 41, 42, 43, 44],
650 [50, 51, 52, 53, 54],
651 [55, 56, 57, 58, 59]]])
652 a2 = compress_nd(x, (0, -2))
653 assert_equal(a, a2)
654
655 # axis=(1, 2)
656 a = compress_nd(x, (1, 2))
657 assert_equal(a, [[[ 0, 2, 3, 4],
658 [10, 12, 13, 14],
659 [15, 17, 18, 19]],
660 [[20, 22, 23, 24],
661 [30, 32, 33, 34],
662 [35, 37, 38, 39]],
663 [[40, 42, 43, 44],
664 [50, 52, 53, 54],
665 [55, 57, 58, 59]]])
666
667 a2 = compress_nd(x, (-2, 2))
668 a3 = compress_nd(x, (1, -1))
669 a4 = compress_nd(x, (-2, -1))
670 assert_equal(a, a2)
671 assert_equal(a, a3)
672 assert_equal(a, a4)
673
674 # axis=(0, 2)
675 a = compress_nd(x, (0, 2))
676 assert_equal(a, [[[ 0, 2, 3, 4],
677 [ 5, 7, 8, 9],
678 [10, 12, 13, 14],
679 [15, 17, 18, 19]],
680 [[40, 42, 43, 44],
681 [45, 47, 48, 49],
682 [50, 52, 53, 54],
683 [55, 57, 58, 59]]])
684
685 a2 = compress_nd(x, (0, -1))
686 assert_equal(a, a2)
687
688 def test_compress_rowcols(self):
689 # Tests compress_rowcols
690 x = array(np.arange(9).reshape(3, 3),
691 mask=[[1, 0, 0], [0, 0, 0], [0, 0, 0]])
692 assert_equal(compress_rowcols(x), [[4, 5], [7, 8]])
693 assert_equal(compress_rowcols(x, 0), [[3, 4, 5], [6, 7, 8]])
694 assert_equal(compress_rowcols(x, 1), [[1, 2], [4, 5], [7, 8]])
695 x = array(x._data, mask=[[0, 0, 0], [0, 1, 0], [0, 0, 0]])
696 assert_equal(compress_rowcols(x), [[0, 2], [6, 8]])
697 assert_equal(compress_rowcols(x, 0), [[0, 1, 2], [6, 7, 8]])
698 assert_equal(compress_rowcols(x, 1), [[0, 2], [3, 5], [6, 8]])
699 x = array(x._data, mask=[[1, 0, 0], [0, 1, 0], [0, 0, 0]])
700 assert_equal(compress_rowcols(x), [[8]])
701 assert_equal(compress_rowcols(x, 0), [[6, 7, 8]])
702 assert_equal(compress_rowcols(x, 1,), [[2], [5], [8]])
703 x = array(x._data, mask=[[1, 0, 0], [0, 1, 0], [0, 0, 1]])
704 assert_equal(compress_rowcols(x).size, 0)
705 assert_equal(compress_rowcols(x, 0).size, 0)
706 assert_equal(compress_rowcols(x, 1).size, 0)
707
708 def test_mask_rowcols(self):
709 # Tests mask_rowcols.
710 x = array(np.arange(9).reshape(3, 3),
711 mask=[[1, 0, 0], [0, 0, 0], [0, 0, 0]])
712 assert_equal(mask_rowcols(x).mask,
713 [[1, 1, 1], [1, 0, 0], [1, 0, 0]])
714 assert_equal(mask_rowcols(x, 0).mask,
715 [[1, 1, 1], [0, 0, 0], [0, 0, 0]])
716 assert_equal(mask_rowcols(x, 1).mask,
717 [[1, 0, 0], [1, 0, 0], [1, 0, 0]])
718 x = array(x._data, mask=[[0, 0, 0], [0, 1, 0], [0, 0, 0]])
719 assert_equal(mask_rowcols(x).mask,
720 [[0, 1, 0], [1, 1, 1], [0, 1, 0]])
721 assert_equal(mask_rowcols(x, 0).mask,
722 [[0, 0, 0], [1, 1, 1], [0, 0, 0]])
723 assert_equal(mask_rowcols(x, 1).mask,
724 [[0, 1, 0], [0, 1, 0], [0, 1, 0]])
725 x = array(x._data, mask=[[1, 0, 0], [0, 1, 0], [0, 0, 0]])
726 assert_equal(mask_rowcols(x).mask,
727 [[1, 1, 1], [1, 1, 1], [1, 1, 0]])
728 assert_equal(mask_rowcols(x, 0).mask,
729 [[1, 1, 1], [1, 1, 1], [0, 0, 0]])
730 assert_equal(mask_rowcols(x, 1,).mask,
731 [[1, 1, 0], [1, 1, 0], [1, 1, 0]])
732 x = array(x._data, mask=[[1, 0, 0], [0, 1, 0], [0, 0, 1]])
733 assert_(mask_rowcols(x).all() is masked)
734 assert_(mask_rowcols(x, 0).all() is masked)
735 assert_(mask_rowcols(x, 1).all() is masked)
736 assert_(mask_rowcols(x).mask.all())
737 assert_(mask_rowcols(x, 0).mask.all())
738 assert_(mask_rowcols(x, 1).mask.all())
739
740 @pytest.mark.parametrize("axis", [None, 0, 1])
741 @pytest.mark.parametrize(["func", "rowcols_axis"],
742 [(np.ma.mask_rows, 0), (np.ma.mask_cols, 1)])
743 def test_mask_row_cols_axis_deprecation(self, axis, func, rowcols_axis):
744 # Test deprecation of the axis argument to `mask_rows` and `mask_cols`
745 x = array(np.arange(9).reshape(3, 3),
746 mask=[[1, 0, 0], [0, 0, 0], [0, 0, 0]])
747
748 with pytest.warns(DeprecationWarning):
749 res = func(x, axis=axis)
750 assert_equal(res, mask_rowcols(x, rowcols_axis))
751
752 def test_dot(self):
753 # Tests dot product
754 n = np.arange(1, 7)
755 #
756 m = [1, 0, 0, 0, 0, 0]
757 a = masked_array(n, mask=m).reshape(2, 3)
758 b = masked_array(n, mask=m).reshape(3, 2)
759 c = dot(a, b, strict=True)
760 assert_equal(c.mask, [[1, 1], [1, 0]])
761 c = dot(b, a, strict=True)
762 assert_equal(c.mask, [[1, 1, 1], [1, 0, 0], [1, 0, 0]])
763 c = dot(a, b, strict=False)
764 assert_equal(c, np.dot(a.filled(0), b.filled(0)))
765 c = dot(b, a, strict=False)
766 assert_equal(c, np.dot(b.filled(0), a.filled(0)))
767 #
768 m = [0, 0, 0, 0, 0, 1]
769 a = masked_array(n, mask=m).reshape(2, 3)
770 b = masked_array(n, mask=m).reshape(3, 2)
771 c = dot(a, b, strict=True)
772 assert_equal(c.mask, [[0, 1], [1, 1]])
773 c = dot(b, a, strict=True)
774 assert_equal(c.mask, [[0, 0, 1], [0, 0, 1], [1, 1, 1]])
775 c = dot(a, b, strict=False)
776 assert_equal(c, np.dot(a.filled(0), b.filled(0)))
777 assert_equal(c, dot(a, b))
778 c = dot(b, a, strict=False)
779 assert_equal(c, np.dot(b.filled(0), a.filled(0)))
780 #
781 m = [0, 0, 0, 0, 0, 0]
782 a = masked_array(n, mask=m).reshape(2, 3)
783 b = masked_array(n, mask=m).reshape(3, 2)
784 c = dot(a, b)
785 assert_equal(c.mask, nomask)
786 c = dot(b, a)
787 assert_equal(c.mask, nomask)
788 #
789 a = masked_array(n, mask=[1, 0, 0, 0, 0, 0]).reshape(2, 3)
790 b = masked_array(n, mask=[0, 0, 0, 0, 0, 0]).reshape(3, 2)
791 c = dot(a, b, strict=True)
792 assert_equal(c.mask, [[1, 1], [0, 0]])
793 c = dot(a, b, strict=False)
794 assert_equal(c, np.dot(a.filled(0), b.filled(0)))
795 c = dot(b, a, strict=True)
796 assert_equal(c.mask, [[1, 0, 0], [1, 0, 0], [1, 0, 0]])
797 c = dot(b, a, strict=False)
798 assert_equal(c, np.dot(b.filled(0), a.filled(0)))
799 #
800 a = masked_array(n, mask=[0, 0, 0, 0, 0, 1]).reshape(2, 3)
801 b = masked_array(n, mask=[0, 0, 0, 0, 0, 0]).reshape(3, 2)
802 c = dot(a, b, strict=True)
803 assert_equal(c.mask, [[0, 0], [1, 1]])
804 c = dot(a, b)
805 assert_equal(c, np.dot(a.filled(0), b.filled(0)))
806 c = dot(b, a, strict=True)
807 assert_equal(c.mask, [[0, 0, 1], [0, 0, 1], [0, 0, 1]])
808 c = dot(b, a, strict=False)
809 assert_equal(c, np.dot(b.filled(0), a.filled(0)))
810 #
811 a = masked_array(n, mask=[0, 0, 0, 0, 0, 1]).reshape(2, 3)
812 b = masked_array(n, mask=[0, 0, 1, 0, 0, 0]).reshape(3, 2)
813 c = dot(a, b, strict=True)
814 assert_equal(c.mask, [[1, 0], [1, 1]])
815 c = dot(a, b, strict=False)
816 assert_equal(c, np.dot(a.filled(0), b.filled(0)))
817 c = dot(b, a, strict=True)
818 assert_equal(c.mask, [[0, 0, 1], [1, 1, 1], [0, 0, 1]])
819 c = dot(b, a, strict=False)
820 assert_equal(c, np.dot(b.filled(0), a.filled(0)))
821 #
822 a = masked_array(np.arange(8).reshape(2, 2, 2),
823 mask=[[[1, 0], [0, 0]], [[0, 0], [0, 0]]])
824 b = masked_array(np.arange(8).reshape(2, 2, 2),
825 mask=[[[0, 0], [0, 0]], [[0, 0], [0, 1]]])
826 c = dot(a, b, strict=True)
827 assert_equal(c.mask,
828 [[[[1, 1], [1, 1]], [[0, 0], [0, 1]]],
829 [[[0, 0], [0, 1]], [[0, 0], [0, 1]]]])
830 c = dot(a, b, strict=False)
831 assert_equal(c.mask,
832 [[[[0, 0], [0, 1]], [[0, 0], [0, 0]]],
833 [[[0, 0], [0, 0]], [[0, 0], [0, 0]]]])
834 c = dot(b, a, strict=True)
835 assert_equal(c.mask,
836 [[[[1, 0], [0, 0]], [[1, 0], [0, 0]]],
837 [[[1, 0], [0, 0]], [[1, 1], [1, 1]]]])
838 c = dot(b, a, strict=False)
839 assert_equal(c.mask,
840 [[[[0, 0], [0, 0]], [[0, 0], [0, 0]]],
841 [[[0, 0], [0, 0]], [[1, 0], [0, 0]]]])
842 #
843 a = masked_array(np.arange(8).reshape(2, 2, 2),
844 mask=[[[1, 0], [0, 0]], [[0, 0], [0, 0]]])
845 b = 5.
846 c = dot(a, b, strict=True)
847 assert_equal(c.mask, [[[1, 0], [0, 0]], [[0, 0], [0, 0]]])
848 c = dot(a, b, strict=False)
849 assert_equal(c.mask, [[[1, 0], [0, 0]], [[0, 0], [0, 0]]])
850 c = dot(b, a, strict=True)
851 assert_equal(c.mask, [[[1, 0], [0, 0]], [[0, 0], [0, 0]]])
852 c = dot(b, a, strict=False)
853 assert_equal(c.mask, [[[1, 0], [0, 0]], [[0, 0], [0, 0]]])
854 #
855 a = masked_array(np.arange(8).reshape(2, 2, 2),
856 mask=[[[1, 0], [0, 0]], [[0, 0], [0, 0]]])
857 b = masked_array(np.arange(2), mask=[0, 1])
858 c = dot(a, b, strict=True)
859 assert_equal(c.mask, [[1, 1], [1, 1]])
860 c = dot(a, b, strict=False)
861 assert_equal(c.mask, [[1, 0], [0, 0]])
862
863 def test_dot_returns_maskedarray(self):
864 # See gh-6611
865 a = np.eye(3)
866 b = array(a)
867 assert_(type(dot(a, a)) is MaskedArray)
868 assert_(type(dot(a, b)) is MaskedArray)
869 assert_(type(dot(b, a)) is MaskedArray)
870 assert_(type(dot(b, b)) is MaskedArray)
871
872 def test_dot_out(self):
873 a = array(np.eye(3))
874 out = array(np.zeros((3, 3)))
875 res = dot(a, a, out=out)
876 assert_(res is out)
877 assert_equal(a, res)
878
879
880class TestApplyAlongAxis:
881 # Tests 2D functions
882 def test_3d(self):
883 a = arange(12.).reshape(2, 2, 3)
884
885 def myfunc(b):
886 return b[1]
887
888 xa = apply_along_axis(myfunc, 2, a)
889 assert_equal(xa, [[1, 4], [7, 10]])
890
891 # Tests kwargs functions
892 def test_3d_kwargs(self):
893 a = arange(12).reshape(2, 2, 3)
894
895 def myfunc(b, offset=0):
896 return b[1 + offset]
897
898 xa = apply_along_axis(myfunc, 2, a, offset=1)
899 assert_equal(xa, [[2, 5], [8, 11]])
900
901
902class TestApplyOverAxes:
903 # Tests apply_over_axes
904 def test_basic(self):
905 a = arange(24).reshape(2, 3, 4)
906 test = apply_over_axes(np.sum, a, [0, 2])
907 ctrl = np.array([[[60], [92], [124]]])
908 assert_equal(test, ctrl)
909 a[(a % 2).astype(bool)] = masked
910 test = apply_over_axes(np.sum, a, [0, 2])
911 ctrl = np.array([[[28], [44], [60]]])
912 assert_equal(test, ctrl)
913
914
915class TestMedian:
916 def test_pytype(self):
917 r = np.ma.median([[np.inf, np.inf], [np.inf, np.inf]], axis=-1)
918 assert_equal(r, np.inf)
919
920 def test_inf(self):
921 # test that even which computes handles inf / x = masked
922 r = np.ma.median(np.ma.masked_array([[np.inf, np.inf],
923 [np.inf, np.inf]]), axis=-1)
924 assert_equal(r, np.inf)
925 r = np.ma.median(np.ma.masked_array([[np.inf, np.inf],
926 [np.inf, np.inf]]), axis=None)
927 assert_equal(r, np.inf)
928 # all masked
929 r = np.ma.median(np.ma.masked_array([[np.inf, np.inf],
930 [np.inf, np.inf]], mask=True),
931 axis=-1)
932 assert_equal(r.mask, True)
933 r = np.ma.median(np.ma.masked_array([[np.inf, np.inf],
934 [np.inf, np.inf]], mask=True),
935 axis=None)
936 assert_equal(r.mask, True)
937
938 def test_non_masked(self):
939 x = np.arange(9)
940 assert_equal(np.ma.median(x), 4.)
941 assert_(type(np.ma.median(x)) is not MaskedArray)
942 x = range(8)
943 assert_equal(np.ma.median(x), 3.5)
944 assert_(type(np.ma.median(x)) is not MaskedArray)
945 x = 5
946 assert_equal(np.ma.median(x), 5.)
947 assert_(type(np.ma.median(x)) is not MaskedArray)
948 # integer
949 x = np.arange(9 * 8).reshape(9, 8)
950 assert_equal(np.ma.median(x, axis=0), np.median(x, axis=0))
951 assert_equal(np.ma.median(x, axis=1), np.median(x, axis=1))
952 assert_(np.ma.median(x, axis=1) is not MaskedArray)
953 # float
954 x = np.arange(9 * 8.).reshape(9, 8)
955 assert_equal(np.ma.median(x, axis=0), np.median(x, axis=0))
956 assert_equal(np.ma.median(x, axis=1), np.median(x, axis=1))
957 assert_(np.ma.median(x, axis=1) is not MaskedArray)
958
959 def test_docstring_examples(self):
960 "test the examples given in the docstring of ma.median"
961 x = array(np.arange(8), mask=[0] * 4 + [1] * 4)
962 assert_equal(np.ma.median(x), 1.5)
963 assert_equal(np.ma.median(x).shape, (), "shape mismatch")
964 assert_(type(np.ma.median(x)) is not MaskedArray)
965 x = array(np.arange(10).reshape(2, 5), mask=[0] * 6 + [1] * 4)
966 assert_equal(np.ma.median(x), 2.5)
967 assert_equal(np.ma.median(x).shape, (), "shape mismatch")
968 assert_(type(np.ma.median(x)) is not MaskedArray)
969 ma_x = np.ma.median(x, axis=-1, overwrite_input=True)
970 assert_equal(ma_x, [2., 5.])
971 assert_equal(ma_x.shape, (2,), "shape mismatch")
972 assert_(type(ma_x) is MaskedArray)
973
974 def test_axis_argument_errors(self):
975 msg = "mask = %s, ndim = %s, axis = %s, overwrite_input = %s"
976 for ndmin in range(5):
977 for mask in [False, True]:
978 x = array(1, ndmin=ndmin, mask=mask)
979
980 # Valid axis values should not raise exception
981 args = itertools.product(range(-ndmin, ndmin), [False, True])
982 for axis, over in args:
983 try:
984 np.ma.median(x, axis=axis, overwrite_input=over)
985 except Exception:
986 raise AssertionError(msg % (mask, ndmin, axis, over))
987
988 # Invalid axis values should raise exception
989 args = itertools.product([-(ndmin + 1), ndmin], [False, True])
990 for axis, over in args:
991 try:
992 np.ma.median(x, axis=axis, overwrite_input=over)
993 except np.exceptions.AxisError:
994 pass
995 else:
996 raise AssertionError(msg % (mask, ndmin, axis, over))
997
998 def test_masked_0d(self):
999 # Check values
1000 x = array(1, mask=False)
1001 assert_equal(np.ma.median(x), 1)
1002 x = array(1, mask=True)
1003 assert_equal(np.ma.median(x), np.ma.masked)
1004
1005 def test_masked_1d(self):
1006 x = array(np.arange(5), mask=True)
1007 assert_equal(np.ma.median(x), np.ma.masked)
1008 assert_equal(np.ma.median(x).shape, (), "shape mismatch")
1009 assert_(type(np.ma.median(x)) is np.ma.core.MaskedConstant)
1010 x = array(np.arange(5), mask=False)
1011 assert_equal(np.ma.median(x), 2.)
1012 assert_equal(np.ma.median(x).shape, (), "shape mismatch")
1013 assert_(type(np.ma.median(x)) is not MaskedArray)
1014 x = array(np.arange(5), mask=[0, 1, 0, 0, 0])
1015 assert_equal(np.ma.median(x), 2.5)
1016 assert_equal(np.ma.median(x).shape, (), "shape mismatch")
1017 assert_(type(np.ma.median(x)) is not MaskedArray)
1018 x = array(np.arange(5), mask=[0, 1, 1, 1, 1])
1019 assert_equal(np.ma.median(x), 0.)
1020 assert_equal(np.ma.median(x).shape, (), "shape mismatch")
1021 assert_(type(np.ma.median(x)) is not MaskedArray)
1022 # integer
1023 x = array(np.arange(5), mask=[0, 1, 1, 0, 0])
1024 assert_equal(np.ma.median(x), 3.)
1025 assert_equal(np.ma.median(x).shape, (), "shape mismatch")
1026 assert_(type(np.ma.median(x)) is not MaskedArray)
1027 # float
1028 x = array(np.arange(5.), mask=[0, 1, 1, 0, 0])
1029 assert_equal(np.ma.median(x), 3.)
1030 assert_equal(np.ma.median(x).shape, (), "shape mismatch")
1031 assert_(type(np.ma.median(x)) is not MaskedArray)
1032 # integer
1033 x = array(np.arange(6), mask=[0, 1, 1, 1, 1, 0])
1034 assert_equal(np.ma.median(x), 2.5)
1035 assert_equal(np.ma.median(x).shape, (), "shape mismatch")
1036 assert_(type(np.ma.median(x)) is not MaskedArray)
1037 # float
1038 x = array(np.arange(6.), mask=[0, 1, 1, 1, 1, 0])
1039 assert_equal(np.ma.median(x), 2.5)
1040 assert_equal(np.ma.median(x).shape, (), "shape mismatch")
1041 assert_(type(np.ma.median(x)) is not MaskedArray)
1042
1043 def test_1d_shape_consistency(self):
1044 assert_equal(np.ma.median(array([1, 2, 3], mask=[0, 0, 0])).shape,
1045 np.ma.median(array([1, 2, 3], mask=[0, 1, 0])).shape)
1046
1047 def test_2d(self):
1048 # Tests median w/ 2D
1049 (n, p) = (101, 30)
1050 x = masked_array(np.linspace(-1., 1., n),)
1051 x[:10] = x[-10:] = masked
1052 z = masked_array(np.empty((n, p), dtype=float))
1053 z[:, 0] = x[:]
1054 idx = np.arange(len(x))
1055 for i in range(1, p):
1056 np.random.shuffle(idx)
1057 z[:, i] = x[idx]
1058 assert_equal(median(z[:, 0]), 0)
1059 assert_equal(median(z), 0)
1060 assert_equal(median(z, axis=0), np.zeros(p))
1061 assert_equal(median(z.T, axis=1), np.zeros(p))
1062
1063 def test_2d_waxis(self):
1064 # Tests median w/ 2D arrays and different axis.
1065 x = masked_array(np.arange(30).reshape(10, 3))
1066 x[:3] = x[-3:] = masked
1067 assert_equal(median(x), 14.5)
1068 assert_(type(np.ma.median(x)) is not MaskedArray)
1069 assert_equal(median(x, axis=0), [13.5, 14.5, 15.5])
1070 assert_(type(np.ma.median(x, axis=0)) is MaskedArray)
1071 assert_equal(median(x, axis=1), [0, 0, 0, 10, 13, 16, 19, 0, 0, 0])
1072 assert_(type(np.ma.median(x, axis=1)) is MaskedArray)
1073 assert_equal(median(x, axis=1).mask, [1, 1, 1, 0, 0, 0, 0, 1, 1, 1])
1074
1075 def test_3d(self):
1076 # Tests median w/ 3D
1077 x = np.ma.arange(24).reshape(3, 4, 2)
1078 x[x % 3 == 0] = masked
1079 assert_equal(median(x, 0), [[12, 9], [6, 15], [12, 9], [18, 15]])
1080 x = x.reshape((4, 3, 2))
1081 assert_equal(median(x, 0), [[99, 10], [11, 99], [13, 14]])
1082 x = np.ma.arange(24).reshape(4, 3, 2)
1083 x[x % 5 == 0] = masked
1084 assert_equal(median(x, 0), [[12, 10], [8, 9], [16, 17]])
1085
1086 def test_neg_axis(self):
1087 x = masked_array(np.arange(30).reshape(10, 3))
1088 x[:3] = x[-3:] = masked
1089 assert_equal(median(x, axis=-1), median(x, axis=1))
1090
1091 def test_out_1d(self):
1092 # integer float even odd
1093 for v in (30, 30., 31, 31.):
1094 x = masked_array(np.arange(v))
1095 x[:3] = x[-3:] = masked
1096 out = masked_array(np.ones(()))
1097 r = median(x, out=out)
1098 if v == 30:
1099 assert_equal(out, 14.5)
1100 else:
1101 assert_equal(out, 15.)
1102 assert_(r is out)
1103 assert_(type(r) is MaskedArray)
1104
1105 def test_out(self):
1106 # integer float even odd
1107 for v in (40, 40., 30, 30.):
1108 x = masked_array(np.arange(v).reshape(10, -1))
1109 x[:3] = x[-3:] = masked
1110 out = masked_array(np.ones(10))
1111 r = median(x, axis=1, out=out)
1112 if v == 30:
1113 e = masked_array([0.] * 3 + [10, 13, 16, 19] + [0.] * 3,
1114 mask=[True] * 3 + [False] * 4 + [True] * 3)
1115 else:
1116 e = masked_array([0.] * 3 + [13.5, 17.5, 21.5, 25.5] + [0.] * 3,
1117 mask=[True] * 3 + [False] * 4 + [True] * 3)
1118 assert_equal(r, e)
1119 assert_(r is out)
1120 assert_(type(r) is MaskedArray)
1121
1122 @pytest.mark.parametrize(
1123 argnames='axis',
1124 argvalues=[
1125 None,
1126 1,
1127 (1, ),
1128 (0, 1),
1129 (-3, -1),
1130 ]
1131 )
1132 def test_keepdims_out(self, axis):
1133 mask = np.zeros((3, 5, 7, 11), dtype=bool)
1134 # Randomly set some elements to True:
1135 w = np.random.random((4, 200)) * np.array(mask.shape)[:, None]
1136 w = w.astype(np.intp)
1137 mask[tuple(w)] = np.nan
1138 d = masked_array(np.ones(mask.shape), mask=mask)
1139 if axis is None:
1140 shape_out = (1,) * d.ndim
1141 else:
1142 axis_norm = normalize_axis_tuple(axis, d.ndim)
1143 shape_out = tuple(
1144 1 if i in axis_norm else d.shape[i] for i in range(d.ndim))
1145 out = masked_array(np.empty(shape_out))
1146 result = median(d, axis=axis, keepdims=True, out=out)
1147 assert result is out
1148 assert_equal(result.shape, shape_out)
1149
1150 def test_single_non_masked_value_on_axis(self):
1151 data = [[1., 0.],
1152 [0., 3.],
1153 [0., 0.]]
1154 masked_arr = np.ma.masked_equal(data, 0)
1155 expected = [1., 3.]
1156 assert_array_equal(np.ma.median(masked_arr, axis=0),
1157 expected)
1158
1159 def test_nan(self):
1160 for mask in (False, np.zeros(6, dtype=bool)):
1161 dm = np.ma.array([[1, np.nan, 3], [1, 2, 3]])
1162 dm.mask = mask
1163
1164 # scalar result
1165 r = np.ma.median(dm, axis=None)
1166 assert_(np.isscalar(r))
1167 assert_array_equal(r, np.nan)
1168 r = np.ma.median(dm.ravel(), axis=0)
1169 assert_(np.isscalar(r))
1170 assert_array_equal(r, np.nan)
1171
1172 r = np.ma.median(dm, axis=0)
1173 assert_equal(type(r), MaskedArray)
1174 assert_array_equal(r, [1, np.nan, 3])
1175 r = np.ma.median(dm, axis=1)
1176 assert_equal(type(r), MaskedArray)
1177 assert_array_equal(r, [np.nan, 2])
1178 r = np.ma.median(dm, axis=-1)
1179 assert_equal(type(r), MaskedArray)
1180 assert_array_equal(r, [np.nan, 2])
1181
1182 dm = np.ma.array([[1, np.nan, 3], [1, 2, 3]])
1183 dm[:, 2] = np.ma.masked
1184 assert_array_equal(np.ma.median(dm, axis=None), np.nan)
1185 assert_array_equal(np.ma.median(dm, axis=0), [1, np.nan, 3])
1186 assert_array_equal(np.ma.median(dm, axis=1), [np.nan, 1.5])
1187
1188 def test_out_nan(self):
1189 o = np.ma.masked_array(np.zeros((4,)))
1190 d = np.ma.masked_array(np.ones((3, 4)))
1191 d[2, 1] = np.nan
1192 d[2, 2] = np.ma.masked
1193 assert_equal(np.ma.median(d, 0, out=o), o)
1194 o = np.ma.masked_array(np.zeros((3,)))
1195 assert_equal(np.ma.median(d, 1, out=o), o)
1196 o = np.ma.masked_array(np.zeros(()))
1197 assert_equal(np.ma.median(d, out=o), o)
1198
1199 def test_nan_behavior(self):
1200 a = np.ma.masked_array(np.arange(24, dtype=float))