# pylint: disable-msg=W0611, W0612, W0511 """Tests suite for MaskedArray. Adapted from the original test_ma by Pierre Gerard-Marchant :author: Pierre Gerard-Marchant :contact: pierregm_at_uga_dot_edu :version: $Id: test_extras.py 3473 2007-10-29 15:18:13Z jarrod.millman $ """ from __future__ import division, absolute_import, print_function import warnings import itertools import numpy as np from numpy.testing import ( TestCase, run_module_suite, assert_warns, clear_and_catch_warnings ) from numpy.ma.testutils import ( assert_, assert_array_equal, assert_equal, assert_almost_equal ) from numpy.ma.core import ( array, arange, masked, MaskedArray, masked_array, getmaskarray, shape, nomask, ones, zeros, count ) from numpy.ma.extras import ( atleast_1d, atleast_2d, atleast_3d, mr_, dot, polyfit, cov, corrcoef, median, average, unique, setxor1d, setdiff1d, union1d, intersect1d, in1d, ediff1d, apply_over_axes, apply_along_axis, compress_nd, compress_rowcols, mask_rowcols, clump_masked, clump_unmasked, flatnotmasked_contiguous, notmasked_contiguous, notmasked_edges, masked_all, masked_all_like, diagflat ) import numpy.ma.extras as mae class TestGeneric(TestCase): # def test_masked_all(self): # Tests masked_all # Standard dtype test = masked_all((2,), dtype=float) control = array([1, 1], mask=[1, 1], dtype=float) assert_equal(test, control) # Flexible dtype dt = np.dtype({'names': ['a', 'b'], 'formats': ['f', 'f']}) test = masked_all((2,), dtype=dt) control = array([(0, 0), (0, 0)], mask=[(1, 1), (1, 1)], dtype=dt) assert_equal(test, control) test = masked_all((2, 2), dtype=dt) control = array([[(0, 0), (0, 0)], [(0, 0), (0, 0)]], mask=[[(1, 1), (1, 1)], [(1, 1), (1, 1)]], dtype=dt) assert_equal(test, control) # Nested dtype dt = np.dtype([('a', 'f'), ('b', [('ba', 'f'), ('bb', 'f')])]) test = masked_all((2,), dtype=dt) control = array([(1, (1, 1)), (1, (1, 1))], mask=[(1, (1, 1)), (1, (1, 1))], dtype=dt) assert_equal(test, control) test = masked_all((2,), dtype=dt) control = array([(1, (1, 1)), (1, (1, 1))], mask=[(1, (1, 1)), (1, (1, 1))], dtype=dt) assert_equal(test, control) test = masked_all((1, 1), dtype=dt) control = array([[(1, (1, 1))]], mask=[[(1, (1, 1))]], dtype=dt) assert_equal(test, control) def test_masked_all_like(self): # Tests masked_all # Standard dtype base = array([1, 2], dtype=float) test = masked_all_like(base) control = array([1, 1], mask=[1, 1], dtype=float) assert_equal(test, control) # Flexible dtype dt = np.dtype({'names': ['a', 'b'], 'formats': ['f', 'f']}) base = array([(0, 0), (0, 0)], mask=[(1, 1), (1, 1)], dtype=dt) test = masked_all_like(base) control = array([(10, 10), (10, 10)], mask=[(1, 1), (1, 1)], dtype=dt) assert_equal(test, control) # Nested dtype dt = np.dtype([('a', 'f'), ('b', [('ba', 'f'), ('bb', 'f')])]) control = array([(1, (1, 1)), (1, (1, 1))], mask=[(1, (1, 1)), (1, (1, 1))], dtype=dt) test = masked_all_like(control) assert_equal(test, control) def check_clump(self, f): for i in range(1, 7): for j in range(2**i): k = np.arange(i, dtype=int) ja = np.full(i, j, dtype=int) a = masked_array(2**k) a.mask = (ja & (2**k)) != 0 s = 0 for sl in f(a): s += a.data[sl].sum() if f == clump_unmasked: assert_equal(a.compressed().sum(), s) else: a.mask = ~a.mask assert_equal(a.compressed().sum(), s) def test_clump_masked(self): # Test clump_masked a = masked_array(np.arange(10)) a[[0, 1, 2, 6, 8, 9]] = masked # test = clump_masked(a) control = [slice(0, 3), slice(6, 7), slice(8, 10)] assert_equal(test, control) self.check_clump(clump_masked) def test_clump_unmasked(self): # Test clump_unmasked a = masked_array(np.arange(10)) a[[0, 1, 2, 6, 8, 9]] = masked test = clump_unmasked(a) control = [slice(3, 6), slice(7, 8), ] assert_equal(test, control) self.check_clump(clump_unmasked) def test_flatnotmasked_contiguous(self): # Test flatnotmasked_contiguous a = arange(10) # No mask test = flatnotmasked_contiguous(a) assert_equal(test, slice(0, a.size)) # Some mask a[(a < 3) | (a > 8) | (a == 5)] = masked test = flatnotmasked_contiguous(a) assert_equal(test, [slice(3, 5), slice(6, 9)]) # a[:] = masked test = flatnotmasked_contiguous(a) assert_equal(test, None) class TestAverage(TestCase): # Several tests of average. Why so many ? Good point... def test_testAverage1(self): # Test of average. ott = array([0., 1., 2., 3.], mask=[True, False, False, False]) assert_equal(2.0, average(ott, axis=0)) assert_equal(2.0, average(ott, weights=[1., 1., 2., 1.])) result, wts = average(ott, weights=[1., 1., 2., 1.], returned=1) assert_equal(2.0, result) self.assertTrue(wts == 4.0) ott[:] = masked assert_equal(average(ott, axis=0).mask, [True]) ott = array([0., 1., 2., 3.], mask=[True, False, False, False]) ott = ott.reshape(2, 2) ott[:, 1] = masked assert_equal(average(ott, axis=0), [2.0, 0.0]) assert_equal(average(ott, axis=1).mask[0], [True]) assert_equal([2., 0.], average(ott, axis=0)) result, wts = average(ott, axis=0, returned=1) assert_equal(wts, [1., 0.]) def test_testAverage2(self): # More tests of average. w1 = [0, 1, 1, 1, 1, 0] w2 = [[0, 1, 1, 1, 1, 0], [1, 0, 0, 0, 0, 1]] x = arange(6, dtype=np.float_) assert_equal(average(x, axis=0), 2.5) assert_equal(average(x, axis=0, weights=w1), 2.5) y = array([arange(6, dtype=np.float_), 2.0 * arange(6)]) assert_equal(average(y, None), np.add.reduce(np.arange(6)) * 3. / 12.) assert_equal(average(y, axis=0), np.arange(6) * 3. / 2.) assert_equal(average(y, axis=1), [average(x, axis=0), average(x, axis=0) * 2.0]) assert_equal(average(y, None, weights=w2), 20. / 6.) assert_equal(average(y, axis=0, weights=w2), [0., 1., 2., 3., 4., 10.]) assert_equal(average(y, axis=1), [average(x, axis=0), average(x, axis=0) * 2.0]) m1 = zeros(6) m2 = [0, 0, 1, 1, 0, 0] m3 = [[0, 0, 1, 1, 0, 0], [0, 1, 1, 1, 1, 0]] m4 = ones(6) m5 = [0, 1, 1, 1, 1, 1] assert_equal(average(masked_array(x, m1), axis=0), 2.5) assert_equal(average(masked_array(x, m2), axis=0), 2.5) assert_equal(average(masked_array(x, m4), axis=0).mask, [True]) assert_equal(average(masked_array(x, m5), axis=0), 0.0) assert_equal(count(average(masked_array(x, m4), axis=0)), 0) z = masked_array(y, m3) assert_equal(average(z, None), 20. / 6.) assert_equal(average(z, axis=0), [0., 1., 99., 99., 4.0, 7.5]) assert_equal(average(z, axis=1), [2.5, 5.0]) assert_equal(average(z, axis=0, weights=w2), [0., 1., 99., 99., 4.0, 10.0]) def test_testAverage3(self): # Yet more tests of average! a = arange(6) b = arange(6) * 3 r1, w1 = average([[a, b], [b, a]], axis=1, returned=1) assert_equal(shape(r1), shape(w1)) assert_equal(r1.shape, w1.shape) r2, w2 = average(ones((2, 2, 3)), axis=0, weights=[3, 1], returned=1) assert_equal(shape(w2), shape(r2)) r2, w2 = average(ones((2, 2, 3)), returned=1) assert_equal(shape(w2), shape(r2)) r2, w2 = average(ones((2, 2, 3)), weights=ones((2, 2, 3)), returned=1) assert_equal(shape(w2), shape(r2)) a2d = array([[1, 2], [0, 4]], float) a2dm = masked_array(a2d, [[False, False], [True, False]]) a2da = average(a2d, axis=0) assert_equal(a2da, [0.5, 3.0]) a2dma = average(a2dm, axis=0) assert_equal(a2dma, [1.0, 3.0]) a2dma = average(a2dm, axis=None) assert_equal(a2dma, 7. / 3.) a2dma = average(a2dm, axis=1) assert_equal(a2dma, [1.5, 4.0]) def test_onintegers_with_mask(self): # Test average on integers with mask a = average(array([1, 2])) assert_equal(a, 1.5) a = average(array([1, 2, 3, 4], mask=[False, False, True, True])) assert_equal(a, 1.5) def test_complex(self): # Test with complex data. # (Regression test for https://github.com/numpy/numpy/issues/2684) mask = np.array([[0, 0, 0, 1, 0], [0, 1, 0, 0, 0]], dtype=bool) a = masked_array([[0, 1+2j, 3+4j, 5+6j, 7+8j], [9j, 0+1j, 2+3j, 4+5j, 7+7j]], mask=mask) av = average(a) expected = np.average(a.compressed()) assert_almost_equal(av.real, expected.real) assert_almost_equal(av.imag, expected.imag) av0 = average(a, axis=0) expected0 = average(a.real, axis=0) + average(a.imag, axis=0)*1j assert_almost_equal(av0.real, expected0.real) assert_almost_equal(av0.imag, expected0.imag) av1 = average(a, axis=1) expected1 = average(a.real, axis=1) + average(a.imag, axis=1)*1j assert_almost_equal(av1.real, expected1.real) assert_almost_equal(av1.imag, expected1.imag) # Test with the 'weights' argument. wts = np.array([[0.5, 1.0, 2.0, 1.0, 0.5], [1.0, 1.0, 1.0, 1.0, 1.0]]) wav = average(a, weights=wts) expected = np.average(a.compressed(), weights=wts[~mask]) assert_almost_equal(wav.real, expected.real) assert_almost_equal(wav.imag, expected.imag) wav0 = average(a, weights=wts, axis=0) expected0 = (average(a.real, weights=wts, axis=0) + average(a.imag, weights=wts, axis=0)*1j) assert_almost_equal(wav0.real, expected0.real) assert_almost_equal(wav0.imag, expected0.imag) wav1 = average(a, weights=wts, axis=1) expected1 = (average(a.real, weights=wts, axis=1) + average(a.imag, weights=wts, axis=1)*1j) assert_almost_equal(wav1.real, expected1.real) assert_almost_equal(wav1.imag, expected1.imag) class TestConcatenator(TestCase): # Tests for mr_, the equivalent of r_ for masked arrays. def test_1d(self): # Tests mr_ on 1D arrays. assert_array_equal(mr_[1, 2, 3, 4, 5, 6], array([1, 2, 3, 4, 5, 6])) b = ones(5) m = [1, 0, 0, 0, 0] d = masked_array(b, mask=m) c = mr_[d, 0, 0, d] self.assertTrue(isinstance(c, MaskedArray)) assert_array_equal(c, [1, 1, 1, 1, 1, 0, 0, 1, 1, 1, 1, 1]) assert_array_equal(c.mask, mr_[m, 0, 0, m]) def test_2d(self): # Tests mr_ on 2D arrays. a_1 = np.random.rand(5, 5) a_2 = np.random.rand(5, 5) m_1 = np.round_(np.random.rand(5, 5), 0) m_2 = np.round_(np.random.rand(5, 5), 0) b_1 = masked_array(a_1, mask=m_1) b_2 = masked_array(a_2, mask=m_2) # append columns d = mr_['1', b_1, b_2] self.assertTrue(d.shape == (5, 10)) assert_array_equal(d[:, :5], b_1) assert_array_equal(d[:, 5:], b_2) assert_array_equal(d.mask, np.r_['1', m_1, m_2]) d = mr_[b_1, b_2] self.assertTrue(d.shape == (10, 5)) assert_array_equal(d[:5,:], b_1) assert_array_equal(d[5:,:], b_2) assert_array_equal(d.mask, np.r_[m_1, m_2]) class TestNotMasked(TestCase): # Tests notmasked_edges and notmasked_contiguous. def test_edges(self): # Tests unmasked_edges data = masked_array(np.arange(25).reshape(5, 5), mask=[[0, 0, 1, 0, 0], [0, 0, 0, 1, 1], [1, 1, 0, 0, 0], [0, 0, 0, 0, 0], [1, 1, 1, 0, 0]],) test = notmasked_edges(data, None) assert_equal(test, [0, 24]) test = notmasked_edges(data, 0) assert_equal(test[0], [(0, 0, 1, 0, 0), (0, 1, 2, 3, 4)]) assert_equal(test[1], [(3, 3, 3, 4, 4), (0, 1, 2, 3, 4)]) test = notmasked_edges(data, 1) assert_equal(test[0], [(0, 1, 2, 3, 4), (0, 0, 2, 0, 3)]) assert_equal(test[1], [(0, 1, 2, 3, 4), (4, 2, 4, 4, 4)]) # test = notmasked_edges(data.data, None) assert_equal(test, [0, 24]) test = notmasked_edges(data.data, 0) assert_equal(test[0], [(0, 0, 0, 0, 0), (0, 1, 2, 3, 4)]) assert_equal(test[1], [(4, 4, 4, 4, 4), (0, 1, 2, 3, 4)]) test = notmasked_edges(data.data, -1) assert_equal(test[0], [(0, 1, 2, 3, 4), (0, 0, 0, 0, 0)]) assert_equal(test[1], [(0, 1, 2, 3, 4), (4, 4, 4, 4, 4)]) # data[-2] = masked test = notmasked_edges(data, 0) assert_equal(test[0], [(0, 0, 1, 0, 0), (0, 1, 2, 3, 4)]) assert_equal(test[1], [(1, 1, 2, 4, 4), (0, 1, 2, 3, 4)]) test = notmasked_edges(data, -1) assert_equal(test[0], [(0, 1, 2, 4), (0, 0, 2, 3)]) assert_equal(test[1], [(0, 1, 2, 4), (4, 2, 4, 4)]) def test_contiguous(self): # Tests notmasked_contiguous a = masked_array(np.arange(24).reshape(3, 8), mask=[[0, 0, 0, 0, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1], [0, 0, 0, 0, 0, 0, 1, 0], ]) tmp = notmasked_contiguous(a, None) assert_equal(tmp[-1], slice(23, 24, None)) assert_equal(tmp[-2], slice(16, 22, None)) assert_equal(tmp[-3], slice(0, 4, None)) # tmp = notmasked_contiguous(a, 0) self.assertTrue(len(tmp[-1]) == 1) self.assertTrue(tmp[-2] is None) assert_equal(tmp[-3], tmp[-1]) self.assertTrue(len(tmp[0]) == 2) # tmp = notmasked_contiguous(a, 1) assert_equal(tmp[0][-1], slice(0, 4, None)) self.assertTrue(tmp[1] is None) assert_equal(tmp[2][-1], slice(7, 8, None)) assert_equal(tmp[2][-2], slice(0, 6, None)) class TestCompressFunctions(TestCase): def test_compress_nd(self): # Tests compress_nd x = np.array(list(range(3*4*5))).reshape(3, 4, 5) m = np.zeros((3,4,5)).astype(bool) m[1,1,1] = True x = array(x, mask=m) # axis=None a = compress_nd(x) assert_equal(a, [[[ 0, 2, 3, 4], [10, 12, 13, 14], [15, 17, 18, 19]], [[40, 42, 43, 44], [50, 52, 53, 54], [55, 57, 58, 59]]]) # axis=0 a = compress_nd(x, 0) assert_equal(a, [[[ 0, 1, 2, 3, 4], [ 5, 6, 7, 8, 9], [10, 11, 12, 13, 14], [15, 16, 17, 18, 19]], [[40, 41, 42, 43, 44], [45, 46, 47, 48, 49], [50, 51, 52, 53, 54], [55, 56, 57, 58, 59]]]) # axis=1 a = compress_nd(x, 1) assert_equal(a, [[[ 0, 1, 2, 3, 4], [10, 11, 12, 13, 14], [15, 16, 17, 18, 19]], [[20, 21, 22, 23, 24], [30, 31, 32, 33, 34], [35, 36, 37, 38, 39]], [[40, 41, 42, 43, 44], [50, 51, 52, 53, 54], [55, 56, 57, 58, 59]]]) a2 = compress_nd(x, (1,)) a3 = compress_nd(x, -2) a4 = compress_nd(x, (-2,)) assert_equal(a, a2) assert_equal(a, a3) assert_equal(a, a4) # axis=2 a = compress_nd(x, 2) assert_equal(a, [[[ 0, 2, 3, 4], [ 5, 7, 8, 9], [10, 12, 13, 14], [15, 17, 18, 19]], [[20, 22, 23, 24], [25, 27, 28, 29], [30, 32, 33, 34], [35, 37, 38, 39]], [[40, 42, 43, 44], [45, 47, 48, 49], [50, 52, 53, 54], [55, 57, 58, 59]]]) a2 = compress_nd(x, (2,)) a3 = compress_nd(x, -1) a4 = compress_nd(x, (-1,)) assert_equal(a, a2) assert_equal(a, a3) assert_equal(a, a4) # axis=(0, 1) a = compress_nd(x, (0, 1)) assert_equal(a, [[[ 0, 1, 2, 3, 4], [10, 11, 12, 13, 14], [15, 16, 17, 18, 19]], [[40, 41, 42, 43, 44], [50, 51, 52, 53, 54], [55, 56, 57, 58, 59]]]) a2 = compress_nd(x, (0, -2)) assert_equal(a, a2) # axis=(1, 2) a = compress_nd(x, (1, 2)) assert_equal(a, [[[ 0, 2, 3, 4], [10, 12, 13, 14], [15, 17, 18, 19]], [[20, 22, 23, 24], [30, 32, 33, 34], [35, 37, 38, 39]], [[40, 42, 43, 44], [50, 52, 53, 54], [55, 57, 58, 59]]]) a2 = compress_nd(x, (-2, 2)) a3 = compress_nd(x, (1, -1)) a4 = compress_nd(x, (-2, -1)) assert_equal(a, a2) assert_equal(a, a3) assert_equal(a, a4) # axis=(0, 2) a = compress_nd(x, (0, 2)) assert_equal(a, [[[ 0, 2, 3, 4], [ 5, 7, 8, 9], [10, 12, 13, 14], [15, 17, 18, 19]], [[40, 42, 43, 44], [45, 47, 48, 49], [50, 52, 53, 54], [55, 57, 58, 59]]]) a2 = compress_nd(x, (0, -1)) assert_equal(a, a2) def test_compress_rowcols(self): # Tests compress_rowcols x = array(np.arange(9).reshape(3, 3), mask=[[1, 0, 0], [0, 0, 0], [0, 0, 0]]) assert_equal(compress_rowcols(x), [[4, 5], [7, 8]]) assert_equal(compress_rowcols(x, 0), [[3, 4, 5], [6, 7, 8]]) assert_equal(compress_rowcols(x, 1), [[1, 2], [4, 5], [7, 8]]) x = array(x._data, mask=[[0, 0, 0], [0, 1, 0], [0, 0, 0]]) assert_equal(compress_rowcols(x), [[0, 2], [6, 8]]) assert_equal(compress_rowcols(x, 0), [[0, 1, 2], [6, 7, 8]]) assert_equal(compress_rowcols(x, 1), [[0, 2], [3, 5], [6, 8]]) x = array(x._data, mask=[[1, 0, 0], [0, 1, 0], [0, 0, 0]]) assert_equal(compress_rowcols(x), [[8]]) assert_equal(compress_rowcols(x, 0), [[6, 7, 8]]) assert_equal(compress_rowcols(x, 1,), [[2], [5], [8]]) x = array(x._data, mask=[[1, 0, 0], [0, 1, 0], [0, 0, 1]]) assert_equal(compress_rowcols(x).size, 0) assert_equal(compress_rowcols(x, 0).size, 0) assert_equal(compress_rowcols(x, 1).size, 0) def test_mask_rowcols(self): # Tests mask_rowcols. x = array(np.arange(9).reshape(3, 3), mask=[[1, 0, 0], [0, 0, 0], [0, 0, 0]]) assert_equal(mask_rowcols(x).mask, [[1, 1, 1], [1, 0, 0], [1, 0, 0]]) assert_equal(mask_rowcols(x, 0).mask, [[1, 1, 1], [0, 0, 0], [0, 0, 0]]) assert_equal(mask_rowcols(x, 1).mask, [[1, 0, 0], [1, 0, 0], [1, 0, 0]]) x = array(x._data, mask=[[0, 0, 0], [0, 1, 0], [0, 0, 0]]) assert_equal(mask_rowcols(x).mask, [[0, 1, 0], [1, 1, 1], [0, 1, 0]]) assert_equal(mask_rowcols(x, 0).mask, [[0, 0, 0], [1, 1, 1], [0, 0, 0]]) assert_equal(mask_rowcols(x, 1).mask, [[0, 1, 0], [0, 1, 0], [0, 1, 0]]) x = array(x._data, mask=[[1, 0, 0], [0, 1, 0], [0, 0, 0]]) assert_equal(mask_rowcols(x).mask, [[1, 1, 1], [1, 1, 1], [1, 1, 0]]) assert_equal(mask_rowcols(x, 0).mask, [[1, 1, 1], [1, 1, 1], [0, 0, 0]]) assert_equal(mask_rowcols(x, 1,).mask, [[1, 1, 0], [1, 1, 0], [1, 1, 0]]) x = array(x._data, mask=[[1, 0, 0], [0, 1, 0], [0, 0, 1]]) self.assertTrue(mask_rowcols(x).all() is masked) self.assertTrue(mask_rowcols(x, 0).all() is masked) self.assertTrue(mask_rowcols(x, 1).all() is masked) self.assertTrue(mask_rowcols(x).mask.all()) self.assertTrue(mask_rowcols(x, 0).mask.all()) self.assertTrue(mask_rowcols(x, 1).mask.all()) def test_dot(self): # Tests dot product n = np.arange(1, 7) # m = [1, 0, 0, 0, 0, 0] a = masked_array(n, mask=m).reshape(2, 3) b = masked_array(n, mask=m).reshape(3, 2) c = dot(a, b, strict=True) assert_equal(c.mask, [[1, 1], [1, 0]]) c = dot(b, a, strict=True) assert_equal(c.mask, [[1, 1, 1], [1, 0, 0], [1, 0, 0]]) c = dot(a, b, strict=False) assert_equal(c, np.dot(a.filled(0), b.filled(0))) c = dot(b, a, strict=False) assert_equal(c, np.dot(b.filled(0), a.filled(0))) # m = [0, 0, 0, 0, 0, 1] a = masked_array(n, mask=m).reshape(2, 3) b = masked_array(n, mask=m).reshape(3, 2) c = dot(a, b, strict=True) assert_equal(c.mask, [[0, 1], [1, 1]]) c = dot(b, a, strict=True) assert_equal(c.mask, [[0, 0, 1], [0, 0, 1], [1, 1, 1]]) c = dot(a, b, strict=False) assert_equal(c, np.dot(a.filled(0), b.filled(0))) assert_equal(c, dot(a, b)) c = dot(b, a, strict=False) assert_equal(c, np.dot(b.filled(0), a.filled(0))) # m = [0, 0, 0, 0, 0, 0] a = masked_array(n, mask=m).reshape(2, 3) b = masked_array(n, mask=m).reshape(3, 2) c = dot(a, b) assert_equal(c.mask, nomask) c = dot(b, a) assert_equal(c.mask, nomask) # a = masked_array(n, mask=[1, 0, 0, 0, 0, 0]).reshape(2, 3) b = masked_array(n, mask=[0, 0, 0, 0, 0, 0]).reshape(3, 2) c = dot(a, b, strict=True) assert_equal(c.mask, [[1, 1], [0, 0]]) c = dot(a, b, strict=False) assert_equal(c, np.dot(a.filled(0), b.filled(0))) c = dot(b, a, strict=True) assert_equal(c.mask, [[1, 0, 0], [1, 0, 0], [1, 0, 0]]) c = dot(b, a, strict=False) assert_equal(c, np.dot(b.filled(0), a.filled(0))) # a = masked_array(n, mask=[0, 0, 0, 0, 0, 1]).reshape(2, 3) b = masked_array(n, mask=[0, 0, 0, 0, 0, 0]).reshape(3, 2) c = dot(a, b, strict=True) assert_equal(c.mask, [[0, 0], [1, 1]]) c = dot(a, b) assert_equal(c, np.dot(a.filled(0), b.filled(0))) c = dot(b, a, strict=True) assert_equal(c.mask, [[0, 0, 1], [0, 0, 1], [0, 0, 1]]) c = dot(b, a, strict=False) assert_equal(c, np.dot(b.filled(0), a.filled(0))) # a = masked_array(n, mask=[0, 0, 0, 0, 0, 1]).reshape(2, 3) b = masked_array(n, mask=[0, 0, 1, 0, 0, 0]).reshape(3, 2) c = dot(a, b, strict=True) assert_equal(c.mask, [[1, 0], [1, 1]]) c = dot(a, b, strict=False) assert_equal(c, np.dot(a.filled(0), b.filled(0))) c = dot(b, a, strict=True) assert_equal(c.mask, [[0, 0, 1], [1, 1, 1], [0, 0, 1]]) c = dot(b, a, strict=False) assert_equal(c, np.dot(b.filled(0), a.filled(0))) def test_dot_returns_maskedarray(self): # See gh-6611 a = np.eye(3) b = array(a) assert_(type(dot(a, a)) is MaskedArray) assert_(type(dot(a, b)) is MaskedArray) assert_(type(dot(b, a)) is MaskedArray) assert_(type(dot(b, b)) is MaskedArray) def test_dot_out(self): a = array(np.eye(3)) out = array(np.zeros((3, 3))) res = dot(a, a, out=out) assert_(res is out) assert_equal(a, res) class TestApplyAlongAxis(TestCase): # Tests 2D functions def test_3d(self): a = arange(12.).reshape(2, 2, 3) def myfunc(b): return b[1] xa = apply_along_axis(myfunc, 2, a) assert_equal(xa, [[1, 4], [7, 10]]) # Tests kwargs functions def test_3d_kwargs(self): a = arange(12).reshape(2, 2, 3) def myfunc(b, offset=0): return b[1+offset] xa = apply_along_axis(myfunc, 2, a, offset=1) assert_equal(xa, [[2, 5], [8, 11]]) class TestApplyOverAxes(TestCase): # Tests apply_over_axes def test_basic(self): a = arange(24).reshape(2, 3, 4) test = apply_over_axes(np.sum, a, [0, 2]) ctrl = np.array([[[60], [92], [124]]]) assert_equal(test, ctrl) a[(a % 2).astype(np.bool)] = masked test = apply_over_axes(np.sum, a, [0, 2]) ctrl = np.array([[[28], [44], [60]]]) assert_equal(test, ctrl) class TestMedian(TestCase): def test_pytype(self): r = np.ma.median([[np.inf, np.inf], [np.inf, np.inf]], axis=-1) assert_equal(r, np.inf) def test_non_masked(self): x = np.arange(9) assert_equal(np.ma.median(x), 4.) assert_(type(np.ma.median(x)) is not MaskedArray) x = range(9) assert_equal(np.ma.median(x), 4.) assert_(type(np.ma.median(x)) is not MaskedArray) x = 5 assert_equal(np.ma.median(x), 5.) assert_(type(np.ma.median(x)) is not MaskedArray) def test_docstring_examples(self): "test the examples given in the docstring of ma.median" x = array(np.arange(8), mask=[0]*4 + [1]*4) assert_equal(np.ma.median(x), 1.5) assert_equal(np.ma.median(x).shape, (), "shape mismatch") assert_(type(np.ma.median(x)) is not MaskedArray) x = array(np.arange(10).reshape(2, 5), mask=[0]*6 + [1]*4) assert_equal(np.ma.median(x), 2.5) assert_equal(np.ma.median(x).shape, (), "shape mismatch") assert_(type(np.ma.median(x)) is not MaskedArray) ma_x = np.ma.median(x, axis=-1, overwrite_input=True) assert_equal(ma_x, [2., 5.]) assert_equal(ma_x.shape, (2,), "shape mismatch") assert_(type(ma_x) is MaskedArray) def test_axis_argument_errors(self): msg = "mask = %s, ndim = %s, axis = %s, overwrite_input = %s" for ndmin in range(5): for mask in [False, True]: x = array(1, ndmin=ndmin, mask=mask) # Valid axis values should not raise exception args = itertools.product(range(-ndmin, ndmin), [False, True]) for axis, over in args: try: np.ma.median(x, axis=axis, overwrite_input=over) except: raise AssertionError(msg % (mask, ndmin, axis, over)) # Invalid axis values should raise exception args = itertools.product([-(ndmin + 1), ndmin], [False, True]) for axis, over in args: try: np.ma.median(x, axis=axis, overwrite_input=over) except IndexError: pass else: raise AssertionError(msg % (mask, ndmin, axis, over)) def test_masked_0d(self): # Check values x = array(1, mask=False) assert_equal(np.ma.median(x), 1) x = array(1, mask=True) assert_equal(np.ma.median(x), np.ma.masked) def test_masked_1d(self): x = array(np.arange(5), mask=True) assert_equal(np.ma.median(x), np.ma.masked) assert_equal(np.ma.median(x).shape, (), "shape mismatch") assert_(type(np.ma.median(x)) is np.ma.core.MaskedConstant) x = array(np.arange(5), mask=False) assert_equal(np.ma.median(x), 2.) assert_equal(np.ma.median(x).shape, (), "shape mismatch") assert_(type(np.ma.median(x)) is not MaskedArray) x = array(np.arange(5), mask=[0,1,0,0,0]) assert_equal(np.ma.median(x), 2.5) assert_equal(np.ma.median(x).shape, (), "shape mismatch") assert_(type(np.ma.median(x)) is not MaskedArray) x = array(np.arange(5), mask=[0,1,1,1,1]) assert_equal(np.ma.median(x), 0.) assert_equal(np.ma.median(x).shape, (), "shape mismatch") assert_(type(np.ma.median(x)) is not MaskedArray) def test_1d_shape_consistency(self): assert_equal(np.ma.median(array([1,2,3],mask=[0,0,0])).shape, np.ma.median(array([1,2,3],mask=[0,1,0])).shape ) def test_2d(self): # Tests median w/ 2D (n, p) = (101, 30) x = masked_array(np.linspace(-1., 1., n),) x[:10] = x[-10:] = masked z = masked_array(np.empty((n, p), dtype=float)) z[:, 0] = x[:] idx = np.arange(len(x)) for i in range(1, p): np.random.shuffle(idx) z[:, i] = x[idx] assert_equal(median(z[:, 0]), 0) assert_equal(median(z), 0) assert_equal(median(z, axis=0), np.zeros(p)) assert_equal(median(z.T, axis=1), np.zeros(p)) def test_2d_waxis(self): # Tests median w/ 2D arrays and different axis. x = masked_array(np.arange(30).reshape(10, 3)) x[:3] = x[-3:] = masked assert_equal(median(x), 14.5) assert_(type(np.ma.median(x)) is not MaskedArray) assert_equal(median(x, axis=0), [13.5, 14.5, 15.5]) assert_(type(np.ma.median(x, axis=0)) is MaskedArray) assert_equal(median(x, axis=1), [0, 0, 0, 10, 13, 16, 19, 0, 0, 0]) assert_(type(np.ma.median(x, axis=1)) is MaskedArray) assert_equal(median(x, axis=1).mask, [1, 1, 1, 0, 0, 0, 0, 1, 1, 1]) def test_3d(self): # Tests median w/ 3D x = np.ma.arange(24).reshape(3, 4, 2) x[x % 3 == 0] = masked assert_equal(median(x, 0), [[12, 9], [6, 15], [12, 9], [18, 15]]) x.shape = (4, 3, 2) assert_equal(median(x, 0), [[99, 10], [11, 99], [13, 14]]) x = np.ma.arange(24).reshape(4, 3, 2) x[x % 5 == 0] = masked assert_equal(median(x, 0), [[12, 10], [8, 9], [16, 17]]) def test_neg_axis(self): x = masked_array(np.arange(30).reshape(10, 3)) x[:3] = x[-3:] = masked assert_equal(median(x, axis=-1), median(x, axis=1)) def test_out(self): x = masked_array(np.arange(30).reshape(10, 3)) x[:3] = x[-3:] = masked out = masked_array(np.ones(10)) r = median(x, axis=1, out=out) assert_equal(r, out) assert_(type(r) == MaskedArray) def test_single_non_masked_value_on_axis(self): data = [[1., 0.], [0., 3.], [0., 0.]] masked_arr = np.ma.masked_equal(data, 0) expected = [1., 3.] assert_array_equal(np.ma.median(masked_arr, axis=0), expected) class TestCov(TestCase): def setUp(self): self.data = array(np.random.rand(12)) def test_1d_wo_missing(self): # Test cov on 1D variable w/o missing values x = self.data assert_almost_equal(np.cov(x), cov(x)) assert_almost_equal(np.cov(x, rowvar=False), cov(x, rowvar=False)) assert_almost_equal(np.cov(x, rowvar=False, bias=True), cov(x, rowvar=False, bias=True)) def test_2d_wo_missing(self): # Test cov on 1 2D variable w/o missing values x = self.data.reshape(3, 4) assert_almost_equal(np.cov(x), cov(x)) assert_almost_equal(np.cov(x, rowvar=False), cov(x, rowvar=False)) assert_almost_equal(np.cov(x, rowvar=False, bias=True), cov(x, rowvar=False, bias=True)) def test_1d_w_missing(self): # Test cov 1 1D variable w/missing values x = self.data x[-1] = masked x -= x.mean() nx = x.compressed() assert_almost_equal(np.cov(nx), cov(x)) assert_almost_equal(np.cov(nx, rowvar=False), cov(x, rowvar=False)) assert_almost_equal(np.cov(nx, rowvar=False, bias=True), cov(x, rowvar=False, bias=True)) # try: cov(x, allow_masked=False) except ValueError: pass # # 2 1D variables w/ missing values nx = x[1:-1] assert_almost_equal(np.cov(nx, nx[::-1]), cov(x, x[::-1])) assert_almost_equal(np.cov(nx, nx[::-1], rowvar=False), cov(x, x[::-1], rowvar=False)) assert_almost_equal(np.cov(nx, nx[::-1], rowvar=False, bias=True), cov(x, x[::-1], rowvar=False, bias=True)) def test_2d_w_missing(self): # Test cov on 2D variable w/ missing value x = self.data x[-1] = masked x = x.reshape(3, 4) valid = np.logical_not(getmaskarray(x)).astype(int) frac = np.dot(valid, valid.T) xf = (x - x.mean(1)[:, None]).filled(0) assert_almost_equal(cov(x), np.cov(xf) * (x.shape[1] - 1) / (frac - 1.)) assert_almost_equal(cov(x, bias=True), np.cov(xf, bias=True) * x.shape[1] / frac) frac = np.dot(valid.T, valid) xf = (x - x.mean(0)).filled(0) assert_almost_equal(cov(x, rowvar=False), (np.cov(xf, rowvar=False) * (x.shape[0] - 1) / (frac - 1.))) assert_almost_equal(cov(x, rowvar=False, bias=True), (np.cov(xf, rowvar=False, bias=True) * x.shape[0] / frac)) class catch_warn_mae(clear_and_catch_warnings): """ Context manager to catch, reset warnings in ma.extras module """ class_modules = (mae,) class TestCorrcoef(TestCase): def setUp(self): self.data = array(np.random.rand(12)) self.data2 = array(np.random.rand(12)) def test_ddof(self): # ddof raises DeprecationWarning x, y = self.data, self.data2 expected = np.corrcoef(x) expected2 = np.corrcoef(x, y) with catch_warn_mae(): warnings.simplefilter("always") assert_warns(DeprecationWarning, corrcoef, x, ddof=-1) warnings.simplefilter("ignore") # ddof has no or negligible effect on the function assert_almost_equal(np.corrcoef(x, ddof=0), corrcoef(x, ddof=0)) assert_almost_equal(corrcoef(x, ddof=-1), expected) assert_almost_equal(corrcoef(x, y, ddof=-1), expected2) assert_almost_equal(corrcoef(x, ddof=3), expected) assert_almost_equal(corrcoef(x, y, ddof=3), expected2) def test_bias(self): x, y = self.data, self.data2 expected = np.corrcoef(x) # bias raises DeprecationWarning with catch_warn_mae(): warnings.simplefilter("always") assert_warns(DeprecationWarning, corrcoef, x, y, True, False) assert_warns(DeprecationWarning, corrcoef, x, y, True, True) assert_warns(DeprecationWarning, corrcoef, x, bias=False) warnings.simplefilter("ignore") # bias has no or negligible effect on the function assert_almost_equal(corrcoef(x, bias=1), expected) def test_1d_wo_missing(self): # Test cov on 1D variable w/o missing values x = self.data assert_almost_equal(np.corrcoef(x), corrcoef(x)) assert_almost_equal(np.corrcoef(x, rowvar=False), corrcoef(x, rowvar=False)) with catch_warn_mae(): warnings.simplefilter("ignore") assert_almost_equal(np.corrcoef(x, rowvar=False, bias=True), corrcoef(x, rowvar=False, bias=True)) def test_2d_wo_missing(self): # Test corrcoef on 1 2D variable w/o missing values x = self.data.reshape(3, 4) assert_almost_equal(np.corrcoef(x), corrcoef(x)) assert_almost_equal(np.corrcoef(x, rowvar=False), corrcoef(x, rowvar=False)) with catch_warn_mae(): warnings.simplefilter("ignore") assert_almost_equal(np.corrcoef(x, rowvar=False, bias=True), corrcoef(x, rowvar=False, bias=True)) def test_1d_w_missing(self): # Test corrcoef 1 1D variable w/missing values x = self.data x[-1] = masked x -= x.mean() nx = x.compressed() assert_almost_equal(np.corrcoef(nx), corrcoef(x)) assert_almost_equal(np.corrcoef(nx, rowvar=False), corrcoef(x, rowvar=False)) with catch_warn_mae(): warnings.simplefilter("ignore") assert_almost_equal(np.corrcoef(nx, rowvar=False, bias=True), corrcoef(x, rowvar=False, bias=True)) try: corrcoef(x, allow_masked=False) except ValueError: pass # 2 1D variables w/ missing values nx = x[1:-1] assert_almost_equal(np.corrcoef(nx, nx[::-1]), corrcoef(x, x[::-1])) assert_almost_equal(np.corrcoef(nx, nx[::-1], rowvar=False), corrcoef(x, x[::-1], rowvar=False)) with catch_warn_mae(): warnings.simplefilter("ignore") # ddof and bias have no or negligible effect on the function assert_almost_equal(np.corrcoef(nx, nx[::-1]), corrcoef(x, x[::-1], bias=1)) assert_almost_equal(np.corrcoef(nx, nx[::-1]), corrcoef(x, x[::-1], ddof=2)) def test_2d_w_missing(self): # Test corrcoef on 2D variable w/ missing value x = self.data x[-1] = masked x = x.reshape(3, 4) test = corrcoef(x) control = np.corrcoef(x) assert_almost_equal(test[:-1, :-1], control[:-1, :-1]) with catch_warn_mae(): warnings.simplefilter("ignore") # ddof and bias have no or negligible effect on the function assert_almost_equal(corrcoef(x, ddof=-2)[:-1, :-1], control[:-1, :-1]) assert_almost_equal(corrcoef(x, ddof=3)[:-1, :-1], control[:-1, :-1]) assert_almost_equal(corrcoef(x, bias=1)[:-1, :-1], control[:-1, :-1]) class TestPolynomial(TestCase): # def test_polyfit(self): # Tests polyfit # On ndarrays x = np.random.rand(10) y = np.random.rand(20).reshape(-1, 2) assert_almost_equal(polyfit(x, y, 3), np.polyfit(x, y, 3)) # ON 1D maskedarrays x = x.view(MaskedArray) x[0] = masked y = y.view(MaskedArray) y[0, 0] = y[-1, -1] = masked # (C, R, K, S, D) = polyfit(x, y[:, 0], 3, full=True) (c, r, k, s, d) = np.polyfit(x[1:], y[1:, 0].compressed(), 3, full=True) for (a, a_) in zip((C, R, K, S, D), (c, r, k, s, d)): assert_almost_equal(a, a_) # (C, R, K, S, D) = polyfit(x, y[:, -1], 3, full=True) (c, r, k, s, d) = np.polyfit(x[1:-1], y[1:-1, -1], 3, full=True) for (a, a_) in zip((C, R, K, S, D), (c, r, k, s, d)): assert_almost_equal(a, a_) # (C, R, K, S, D) = polyfit(x, y, 3, full=True) (c, r, k, s, d) = np.polyfit(x[1:-1], y[1:-1,:], 3, full=True) for (a, a_) in zip((C, R, K, S, D), (c, r, k, s, d)): assert_almost_equal(a, a_) # w = np.random.rand(10) + 1 wo = w.copy() xs = x[1:-1] ys = y[1:-1] ws = w[1:-1] (C, R, K, S, D) = polyfit(x, y, 3, full=True, w=w) (c, r, k, s, d) = np.polyfit(xs, ys, 3, full=True, w=ws) assert_equal(w, wo) for (a, a_) in zip((C, R, K, S, D), (c, r, k, s, d)): assert_almost_equal(a, a_) def test_polyfit_with_masked_NaNs(self): x = np.random.rand(10) y = np.random.rand(20).reshape(-1, 2) x[0] = np.nan y[-1,-1] = np.nan x = x.view(MaskedArray) y = y.view(MaskedArray) x[0] = masked y[-1,-1] = masked (C, R, K, S, D) = polyfit(x, y, 3, full=True) (c, r, k, s, d) = np.polyfit(x[1:-1], y[1:-1,:], 3, full=True) for (a, a_) in zip((C, R, K, S, D), (c, r, k, s, d)): assert_almost_equal(a, a_) class TestArraySetOps(TestCase): def test_unique_onlist(self): # Test unique on list data = [1, 1, 1, 2, 2, 3] test = unique(data, return_index=True, return_inverse=True) self.assertTrue(isinstance(test[0], MaskedArray)) assert_equal(test[0], masked_array([1, 2, 3], mask=[0, 0, 0])) assert_equal(test[1], [0, 3, 5]) assert_equal(test[2], [0, 0, 0, 1, 1, 2]) def test_unique_onmaskedarray(self): # Test unique on masked data w/use_mask=True data = masked_array([1, 1, 1, 2, 2, 3], mask=[0, 0, 1, 0, 1, 0]) test = unique(data, return_index=True, return_inverse=True) assert_equal(test[0], masked_array([1, 2, 3, -1], mask=[0, 0, 0, 1])) assert_equal(test[1], [0, 3, 5, 2]) assert_equal(test[2], [0, 0, 3, 1, 3, 2]) # data.fill_value = 3 data = masked_array(data=[1, 1, 1, 2, 2, 3], mask=[0, 0, 1, 0, 1, 0], fill_value=3) test = unique(data, return_index=True, return_inverse=True) assert_equal(test[0], masked_array([1, 2, 3, -1], mask=[0, 0, 0, 1])) assert_equal(test[1], [0, 3, 5, 2]) assert_equal(test[2], [0, 0, 3, 1, 3, 2]) def test_unique_allmasked(self): # Test all masked data = masked_array([1, 1, 1], mask=True) test = unique(data, return_index=True, return_inverse=True) assert_equal(test[0], masked_array([1, ], mask=[True])) assert_equal(test[1], [0]) assert_equal(test[2], [0, 0, 0]) # # Test masked data = masked test = unique(data, return_index=True, return_inverse=True) assert_equal(test[0], masked_array(masked)) assert_equal(test[1], [0]) assert_equal(test[2], [0]) def test_ediff1d(self): # Tests mediff1d x = masked_array(np.arange(5), mask=[1, 0, 0, 0, 1]) control = array([1, 1, 1, 4], mask=[1, 0, 0, 1]) test = ediff1d(x) assert_equal(test, control) assert_equal(test.data, control.data) assert_equal(test.mask, control.mask) def test_ediff1d_tobegin(self): # Test ediff1d w/ to_begin x = masked_array(np.arange(5), mask=[1, 0, 0, 0, 1]) test = ediff1d(x, to_begin=masked) control = array([0, 1, 1, 1, 4], mask=[1, 1, 0, 0, 1]) assert_equal(test, control) assert_equal(test.data, control.data) assert_equal(test.mask, control.mask) # test = ediff1d(x, to_begin=[1, 2, 3]) control = array([1, 2, 3, 1, 1, 1, 4], mask=[0, 0, 0, 1, 0, 0, 1]) assert_equal(test, control) assert_equal(test.data, control.data) assert_equal(test.mask, control.mask) def test_ediff1d_toend(self): # Test ediff1d w/ to_end x = masked_array(np.arange(5), mask=[1, 0, 0, 0, 1]) test = ediff1d(x, to_end=masked) control = array([1, 1, 1, 4, 0], mask=[1, 0, 0, 1, 1]) assert_equal(test, control) assert_equal(test.data, control.data) assert_equal(test.mask, control.mask) # test = ediff1d(x, to_end=[1, 2, 3]) control = array([1, 1, 1, 4, 1, 2, 3], mask=[1, 0, 0, 1, 0, 0, 0]) assert_equal(test, control) assert_equal(test.data, control.data) assert_equal(test.mask, control.mask) def test_ediff1d_tobegin_toend(self): # Test ediff1d w/ to_begin and to_end x = masked_array(np.arange(5), mask=[1, 0, 0, 0, 1]) test = ediff1d(x, to_end=masked, to_begin=masked) control = array([0, 1, 1, 1, 4, 0], mask=[1, 1, 0, 0, 1, 1]) assert_equal(test, control) assert_equal(test.data, control.data) assert_equal(test.mask, control.mask) # test = ediff1d(x, to_end=[1, 2, 3], to_begin=masked) control = array([0, 1, 1, 1, 4, 1, 2, 3], mask=[1, 1, 0, 0, 1, 0, 0, 0]) assert_equal(test, control) assert_equal(test.data, control.data) assert_equal(test.mask, control.mask) def test_ediff1d_ndarray(self): # Test ediff1d w/ a ndarray x = np.arange(5) test = ediff1d(x) control = array([1, 1, 1, 1], mask=[0, 0, 0, 0]) assert_equal(test, control) self.assertTrue(isinstance(test, MaskedArray)) assert_equal(test.data, control.data) assert_equal(test.mask, control.mask) # test = ediff1d(x, to_end=masked, to_begin=masked) control = array([0, 1, 1, 1, 1, 0], mask=[1, 0, 0, 0, 0, 1]) self.assertTrue(isinstance(test, MaskedArray)) assert_equal(test.data, control.data) assert_equal(test.mask, control.mask) def test_intersect1d(self): # Test intersect1d x = array([1, 3, 3, 3], mask=[0, 0, 0, 1]) y = array([3, 1, 1, 1], mask=[0, 0, 0, 1]) test = intersect1d(x, y) control = array([1, 3, -1], mask=[0, 0, 1]) assert_equal(test, control) def test_setxor1d(self): # Test setxor1d a = array([1, 2, 5, 7, -1], mask=[0, 0, 0, 0, 1]) b = array([1, 2, 3, 4, 5, -1], mask=[0, 0, 0, 0, 0, 1]) test = setxor1d(a, b) assert_equal(test, array([3, 4, 7])) # a = array([1, 2, 5, 7, -1], mask=[0, 0, 0, 0, 1]) b = [1, 2, 3, 4, 5] test = setxor1d(a, b) assert_equal(test, array([3, 4, 7, -1], mask=[0, 0, 0, 1])) # a = array([1, 2, 3]) b = array([6, 5, 4]) test = setxor1d(a, b) assert_(isinstance(test, MaskedArray)) assert_equal(test, [1, 2, 3, 4, 5, 6]) # a = array([1, 8, 2, 3], mask=[0, 1, 0, 0]) b = array([6, 5, 4, 8], mask=[0, 0, 0, 1]) test = setxor1d(a, b) assert_(isinstance(test, MaskedArray)) assert_equal(test, [1, 2, 3, 4, 5, 6]) # assert_array_equal([], setxor1d([], [])) def test_in1d(self): # Test in1d a = array([1, 2, 5, 7, -1], mask=[0, 0, 0, 0, 1]) b = array([1, 2, 3, 4, 5, -1], mask=[0, 0, 0, 0, 0, 1]) test = in1d(a, b) assert_equal(test, [True, True, True, False, True]) # a = array([5, 5, 2, 1, -1], mask=[0, 0, 0, 0, 1]) b = array([1, 5, -1], mask=[0, 0, 1]) test = in1d(a, b) assert_equal(test, [True, True, False, True, True]) # assert_array_equal([], in1d([], [])) def test_in1d_invert(self): # Test in1d's invert parameter a = array([1, 2, 5, 7, -1], mask=[0, 0, 0, 0, 1]) b = array([1, 2, 3, 4, 5, -1], mask=[0, 0, 0, 0, 0, 1]) assert_equal(np.invert(in1d(a, b)), in1d(a, b, invert=True)) a = array([5, 5, 2, 1, -1], mask=[0, 0, 0, 0, 1]) b = array([1, 5, -1], mask=[0, 0, 1]) assert_equal(np.invert(in1d(a, b)), in1d(a, b, invert=True)) assert_array_equal([], in1d([], [], invert=True)) def test_union1d(self): # Test union1d a = array([1, 2, 5, 7, 5, -1], mask=[0, 0, 0, 0, 0, 1]) b = array([1, 2, 3, 4, 5, -1], mask=[0, 0, 0, 0, 0, 1]) test = union1d(a, b) control = array([1, 2, 3, 4, 5, 7, -1], mask=[0, 0, 0, 0, 0, 0, 1]) assert_equal(test, control) # assert_array_equal([], union1d([], [])) def test_setdiff1d(self): # Test setdiff1d a = array([6, 5, 4, 7, 7, 1, 2, 1], mask=[0, 0, 0, 0, 0, 0, 0, 1]) b = array([2, 4, 3, 3, 2, 1, 5]) test = setdiff1d(a, b) assert_equal(test, array([6, 7, -1], mask=[0, 0, 1])) # a = arange(10) b = arange(8) assert_equal(setdiff1d(a, b), array([8, 9])) a = array([], np.uint32, mask=[]) assert_equal(setdiff1d(a, []).dtype, np.uint32) def test_setdiff1d_char_array(self): # Test setdiff1d_charray a = np.array(['a', 'b', 'c']) b = np.array(['a', 'b', 's']) assert_array_equal(setdiff1d(a, b), np.array(['c'])) class TestShapeBase(TestCase): def test_atleast2d(self): # Test atleast_2d a = masked_array([0, 1, 2], mask=[0, 1, 0]) b = atleast_2d(a) assert_equal(b.shape, (1, 3)) assert_equal(b.mask.shape, b.data.shape) assert_equal(a.shape, (3,)) assert_equal(a.mask.shape, a.data.shape) def test_shape_scalar(self): # the atleast and diagflat function should work with scalars # GitHub issue #3367 b = atleast_1d(1.0) assert_equal(b.shape, (1, )) assert_equal(b.mask.shape, b.data.shape) b = atleast_2d(1.0) assert_equal(b.shape, (1, 1)) assert_equal(b.mask.shape, b.data.shape) b = atleast_3d(1.0) assert_equal(b.shape, (1, 1, 1)) assert_equal(b.mask.shape, b.data.shape) b = diagflat(1.0) assert_equal(b.shape, (1, 1)) assert_equal(b.mask.shape, b.data.shape) if __name__ == "__main__": run_module_suite()