Chris@87: # pylint: disable-msg=W0611, W0612, W0511 Chris@87: """Tests suite for MaskedArray. Chris@87: Adapted from the original test_ma by Pierre Gerard-Marchant Chris@87: Chris@87: :author: Pierre Gerard-Marchant Chris@87: :contact: pierregm_at_uga_dot_edu Chris@87: :version: $Id: test_extras.py 3473 2007-10-29 15:18:13Z jarrod.millman $ Chris@87: Chris@87: """ Chris@87: from __future__ import division, absolute_import, print_function Chris@87: Chris@87: __author__ = "Pierre GF Gerard-Marchant ($Author: jarrod.millman $)" Chris@87: __version__ = '1.0' Chris@87: __revision__ = "$Revision: 3473 $" Chris@87: __date__ = '$Date: 2007-10-29 17:18:13 +0200 (Mon, 29 Oct 2007) $' Chris@87: Chris@87: import numpy as np Chris@87: from numpy.testing import TestCase, run_module_suite Chris@87: from numpy.ma.testutils import (rand, assert_, assert_array_equal, Chris@87: assert_equal, assert_almost_equal) Chris@87: from numpy.ma.core import (array, arange, masked, MaskedArray, masked_array, Chris@87: getmaskarray, shape, nomask, ones, zeros, count) Chris@87: from numpy.ma.extras import ( Chris@87: atleast_2d, mr_, dot, polyfit, Chris@87: cov, corrcoef, median, average, Chris@87: unique, setxor1d, setdiff1d, union1d, intersect1d, in1d, ediff1d, Chris@87: apply_over_axes, apply_along_axis, Chris@87: compress_rowcols, mask_rowcols, Chris@87: clump_masked, clump_unmasked, Chris@87: flatnotmasked_contiguous, notmasked_contiguous, notmasked_edges, Chris@87: masked_all, masked_all_like) Chris@87: Chris@87: Chris@87: class TestGeneric(TestCase): Chris@87: # Chris@87: def test_masked_all(self): Chris@87: # Tests masked_all Chris@87: # Standard dtype Chris@87: test = masked_all((2,), dtype=float) Chris@87: control = array([1, 1], mask=[1, 1], dtype=float) Chris@87: assert_equal(test, control) Chris@87: # Flexible dtype Chris@87: dt = np.dtype({'names': ['a', 'b'], 'formats': ['f', 'f']}) Chris@87: test = masked_all((2,), dtype=dt) Chris@87: control = array([(0, 0), (0, 0)], mask=[(1, 1), (1, 1)], dtype=dt) Chris@87: assert_equal(test, control) Chris@87: test = masked_all((2, 2), dtype=dt) Chris@87: control = array([[(0, 0), (0, 0)], [(0, 0), (0, 0)]], Chris@87: mask=[[(1, 1), (1, 1)], [(1, 1), (1, 1)]], Chris@87: dtype=dt) Chris@87: assert_equal(test, control) Chris@87: # Nested dtype Chris@87: dt = np.dtype([('a', 'f'), ('b', [('ba', 'f'), ('bb', 'f')])]) Chris@87: test = masked_all((2,), dtype=dt) Chris@87: control = array([(1, (1, 1)), (1, (1, 1))], Chris@87: mask=[(1, (1, 1)), (1, (1, 1))], dtype=dt) Chris@87: assert_equal(test, control) Chris@87: test = masked_all((2,), dtype=dt) Chris@87: control = array([(1, (1, 1)), (1, (1, 1))], Chris@87: mask=[(1, (1, 1)), (1, (1, 1))], dtype=dt) Chris@87: assert_equal(test, control) Chris@87: test = masked_all((1, 1), dtype=dt) Chris@87: control = array([[(1, (1, 1))]], mask=[[(1, (1, 1))]], dtype=dt) Chris@87: assert_equal(test, control) Chris@87: Chris@87: def test_masked_all_like(self): Chris@87: # Tests masked_all Chris@87: # Standard dtype Chris@87: base = array([1, 2], dtype=float) Chris@87: test = masked_all_like(base) Chris@87: control = array([1, 1], mask=[1, 1], dtype=float) Chris@87: assert_equal(test, control) Chris@87: # Flexible dtype Chris@87: dt = np.dtype({'names': ['a', 'b'], 'formats': ['f', 'f']}) Chris@87: base = array([(0, 0), (0, 0)], mask=[(1, 1), (1, 1)], dtype=dt) Chris@87: test = masked_all_like(base) Chris@87: control = array([(10, 10), (10, 10)], mask=[(1, 1), (1, 1)], dtype=dt) Chris@87: assert_equal(test, control) Chris@87: # Nested dtype Chris@87: dt = np.dtype([('a', 'f'), ('b', [('ba', 'f'), ('bb', 'f')])]) Chris@87: control = array([(1, (1, 1)), (1, (1, 1))], Chris@87: mask=[(1, (1, 1)), (1, (1, 1))], dtype=dt) Chris@87: test = masked_all_like(control) Chris@87: assert_equal(test, control) Chris@87: Chris@87: def test_clump_masked(self): Chris@87: # Test clump_masked Chris@87: a = masked_array(np.arange(10)) Chris@87: a[[0, 1, 2, 6, 8, 9]] = masked Chris@87: # Chris@87: test = clump_masked(a) Chris@87: control = [slice(0, 3), slice(6, 7), slice(8, 10)] Chris@87: assert_equal(test, control) Chris@87: Chris@87: def test_clump_unmasked(self): Chris@87: # Test clump_unmasked Chris@87: a = masked_array(np.arange(10)) Chris@87: a[[0, 1, 2, 6, 8, 9]] = masked Chris@87: test = clump_unmasked(a) Chris@87: control = [slice(3, 6), slice(7, 8), ] Chris@87: assert_equal(test, control) Chris@87: Chris@87: def test_flatnotmasked_contiguous(self): Chris@87: # Test flatnotmasked_contiguous Chris@87: a = arange(10) Chris@87: # No mask Chris@87: test = flatnotmasked_contiguous(a) Chris@87: assert_equal(test, slice(0, a.size)) Chris@87: # Some mask Chris@87: a[(a < 3) | (a > 8) | (a == 5)] = masked Chris@87: test = flatnotmasked_contiguous(a) Chris@87: assert_equal(test, [slice(3, 5), slice(6, 9)]) Chris@87: # Chris@87: a[:] = masked Chris@87: test = flatnotmasked_contiguous(a) Chris@87: assert_equal(test, None) Chris@87: Chris@87: Chris@87: class TestAverage(TestCase): Chris@87: # Several tests of average. Why so many ? Good point... Chris@87: def test_testAverage1(self): Chris@87: # Test of average. Chris@87: ott = array([0., 1., 2., 3.], mask=[True, False, False, False]) Chris@87: assert_equal(2.0, average(ott, axis=0)) Chris@87: assert_equal(2.0, average(ott, weights=[1., 1., 2., 1.])) Chris@87: result, wts = average(ott, weights=[1., 1., 2., 1.], returned=1) Chris@87: assert_equal(2.0, result) Chris@87: self.assertTrue(wts == 4.0) Chris@87: ott[:] = masked Chris@87: assert_equal(average(ott, axis=0).mask, [True]) Chris@87: ott = array([0., 1., 2., 3.], mask=[True, False, False, False]) Chris@87: ott = ott.reshape(2, 2) Chris@87: ott[:, 1] = masked Chris@87: assert_equal(average(ott, axis=0), [2.0, 0.0]) Chris@87: assert_equal(average(ott, axis=1).mask[0], [True]) Chris@87: assert_equal([2., 0.], average(ott, axis=0)) Chris@87: result, wts = average(ott, axis=0, returned=1) Chris@87: assert_equal(wts, [1., 0.]) Chris@87: Chris@87: def test_testAverage2(self): Chris@87: # More tests of average. Chris@87: w1 = [0, 1, 1, 1, 1, 0] Chris@87: w2 = [[0, 1, 1, 1, 1, 0], [1, 0, 0, 0, 0, 1]] Chris@87: x = arange(6, dtype=np.float_) Chris@87: assert_equal(average(x, axis=0), 2.5) Chris@87: assert_equal(average(x, axis=0, weights=w1), 2.5) Chris@87: y = array([arange(6, dtype=np.float_), 2.0 * arange(6)]) Chris@87: assert_equal(average(y, None), np.add.reduce(np.arange(6)) * 3. / 12.) Chris@87: assert_equal(average(y, axis=0), np.arange(6) * 3. / 2.) Chris@87: assert_equal(average(y, axis=1), Chris@87: [average(x, axis=0), average(x, axis=0) * 2.0]) Chris@87: assert_equal(average(y, None, weights=w2), 20. / 6.) Chris@87: assert_equal(average(y, axis=0, weights=w2), Chris@87: [0., 1., 2., 3., 4., 10.]) Chris@87: assert_equal(average(y, axis=1), Chris@87: [average(x, axis=0), average(x, axis=0) * 2.0]) Chris@87: m1 = zeros(6) Chris@87: m2 = [0, 0, 1, 1, 0, 0] Chris@87: m3 = [[0, 0, 1, 1, 0, 0], [0, 1, 1, 1, 1, 0]] Chris@87: m4 = ones(6) Chris@87: m5 = [0, 1, 1, 1, 1, 1] Chris@87: assert_equal(average(masked_array(x, m1), axis=0), 2.5) Chris@87: assert_equal(average(masked_array(x, m2), axis=0), 2.5) Chris@87: assert_equal(average(masked_array(x, m4), axis=0).mask, [True]) Chris@87: assert_equal(average(masked_array(x, m5), axis=0), 0.0) Chris@87: assert_equal(count(average(masked_array(x, m4), axis=0)), 0) Chris@87: z = masked_array(y, m3) Chris@87: assert_equal(average(z, None), 20. / 6.) Chris@87: assert_equal(average(z, axis=0), [0., 1., 99., 99., 4.0, 7.5]) Chris@87: assert_equal(average(z, axis=1), [2.5, 5.0]) Chris@87: assert_equal(average(z, axis=0, weights=w2), Chris@87: [0., 1., 99., 99., 4.0, 10.0]) Chris@87: Chris@87: def test_testAverage3(self): Chris@87: # Yet more tests of average! Chris@87: a = arange(6) Chris@87: b = arange(6) * 3 Chris@87: r1, w1 = average([[a, b], [b, a]], axis=1, returned=1) Chris@87: assert_equal(shape(r1), shape(w1)) Chris@87: assert_equal(r1.shape, w1.shape) Chris@87: r2, w2 = average(ones((2, 2, 3)), axis=0, weights=[3, 1], returned=1) Chris@87: assert_equal(shape(w2), shape(r2)) Chris@87: r2, w2 = average(ones((2, 2, 3)), returned=1) Chris@87: assert_equal(shape(w2), shape(r2)) Chris@87: r2, w2 = average(ones((2, 2, 3)), weights=ones((2, 2, 3)), returned=1) Chris@87: assert_equal(shape(w2), shape(r2)) Chris@87: a2d = array([[1, 2], [0, 4]], float) Chris@87: a2dm = masked_array(a2d, [[False, False], [True, False]]) Chris@87: a2da = average(a2d, axis=0) Chris@87: assert_equal(a2da, [0.5, 3.0]) Chris@87: a2dma = average(a2dm, axis=0) Chris@87: assert_equal(a2dma, [1.0, 3.0]) Chris@87: a2dma = average(a2dm, axis=None) Chris@87: assert_equal(a2dma, 7. / 3.) Chris@87: a2dma = average(a2dm, axis=1) Chris@87: assert_equal(a2dma, [1.5, 4.0]) Chris@87: Chris@87: def test_onintegers_with_mask(self): Chris@87: # Test average on integers with mask Chris@87: a = average(array([1, 2])) Chris@87: assert_equal(a, 1.5) Chris@87: a = average(array([1, 2, 3, 4], mask=[False, False, True, True])) Chris@87: assert_equal(a, 1.5) Chris@87: Chris@87: def test_complex(self): Chris@87: # Test with complex data. Chris@87: # (Regression test for https://github.com/numpy/numpy/issues/2684) Chris@87: mask = np.array([[0, 0, 0, 1, 0], Chris@87: [0, 1, 0, 0, 0]], dtype=bool) Chris@87: a = masked_array([[0, 1+2j, 3+4j, 5+6j, 7+8j], Chris@87: [9j, 0+1j, 2+3j, 4+5j, 7+7j]], Chris@87: mask=mask) Chris@87: Chris@87: av = average(a) Chris@87: expected = np.average(a.compressed()) Chris@87: assert_almost_equal(av.real, expected.real) Chris@87: assert_almost_equal(av.imag, expected.imag) Chris@87: Chris@87: av0 = average(a, axis=0) Chris@87: expected0 = average(a.real, axis=0) + average(a.imag, axis=0)*1j Chris@87: assert_almost_equal(av0.real, expected0.real) Chris@87: assert_almost_equal(av0.imag, expected0.imag) Chris@87: Chris@87: av1 = average(a, axis=1) Chris@87: expected1 = average(a.real, axis=1) + average(a.imag, axis=1)*1j Chris@87: assert_almost_equal(av1.real, expected1.real) Chris@87: assert_almost_equal(av1.imag, expected1.imag) Chris@87: Chris@87: # Test with the 'weights' argument. Chris@87: wts = np.array([[0.5, 1.0, 2.0, 1.0, 0.5], Chris@87: [1.0, 1.0, 1.0, 1.0, 1.0]]) Chris@87: wav = average(a, weights=wts) Chris@87: expected = np.average(a.compressed(), weights=wts[~mask]) Chris@87: assert_almost_equal(wav.real, expected.real) Chris@87: assert_almost_equal(wav.imag, expected.imag) Chris@87: Chris@87: wav0 = average(a, weights=wts, axis=0) Chris@87: expected0 = (average(a.real, weights=wts, axis=0) + Chris@87: average(a.imag, weights=wts, axis=0)*1j) Chris@87: assert_almost_equal(wav0.real, expected0.real) Chris@87: assert_almost_equal(wav0.imag, expected0.imag) Chris@87: Chris@87: wav1 = average(a, weights=wts, axis=1) Chris@87: expected1 = (average(a.real, weights=wts, axis=1) + Chris@87: average(a.imag, weights=wts, axis=1)*1j) Chris@87: assert_almost_equal(wav1.real, expected1.real) Chris@87: assert_almost_equal(wav1.imag, expected1.imag) Chris@87: Chris@87: Chris@87: class TestConcatenator(TestCase): Chris@87: # Tests for mr_, the equivalent of r_ for masked arrays. Chris@87: Chris@87: def test_1d(self): Chris@87: # Tests mr_ on 1D arrays. Chris@87: assert_array_equal(mr_[1, 2, 3, 4, 5, 6], array([1, 2, 3, 4, 5, 6])) Chris@87: b = ones(5) Chris@87: m = [1, 0, 0, 0, 0] Chris@87: d = masked_array(b, mask=m) Chris@87: c = mr_[d, 0, 0, d] Chris@87: self.assertTrue(isinstance(c, MaskedArray)) Chris@87: assert_array_equal(c, [1, 1, 1, 1, 1, 0, 0, 1, 1, 1, 1, 1]) Chris@87: assert_array_equal(c.mask, mr_[m, 0, 0, m]) Chris@87: Chris@87: def test_2d(self): Chris@87: # Tests mr_ on 2D arrays. Chris@87: a_1 = rand(5, 5) Chris@87: a_2 = rand(5, 5) Chris@87: m_1 = np.round_(rand(5, 5), 0) Chris@87: m_2 = np.round_(rand(5, 5), 0) Chris@87: b_1 = masked_array(a_1, mask=m_1) Chris@87: b_2 = masked_array(a_2, mask=m_2) Chris@87: # append columns Chris@87: d = mr_['1', b_1, b_2] Chris@87: self.assertTrue(d.shape == (5, 10)) Chris@87: assert_array_equal(d[:, :5], b_1) Chris@87: assert_array_equal(d[:, 5:], b_2) Chris@87: assert_array_equal(d.mask, np.r_['1', m_1, m_2]) Chris@87: d = mr_[b_1, b_2] Chris@87: self.assertTrue(d.shape == (10, 5)) Chris@87: assert_array_equal(d[:5,:], b_1) Chris@87: assert_array_equal(d[5:,:], b_2) Chris@87: assert_array_equal(d.mask, np.r_[m_1, m_2]) Chris@87: Chris@87: Chris@87: class TestNotMasked(TestCase): Chris@87: # Tests notmasked_edges and notmasked_contiguous. Chris@87: Chris@87: def test_edges(self): Chris@87: # Tests unmasked_edges Chris@87: data = masked_array(np.arange(25).reshape(5, 5), Chris@87: mask=[[0, 0, 1, 0, 0], Chris@87: [0, 0, 0, 1, 1], Chris@87: [1, 1, 0, 0, 0], Chris@87: [0, 0, 0, 0, 0], Chris@87: [1, 1, 1, 0, 0]],) Chris@87: test = notmasked_edges(data, None) Chris@87: assert_equal(test, [0, 24]) Chris@87: test = notmasked_edges(data, 0) Chris@87: assert_equal(test[0], [(0, 0, 1, 0, 0), (0, 1, 2, 3, 4)]) Chris@87: assert_equal(test[1], [(3, 3, 3, 4, 4), (0, 1, 2, 3, 4)]) Chris@87: test = notmasked_edges(data, 1) Chris@87: assert_equal(test[0], [(0, 1, 2, 3, 4), (0, 0, 2, 0, 3)]) Chris@87: assert_equal(test[1], [(0, 1, 2, 3, 4), (4, 2, 4, 4, 4)]) Chris@87: # Chris@87: test = notmasked_edges(data.data, None) Chris@87: assert_equal(test, [0, 24]) Chris@87: test = notmasked_edges(data.data, 0) Chris@87: assert_equal(test[0], [(0, 0, 0, 0, 0), (0, 1, 2, 3, 4)]) Chris@87: assert_equal(test[1], [(4, 4, 4, 4, 4), (0, 1, 2, 3, 4)]) Chris@87: test = notmasked_edges(data.data, -1) Chris@87: assert_equal(test[0], [(0, 1, 2, 3, 4), (0, 0, 0, 0, 0)]) Chris@87: assert_equal(test[1], [(0, 1, 2, 3, 4), (4, 4, 4, 4, 4)]) Chris@87: # Chris@87: data[-2] = masked Chris@87: test = notmasked_edges(data, 0) Chris@87: assert_equal(test[0], [(0, 0, 1, 0, 0), (0, 1, 2, 3, 4)]) Chris@87: assert_equal(test[1], [(1, 1, 2, 4, 4), (0, 1, 2, 3, 4)]) Chris@87: test = notmasked_edges(data, -1) Chris@87: assert_equal(test[0], [(0, 1, 2, 4), (0, 0, 2, 3)]) Chris@87: assert_equal(test[1], [(0, 1, 2, 4), (4, 2, 4, 4)]) Chris@87: Chris@87: def test_contiguous(self): Chris@87: # Tests notmasked_contiguous Chris@87: a = masked_array(np.arange(24).reshape(3, 8), Chris@87: mask=[[0, 0, 0, 0, 1, 1, 1, 1], Chris@87: [1, 1, 1, 1, 1, 1, 1, 1], Chris@87: [0, 0, 0, 0, 0, 0, 1, 0], ]) Chris@87: tmp = notmasked_contiguous(a, None) Chris@87: assert_equal(tmp[-1], slice(23, 24, None)) Chris@87: assert_equal(tmp[-2], slice(16, 22, None)) Chris@87: assert_equal(tmp[-3], slice(0, 4, None)) Chris@87: # Chris@87: tmp = notmasked_contiguous(a, 0) Chris@87: self.assertTrue(len(tmp[-1]) == 1) Chris@87: self.assertTrue(tmp[-2] is None) Chris@87: assert_equal(tmp[-3], tmp[-1]) Chris@87: self.assertTrue(len(tmp[0]) == 2) Chris@87: # Chris@87: tmp = notmasked_contiguous(a, 1) Chris@87: assert_equal(tmp[0][-1], slice(0, 4, None)) Chris@87: self.assertTrue(tmp[1] is None) Chris@87: assert_equal(tmp[2][-1], slice(7, 8, None)) Chris@87: assert_equal(tmp[2][-2], slice(0, 6, None)) Chris@87: Chris@87: Chris@87: class Test2DFunctions(TestCase): Chris@87: # Tests 2D functions Chris@87: def test_compress2d(self): Chris@87: # Tests compress2d Chris@87: x = array(np.arange(9).reshape(3, 3), Chris@87: mask=[[1, 0, 0], [0, 0, 0], [0, 0, 0]]) Chris@87: assert_equal(compress_rowcols(x), [[4, 5], [7, 8]]) Chris@87: assert_equal(compress_rowcols(x, 0), [[3, 4, 5], [6, 7, 8]]) Chris@87: assert_equal(compress_rowcols(x, 1), [[1, 2], [4, 5], [7, 8]]) Chris@87: x = array(x._data, mask=[[0, 0, 0], [0, 1, 0], [0, 0, 0]]) Chris@87: assert_equal(compress_rowcols(x), [[0, 2], [6, 8]]) Chris@87: assert_equal(compress_rowcols(x, 0), [[0, 1, 2], [6, 7, 8]]) Chris@87: assert_equal(compress_rowcols(x, 1), [[0, 2], [3, 5], [6, 8]]) Chris@87: x = array(x._data, mask=[[1, 0, 0], [0, 1, 0], [0, 0, 0]]) Chris@87: assert_equal(compress_rowcols(x), [[8]]) Chris@87: assert_equal(compress_rowcols(x, 0), [[6, 7, 8]]) Chris@87: assert_equal(compress_rowcols(x, 1,), [[2], [5], [8]]) Chris@87: x = array(x._data, mask=[[1, 0, 0], [0, 1, 0], [0, 0, 1]]) Chris@87: assert_equal(compress_rowcols(x).size, 0) Chris@87: assert_equal(compress_rowcols(x, 0).size, 0) Chris@87: assert_equal(compress_rowcols(x, 1).size, 0) Chris@87: Chris@87: def test_mask_rowcols(self): Chris@87: # Tests mask_rowcols. Chris@87: x = array(np.arange(9).reshape(3, 3), Chris@87: mask=[[1, 0, 0], [0, 0, 0], [0, 0, 0]]) Chris@87: assert_equal(mask_rowcols(x).mask, Chris@87: [[1, 1, 1], [1, 0, 0], [1, 0, 0]]) Chris@87: assert_equal(mask_rowcols(x, 0).mask, Chris@87: [[1, 1, 1], [0, 0, 0], [0, 0, 0]]) Chris@87: assert_equal(mask_rowcols(x, 1).mask, Chris@87: [[1, 0, 0], [1, 0, 0], [1, 0, 0]]) Chris@87: x = array(x._data, mask=[[0, 0, 0], [0, 1, 0], [0, 0, 0]]) Chris@87: assert_equal(mask_rowcols(x).mask, Chris@87: [[0, 1, 0], [1, 1, 1], [0, 1, 0]]) Chris@87: assert_equal(mask_rowcols(x, 0).mask, Chris@87: [[0, 0, 0], [1, 1, 1], [0, 0, 0]]) Chris@87: assert_equal(mask_rowcols(x, 1).mask, Chris@87: [[0, 1, 0], [0, 1, 0], [0, 1, 0]]) Chris@87: x = array(x._data, mask=[[1, 0, 0], [0, 1, 0], [0, 0, 0]]) Chris@87: assert_equal(mask_rowcols(x).mask, Chris@87: [[1, 1, 1], [1, 1, 1], [1, 1, 0]]) Chris@87: assert_equal(mask_rowcols(x, 0).mask, Chris@87: [[1, 1, 1], [1, 1, 1], [0, 0, 0]]) Chris@87: assert_equal(mask_rowcols(x, 1,).mask, Chris@87: [[1, 1, 0], [1, 1, 0], [1, 1, 0]]) Chris@87: x = array(x._data, mask=[[1, 0, 0], [0, 1, 0], [0, 0, 1]]) Chris@87: self.assertTrue(mask_rowcols(x).all() is masked) Chris@87: self.assertTrue(mask_rowcols(x, 0).all() is masked) Chris@87: self.assertTrue(mask_rowcols(x, 1).all() is masked) Chris@87: self.assertTrue(mask_rowcols(x).mask.all()) Chris@87: self.assertTrue(mask_rowcols(x, 0).mask.all()) Chris@87: self.assertTrue(mask_rowcols(x, 1).mask.all()) Chris@87: Chris@87: def test_dot(self): Chris@87: # Tests dot product Chris@87: n = np.arange(1, 7) Chris@87: # Chris@87: m = [1, 0, 0, 0, 0, 0] Chris@87: a = masked_array(n, mask=m).reshape(2, 3) Chris@87: b = masked_array(n, mask=m).reshape(3, 2) Chris@87: c = dot(a, b, True) Chris@87: assert_equal(c.mask, [[1, 1], [1, 0]]) Chris@87: c = dot(b, a, True) Chris@87: assert_equal(c.mask, [[1, 1, 1], [1, 0, 0], [1, 0, 0]]) Chris@87: c = dot(a, b, False) Chris@87: assert_equal(c, np.dot(a.filled(0), b.filled(0))) Chris@87: c = dot(b, a, False) Chris@87: assert_equal(c, np.dot(b.filled(0), a.filled(0))) Chris@87: # Chris@87: m = [0, 0, 0, 0, 0, 1] Chris@87: a = masked_array(n, mask=m).reshape(2, 3) Chris@87: b = masked_array(n, mask=m).reshape(3, 2) Chris@87: c = dot(a, b, True) Chris@87: assert_equal(c.mask, [[0, 1], [1, 1]]) Chris@87: c = dot(b, a, True) Chris@87: assert_equal(c.mask, [[0, 0, 1], [0, 0, 1], [1, 1, 1]]) Chris@87: c = dot(a, b, False) Chris@87: assert_equal(c, np.dot(a.filled(0), b.filled(0))) Chris@87: assert_equal(c, dot(a, b)) Chris@87: c = dot(b, a, False) Chris@87: assert_equal(c, np.dot(b.filled(0), a.filled(0))) Chris@87: # Chris@87: m = [0, 0, 0, 0, 0, 0] Chris@87: a = masked_array(n, mask=m).reshape(2, 3) Chris@87: b = masked_array(n, mask=m).reshape(3, 2) Chris@87: c = dot(a, b) Chris@87: assert_equal(c.mask, nomask) Chris@87: c = dot(b, a) Chris@87: assert_equal(c.mask, nomask) Chris@87: # Chris@87: a = masked_array(n, mask=[1, 0, 0, 0, 0, 0]).reshape(2, 3) Chris@87: b = masked_array(n, mask=[0, 0, 0, 0, 0, 0]).reshape(3, 2) Chris@87: c = dot(a, b, True) Chris@87: assert_equal(c.mask, [[1, 1], [0, 0]]) Chris@87: c = dot(a, b, False) Chris@87: assert_equal(c, np.dot(a.filled(0), b.filled(0))) Chris@87: c = dot(b, a, True) Chris@87: assert_equal(c.mask, [[1, 0, 0], [1, 0, 0], [1, 0, 0]]) Chris@87: c = dot(b, a, False) Chris@87: assert_equal(c, np.dot(b.filled(0), a.filled(0))) Chris@87: # Chris@87: a = masked_array(n, mask=[0, 0, 0, 0, 0, 1]).reshape(2, 3) Chris@87: b = masked_array(n, mask=[0, 0, 0, 0, 0, 0]).reshape(3, 2) Chris@87: c = dot(a, b, True) Chris@87: assert_equal(c.mask, [[0, 0], [1, 1]]) Chris@87: c = dot(a, b) Chris@87: assert_equal(c, np.dot(a.filled(0), b.filled(0))) Chris@87: c = dot(b, a, True) Chris@87: assert_equal(c.mask, [[0, 0, 1], [0, 0, 1], [0, 0, 1]]) Chris@87: c = dot(b, a, False) Chris@87: assert_equal(c, np.dot(b.filled(0), a.filled(0))) Chris@87: # Chris@87: a = masked_array(n, mask=[0, 0, 0, 0, 0, 1]).reshape(2, 3) Chris@87: b = masked_array(n, mask=[0, 0, 1, 0, 0, 0]).reshape(3, 2) Chris@87: c = dot(a, b, True) Chris@87: assert_equal(c.mask, [[1, 0], [1, 1]]) Chris@87: c = dot(a, b, False) Chris@87: assert_equal(c, np.dot(a.filled(0), b.filled(0))) Chris@87: c = dot(b, a, True) Chris@87: assert_equal(c.mask, [[0, 0, 1], [1, 1, 1], [0, 0, 1]]) Chris@87: c = dot(b, a, False) Chris@87: assert_equal(c, np.dot(b.filled(0), a.filled(0))) Chris@87: Chris@87: Chris@87: class TestApplyAlongAxis(TestCase): Chris@87: # Tests 2D functions Chris@87: def test_3d(self): Chris@87: a = arange(12.).reshape(2, 2, 3) Chris@87: Chris@87: def myfunc(b): Chris@87: return b[1] Chris@87: Chris@87: xa = apply_along_axis(myfunc, 2, a) Chris@87: assert_equal(xa, [[1, 4], [7, 10]]) Chris@87: Chris@87: # Tests kwargs functions Chris@87: def test_3d_kwargs(self): Chris@87: a = arange(12).reshape(2, 2, 3) Chris@87: Chris@87: def myfunc(b, offset=0): Chris@87: return b[1+offset] Chris@87: Chris@87: xa = apply_along_axis(myfunc, 2, a, offset=1) Chris@87: assert_equal(xa, [[2, 5], [8, 11]]) Chris@87: Chris@87: Chris@87: class TestApplyOverAxes(TestCase): Chris@87: # Tests apply_over_axes Chris@87: def test_basic(self): Chris@87: a = arange(24).reshape(2, 3, 4) Chris@87: test = apply_over_axes(np.sum, a, [0, 2]) Chris@87: ctrl = np.array([[[60], [92], [124]]]) Chris@87: assert_equal(test, ctrl) Chris@87: a[(a % 2).astype(np.bool)] = masked Chris@87: test = apply_over_axes(np.sum, a, [0, 2]) Chris@87: ctrl = np.array([[[28], [44], [60]]]) Chris@87: assert_equal(test, ctrl) Chris@87: Chris@87: Chris@87: class TestMedian(TestCase): Chris@87: def test_pytype(self): Chris@87: r = np.ma.median([[np.inf, np.inf], [np.inf, np.inf]], axis=-1) Chris@87: assert_equal(r, np.inf) Chris@87: Chris@87: def test_2d(self): Chris@87: # Tests median w/ 2D Chris@87: (n, p) = (101, 30) Chris@87: x = masked_array(np.linspace(-1., 1., n),) Chris@87: x[:10] = x[-10:] = masked Chris@87: z = masked_array(np.empty((n, p), dtype=float)) Chris@87: z[:, 0] = x[:] Chris@87: idx = np.arange(len(x)) Chris@87: for i in range(1, p): Chris@87: np.random.shuffle(idx) Chris@87: z[:, i] = x[idx] Chris@87: assert_equal(median(z[:, 0]), 0) Chris@87: assert_equal(median(z), 0) Chris@87: assert_equal(median(z, axis=0), np.zeros(p)) Chris@87: assert_equal(median(z.T, axis=1), np.zeros(p)) Chris@87: Chris@87: def test_2d_waxis(self): Chris@87: # Tests median w/ 2D arrays and different axis. Chris@87: x = masked_array(np.arange(30).reshape(10, 3)) Chris@87: x[:3] = x[-3:] = masked Chris@87: assert_equal(median(x), 14.5) Chris@87: assert_equal(median(x, axis=0), [13.5, 14.5, 15.5]) Chris@87: assert_equal(median(x, axis=1), [0, 0, 0, 10, 13, 16, 19, 0, 0, 0]) Chris@87: assert_equal(median(x, axis=1).mask, [1, 1, 1, 0, 0, 0, 0, 1, 1, 1]) Chris@87: Chris@87: def test_3d(self): Chris@87: # Tests median w/ 3D Chris@87: x = np.ma.arange(24).reshape(3, 4, 2) Chris@87: x[x % 3 == 0] = masked Chris@87: assert_equal(median(x, 0), [[12, 9], [6, 15], [12, 9], [18, 15]]) Chris@87: x.shape = (4, 3, 2) Chris@87: assert_equal(median(x, 0), [[99, 10], [11, 99], [13, 14]]) Chris@87: x = np.ma.arange(24).reshape(4, 3, 2) Chris@87: x[x % 5 == 0] = masked Chris@87: assert_equal(median(x, 0), [[12, 10], [8, 9], [16, 17]]) Chris@87: Chris@87: def test_neg_axis(self): Chris@87: x = masked_array(np.arange(30).reshape(10, 3)) Chris@87: x[:3] = x[-3:] = masked Chris@87: assert_equal(median(x, axis=-1), median(x, axis=1)) Chris@87: Chris@87: def test_out(self): Chris@87: x = masked_array(np.arange(30).reshape(10, 3)) Chris@87: x[:3] = x[-3:] = masked Chris@87: out = masked_array(np.ones(10)) Chris@87: r = median(x, axis=1, out=out) Chris@87: assert_equal(r, out) Chris@87: assert_(type(r) == MaskedArray) Chris@87: Chris@87: Chris@87: class TestCov(TestCase): Chris@87: Chris@87: def setUp(self): Chris@87: self.data = array(np.random.rand(12)) Chris@87: Chris@87: def test_1d_wo_missing(self): Chris@87: # Test cov on 1D variable w/o missing values Chris@87: x = self.data Chris@87: assert_almost_equal(np.cov(x), cov(x)) Chris@87: assert_almost_equal(np.cov(x, rowvar=False), cov(x, rowvar=False)) Chris@87: assert_almost_equal(np.cov(x, rowvar=False, bias=True), Chris@87: cov(x, rowvar=False, bias=True)) Chris@87: Chris@87: def test_2d_wo_missing(self): Chris@87: # Test cov on 1 2D variable w/o missing values Chris@87: x = self.data.reshape(3, 4) Chris@87: assert_almost_equal(np.cov(x), cov(x)) Chris@87: assert_almost_equal(np.cov(x, rowvar=False), cov(x, rowvar=False)) Chris@87: assert_almost_equal(np.cov(x, rowvar=False, bias=True), Chris@87: cov(x, rowvar=False, bias=True)) Chris@87: Chris@87: def test_1d_w_missing(self): Chris@87: # Test cov 1 1D variable w/missing values Chris@87: x = self.data Chris@87: x[-1] = masked Chris@87: x -= x.mean() Chris@87: nx = x.compressed() Chris@87: assert_almost_equal(np.cov(nx), cov(x)) Chris@87: assert_almost_equal(np.cov(nx, rowvar=False), cov(x, rowvar=False)) Chris@87: assert_almost_equal(np.cov(nx, rowvar=False, bias=True), Chris@87: cov(x, rowvar=False, bias=True)) Chris@87: # Chris@87: try: Chris@87: cov(x, allow_masked=False) Chris@87: except ValueError: Chris@87: pass Chris@87: # Chris@87: # 2 1D variables w/ missing values Chris@87: nx = x[1:-1] Chris@87: assert_almost_equal(np.cov(nx, nx[::-1]), cov(x, x[::-1])) Chris@87: assert_almost_equal(np.cov(nx, nx[::-1], rowvar=False), Chris@87: cov(x, x[::-1], rowvar=False)) Chris@87: assert_almost_equal(np.cov(nx, nx[::-1], rowvar=False, bias=True), Chris@87: cov(x, x[::-1], rowvar=False, bias=True)) Chris@87: Chris@87: def test_2d_w_missing(self): Chris@87: # Test cov on 2D variable w/ missing value Chris@87: x = self.data Chris@87: x[-1] = masked Chris@87: x = x.reshape(3, 4) Chris@87: valid = np.logical_not(getmaskarray(x)).astype(int) Chris@87: frac = np.dot(valid, valid.T) Chris@87: xf = (x - x.mean(1)[:, None]).filled(0) Chris@87: assert_almost_equal(cov(x), Chris@87: np.cov(xf) * (x.shape[1] - 1) / (frac - 1.)) Chris@87: assert_almost_equal(cov(x, bias=True), Chris@87: np.cov(xf, bias=True) * x.shape[1] / frac) Chris@87: frac = np.dot(valid.T, valid) Chris@87: xf = (x - x.mean(0)).filled(0) Chris@87: assert_almost_equal(cov(x, rowvar=False), Chris@87: (np.cov(xf, rowvar=False) * Chris@87: (x.shape[0] - 1) / (frac - 1.))) Chris@87: assert_almost_equal(cov(x, rowvar=False, bias=True), Chris@87: (np.cov(xf, rowvar=False, bias=True) * Chris@87: x.shape[0] / frac)) Chris@87: Chris@87: Chris@87: class TestCorrcoef(TestCase): Chris@87: Chris@87: def setUp(self): Chris@87: self.data = array(np.random.rand(12)) Chris@87: Chris@87: def test_ddof(self): Chris@87: # Test ddof keyword Chris@87: x = self.data Chris@87: assert_almost_equal(np.corrcoef(x, ddof=0), corrcoef(x, ddof=0)) Chris@87: Chris@87: def test_1d_wo_missing(self): Chris@87: # Test cov on 1D variable w/o missing values Chris@87: x = self.data Chris@87: assert_almost_equal(np.corrcoef(x), corrcoef(x)) Chris@87: assert_almost_equal(np.corrcoef(x, rowvar=False), Chris@87: corrcoef(x, rowvar=False)) Chris@87: assert_almost_equal(np.corrcoef(x, rowvar=False, bias=True), Chris@87: corrcoef(x, rowvar=False, bias=True)) Chris@87: Chris@87: def test_2d_wo_missing(self): Chris@87: # Test corrcoef on 1 2D variable w/o missing values Chris@87: x = self.data.reshape(3, 4) Chris@87: assert_almost_equal(np.corrcoef(x), corrcoef(x)) Chris@87: assert_almost_equal(np.corrcoef(x, rowvar=False), Chris@87: corrcoef(x, rowvar=False)) Chris@87: assert_almost_equal(np.corrcoef(x, rowvar=False, bias=True), Chris@87: corrcoef(x, rowvar=False, bias=True)) Chris@87: Chris@87: def test_1d_w_missing(self): Chris@87: # Test corrcoef 1 1D variable w/missing values Chris@87: x = self.data Chris@87: x[-1] = masked Chris@87: x -= x.mean() Chris@87: nx = x.compressed() Chris@87: assert_almost_equal(np.corrcoef(nx), corrcoef(x)) Chris@87: assert_almost_equal(np.corrcoef(nx, rowvar=False), Chris@87: corrcoef(x, rowvar=False)) Chris@87: assert_almost_equal(np.corrcoef(nx, rowvar=False, bias=True), Chris@87: corrcoef(x, rowvar=False, bias=True)) Chris@87: # Chris@87: try: Chris@87: corrcoef(x, allow_masked=False) Chris@87: except ValueError: Chris@87: pass Chris@87: # Chris@87: # 2 1D variables w/ missing values Chris@87: nx = x[1:-1] Chris@87: assert_almost_equal(np.corrcoef(nx, nx[::-1]), corrcoef(x, x[::-1])) Chris@87: assert_almost_equal(np.corrcoef(nx, nx[::-1], rowvar=False), Chris@87: corrcoef(x, x[::-1], rowvar=False)) Chris@87: assert_almost_equal(np.corrcoef(nx, nx[::-1], rowvar=False, bias=True), Chris@87: corrcoef(x, x[::-1], rowvar=False, bias=True)) Chris@87: Chris@87: def test_2d_w_missing(self): Chris@87: # Test corrcoef on 2D variable w/ missing value Chris@87: x = self.data Chris@87: x[-1] = masked Chris@87: x = x.reshape(3, 4) Chris@87: Chris@87: test = corrcoef(x) Chris@87: control = np.corrcoef(x) Chris@87: assert_almost_equal(test[:-1, :-1], control[:-1, :-1]) Chris@87: Chris@87: Chris@87: class TestPolynomial(TestCase): Chris@87: # Chris@87: def test_polyfit(self): Chris@87: # Tests polyfit Chris@87: # On ndarrays Chris@87: x = np.random.rand(10) Chris@87: y = np.random.rand(20).reshape(-1, 2) Chris@87: assert_almost_equal(polyfit(x, y, 3), np.polyfit(x, y, 3)) Chris@87: # ON 1D maskedarrays Chris@87: x = x.view(MaskedArray) Chris@87: x[0] = masked Chris@87: y = y.view(MaskedArray) Chris@87: y[0, 0] = y[-1, -1] = masked Chris@87: # Chris@87: (C, R, K, S, D) = polyfit(x, y[:, 0], 3, full=True) Chris@87: (c, r, k, s, d) = np.polyfit(x[1:], y[1:, 0].compressed(), 3, Chris@87: full=True) Chris@87: for (a, a_) in zip((C, R, K, S, D), (c, r, k, s, d)): Chris@87: assert_almost_equal(a, a_) Chris@87: # Chris@87: (C, R, K, S, D) = polyfit(x, y[:, -1], 3, full=True) Chris@87: (c, r, k, s, d) = np.polyfit(x[1:-1], y[1:-1, -1], 3, full=True) Chris@87: for (a, a_) in zip((C, R, K, S, D), (c, r, k, s, d)): Chris@87: assert_almost_equal(a, a_) Chris@87: # Chris@87: (C, R, K, S, D) = polyfit(x, y, 3, full=True) Chris@87: (c, r, k, s, d) = np.polyfit(x[1:-1], y[1:-1,:], 3, full=True) Chris@87: for (a, a_) in zip((C, R, K, S, D), (c, r, k, s, d)): Chris@87: assert_almost_equal(a, a_) Chris@87: # Chris@87: w = np.random.rand(10) + 1 Chris@87: wo = w.copy() Chris@87: xs = x[1:-1] Chris@87: ys = y[1:-1] Chris@87: ws = w[1:-1] Chris@87: (C, R, K, S, D) = polyfit(x, y, 3, full=True, w=w) Chris@87: (c, r, k, s, d) = np.polyfit(xs, ys, 3, full=True, w=ws) Chris@87: assert_equal(w, wo) Chris@87: for (a, a_) in zip((C, R, K, S, D), (c, r, k, s, d)): Chris@87: assert_almost_equal(a, a_) Chris@87: Chris@87: Chris@87: class TestArraySetOps(TestCase): Chris@87: Chris@87: def test_unique_onlist(self): Chris@87: # Test unique on list Chris@87: data = [1, 1, 1, 2, 2, 3] Chris@87: test = unique(data, return_index=True, return_inverse=True) Chris@87: self.assertTrue(isinstance(test[0], MaskedArray)) Chris@87: assert_equal(test[0], masked_array([1, 2, 3], mask=[0, 0, 0])) Chris@87: assert_equal(test[1], [0, 3, 5]) Chris@87: assert_equal(test[2], [0, 0, 0, 1, 1, 2]) Chris@87: Chris@87: def test_unique_onmaskedarray(self): Chris@87: # Test unique on masked data w/use_mask=True Chris@87: data = masked_array([1, 1, 1, 2, 2, 3], mask=[0, 0, 1, 0, 1, 0]) Chris@87: test = unique(data, return_index=True, return_inverse=True) Chris@87: assert_equal(test[0], masked_array([1, 2, 3, -1], mask=[0, 0, 0, 1])) Chris@87: assert_equal(test[1], [0, 3, 5, 2]) Chris@87: assert_equal(test[2], [0, 0, 3, 1, 3, 2]) Chris@87: # Chris@87: data.fill_value = 3 Chris@87: data = masked_array(data=[1, 1, 1, 2, 2, 3], Chris@87: mask=[0, 0, 1, 0, 1, 0], fill_value=3) Chris@87: test = unique(data, return_index=True, return_inverse=True) Chris@87: assert_equal(test[0], masked_array([1, 2, 3, -1], mask=[0, 0, 0, 1])) Chris@87: assert_equal(test[1], [0, 3, 5, 2]) Chris@87: assert_equal(test[2], [0, 0, 3, 1, 3, 2]) Chris@87: Chris@87: def test_unique_allmasked(self): Chris@87: # Test all masked Chris@87: data = masked_array([1, 1, 1], mask=True) Chris@87: test = unique(data, return_index=True, return_inverse=True) Chris@87: assert_equal(test[0], masked_array([1, ], mask=[True])) Chris@87: assert_equal(test[1], [0]) Chris@87: assert_equal(test[2], [0, 0, 0]) Chris@87: # Chris@87: # Test masked Chris@87: data = masked Chris@87: test = unique(data, return_index=True, return_inverse=True) Chris@87: assert_equal(test[0], masked_array(masked)) Chris@87: assert_equal(test[1], [0]) Chris@87: assert_equal(test[2], [0]) Chris@87: Chris@87: def test_ediff1d(self): Chris@87: # Tests mediff1d Chris@87: x = masked_array(np.arange(5), mask=[1, 0, 0, 0, 1]) Chris@87: control = array([1, 1, 1, 4], mask=[1, 0, 0, 1]) Chris@87: test = ediff1d(x) Chris@87: assert_equal(test, control) Chris@87: assert_equal(test.data, control.data) Chris@87: assert_equal(test.mask, control.mask) Chris@87: Chris@87: def test_ediff1d_tobegin(self): Chris@87: # Test ediff1d w/ to_begin Chris@87: x = masked_array(np.arange(5), mask=[1, 0, 0, 0, 1]) Chris@87: test = ediff1d(x, to_begin=masked) Chris@87: control = array([0, 1, 1, 1, 4], mask=[1, 1, 0, 0, 1]) Chris@87: assert_equal(test, control) Chris@87: assert_equal(test.data, control.data) Chris@87: assert_equal(test.mask, control.mask) Chris@87: # Chris@87: test = ediff1d(x, to_begin=[1, 2, 3]) Chris@87: control = array([1, 2, 3, 1, 1, 1, 4], mask=[0, 0, 0, 1, 0, 0, 1]) Chris@87: assert_equal(test, control) Chris@87: assert_equal(test.data, control.data) Chris@87: assert_equal(test.mask, control.mask) Chris@87: Chris@87: def test_ediff1d_toend(self): Chris@87: # Test ediff1d w/ to_end Chris@87: x = masked_array(np.arange(5), mask=[1, 0, 0, 0, 1]) Chris@87: test = ediff1d(x, to_end=masked) Chris@87: control = array([1, 1, 1, 4, 0], mask=[1, 0, 0, 1, 1]) Chris@87: assert_equal(test, control) Chris@87: assert_equal(test.data, control.data) Chris@87: assert_equal(test.mask, control.mask) Chris@87: # Chris@87: test = ediff1d(x, to_end=[1, 2, 3]) Chris@87: control = array([1, 1, 1, 4, 1, 2, 3], mask=[1, 0, 0, 1, 0, 0, 0]) Chris@87: assert_equal(test, control) Chris@87: assert_equal(test.data, control.data) Chris@87: assert_equal(test.mask, control.mask) Chris@87: Chris@87: def test_ediff1d_tobegin_toend(self): Chris@87: # Test ediff1d w/ to_begin and to_end Chris@87: x = masked_array(np.arange(5), mask=[1, 0, 0, 0, 1]) Chris@87: test = ediff1d(x, to_end=masked, to_begin=masked) Chris@87: control = array([0, 1, 1, 1, 4, 0], mask=[1, 1, 0, 0, 1, 1]) Chris@87: assert_equal(test, control) Chris@87: assert_equal(test.data, control.data) Chris@87: assert_equal(test.mask, control.mask) Chris@87: # Chris@87: test = ediff1d(x, to_end=[1, 2, 3], to_begin=masked) Chris@87: control = array([0, 1, 1, 1, 4, 1, 2, 3], Chris@87: mask=[1, 1, 0, 0, 1, 0, 0, 0]) Chris@87: assert_equal(test, control) Chris@87: assert_equal(test.data, control.data) Chris@87: assert_equal(test.mask, control.mask) Chris@87: Chris@87: def test_ediff1d_ndarray(self): Chris@87: # Test ediff1d w/ a ndarray Chris@87: x = np.arange(5) Chris@87: test = ediff1d(x) Chris@87: control = array([1, 1, 1, 1], mask=[0, 0, 0, 0]) Chris@87: assert_equal(test, control) Chris@87: self.assertTrue(isinstance(test, MaskedArray)) Chris@87: assert_equal(test.data, control.data) Chris@87: assert_equal(test.mask, control.mask) Chris@87: # Chris@87: test = ediff1d(x, to_end=masked, to_begin=masked) Chris@87: control = array([0, 1, 1, 1, 1, 0], mask=[1, 0, 0, 0, 0, 1]) Chris@87: self.assertTrue(isinstance(test, MaskedArray)) Chris@87: assert_equal(test.data, control.data) Chris@87: assert_equal(test.mask, control.mask) Chris@87: Chris@87: def test_intersect1d(self): Chris@87: # Test intersect1d Chris@87: x = array([1, 3, 3, 3], mask=[0, 0, 0, 1]) Chris@87: y = array([3, 1, 1, 1], mask=[0, 0, 0, 1]) Chris@87: test = intersect1d(x, y) Chris@87: control = array([1, 3, -1], mask=[0, 0, 1]) Chris@87: assert_equal(test, control) Chris@87: Chris@87: def test_setxor1d(self): Chris@87: # Test setxor1d Chris@87: a = array([1, 2, 5, 7, -1], mask=[0, 0, 0, 0, 1]) Chris@87: b = array([1, 2, 3, 4, 5, -1], mask=[0, 0, 0, 0, 0, 1]) Chris@87: test = setxor1d(a, b) Chris@87: assert_equal(test, array([3, 4, 7])) Chris@87: # Chris@87: a = array([1, 2, 5, 7, -1], mask=[0, 0, 0, 0, 1]) Chris@87: b = [1, 2, 3, 4, 5] Chris@87: test = setxor1d(a, b) Chris@87: assert_equal(test, array([3, 4, 7, -1], mask=[0, 0, 0, 1])) Chris@87: # Chris@87: a = array([1, 2, 3]) Chris@87: b = array([6, 5, 4]) Chris@87: test = setxor1d(a, b) Chris@87: assert_(isinstance(test, MaskedArray)) Chris@87: assert_equal(test, [1, 2, 3, 4, 5, 6]) Chris@87: # Chris@87: a = array([1, 8, 2, 3], mask=[0, 1, 0, 0]) Chris@87: b = array([6, 5, 4, 8], mask=[0, 0, 0, 1]) Chris@87: test = setxor1d(a, b) Chris@87: assert_(isinstance(test, MaskedArray)) Chris@87: assert_equal(test, [1, 2, 3, 4, 5, 6]) Chris@87: # Chris@87: assert_array_equal([], setxor1d([], [])) Chris@87: Chris@87: def test_in1d(self): Chris@87: # Test in1d Chris@87: a = array([1, 2, 5, 7, -1], mask=[0, 0, 0, 0, 1]) Chris@87: b = array([1, 2, 3, 4, 5, -1], mask=[0, 0, 0, 0, 0, 1]) Chris@87: test = in1d(a, b) Chris@87: assert_equal(test, [True, True, True, False, True]) Chris@87: # Chris@87: a = array([5, 5, 2, 1, -1], mask=[0, 0, 0, 0, 1]) Chris@87: b = array([1, 5, -1], mask=[0, 0, 1]) Chris@87: test = in1d(a, b) Chris@87: assert_equal(test, [True, True, False, True, True]) Chris@87: # Chris@87: assert_array_equal([], in1d([], [])) Chris@87: Chris@87: def test_in1d_invert(self): Chris@87: # Test in1d's invert parameter Chris@87: a = array([1, 2, 5, 7, -1], mask=[0, 0, 0, 0, 1]) Chris@87: b = array([1, 2, 3, 4, 5, -1], mask=[0, 0, 0, 0, 0, 1]) Chris@87: assert_equal(np.invert(in1d(a, b)), in1d(a, b, invert=True)) Chris@87: Chris@87: a = array([5, 5, 2, 1, -1], mask=[0, 0, 0, 0, 1]) Chris@87: b = array([1, 5, -1], mask=[0, 0, 1]) Chris@87: assert_equal(np.invert(in1d(a, b)), in1d(a, b, invert=True)) Chris@87: Chris@87: assert_array_equal([], in1d([], [], invert=True)) Chris@87: Chris@87: def test_union1d(self): Chris@87: # Test union1d Chris@87: a = array([1, 2, 5, 7, 5, -1], mask=[0, 0, 0, 0, 0, 1]) Chris@87: b = array([1, 2, 3, 4, 5, -1], mask=[0, 0, 0, 0, 0, 1]) Chris@87: test = union1d(a, b) Chris@87: control = array([1, 2, 3, 4, 5, 7, -1], mask=[0, 0, 0, 0, 0, 0, 1]) Chris@87: assert_equal(test, control) Chris@87: # Chris@87: assert_array_equal([], union1d([], [])) Chris@87: Chris@87: def test_setdiff1d(self): Chris@87: # Test setdiff1d Chris@87: a = array([6, 5, 4, 7, 7, 1, 2, 1], mask=[0, 0, 0, 0, 0, 0, 0, 1]) Chris@87: b = array([2, 4, 3, 3, 2, 1, 5]) Chris@87: test = setdiff1d(a, b) Chris@87: assert_equal(test, array([6, 7, -1], mask=[0, 0, 1])) Chris@87: # Chris@87: a = arange(10) Chris@87: b = arange(8) Chris@87: assert_equal(setdiff1d(a, b), array([8, 9])) Chris@87: Chris@87: def test_setdiff1d_char_array(self): Chris@87: # Test setdiff1d_charray Chris@87: a = np.array(['a', 'b', 'c']) Chris@87: b = np.array(['a', 'b', 's']) Chris@87: assert_array_equal(setdiff1d(a, b), np.array(['c'])) Chris@87: Chris@87: Chris@87: class TestShapeBase(TestCase): Chris@87: # Chris@87: def test_atleast2d(self): Chris@87: # Test atleast_2d Chris@87: a = masked_array([0, 1, 2], mask=[0, 1, 0]) Chris@87: b = atleast_2d(a) Chris@87: assert_equal(b.shape, (1, 3)) Chris@87: assert_equal(b.mask.shape, b.data.shape) Chris@87: assert_equal(a.shape, (3,)) Chris@87: assert_equal(a.mask.shape, a.data.shape) Chris@87: Chris@87: Chris@87: ############################################################################### Chris@87: #------------------------------------------------------------------------------ Chris@87: if __name__ == "__main__": Chris@87: run_module_suite()