Chris@87: # pylint: disable-msg=W0611, W0612, W0511,R0201
Chris@87: """Tests suite for MaskedArray & subclassing.
Chris@87: 
Chris@87: :author: Pierre Gerard-Marchant
Chris@87: :contact: pierregm_at_uga_dot_edu
Chris@87: :version: $Id: test_subclassing.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 *
Chris@87: from numpy.ma.testutils import *
Chris@87: from numpy.ma.core import *
Chris@87: 
Chris@87: 
Chris@87: class SubArray(np.ndarray):
Chris@87:     # Defines a generic np.ndarray subclass, that stores some metadata
Chris@87:     # in the  dictionary `info`.
Chris@87:     def __new__(cls,arr,info={}):
Chris@87:         x = np.asanyarray(arr).view(cls)
Chris@87:         x.info = info
Chris@87:         return x
Chris@87: 
Chris@87:     def __array_finalize__(self, obj):
Chris@87:         self.info = getattr(obj, 'info', {})
Chris@87:         return
Chris@87: 
Chris@87:     def __add__(self, other):
Chris@87:         result = np.ndarray.__add__(self, other)
Chris@87:         result.info.update({'added':result.info.pop('added', 0)+1})
Chris@87:         return result
Chris@87: 
Chris@87: subarray = SubArray
Chris@87: 
Chris@87: 
Chris@87: class MSubArray(SubArray, MaskedArray):
Chris@87: 
Chris@87:     def __new__(cls, data, info={}, mask=nomask):
Chris@87:         subarr = SubArray(data, info)
Chris@87:         _data = MaskedArray.__new__(cls, data=subarr, mask=mask)
Chris@87:         _data.info = subarr.info
Chris@87:         return _data
Chris@87: 
Chris@87:     def __array_finalize__(self, obj):
Chris@87:         MaskedArray.__array_finalize__(self, obj)
Chris@87:         SubArray.__array_finalize__(self, obj)
Chris@87:         return
Chris@87: 
Chris@87:     def _get_series(self):
Chris@87:         _view = self.view(MaskedArray)
Chris@87:         _view._sharedmask = False
Chris@87:         return _view
Chris@87:     _series = property(fget=_get_series)
Chris@87: 
Chris@87: msubarray = MSubArray
Chris@87: 
Chris@87: 
Chris@87: class MMatrix(MaskedArray, np.matrix,):
Chris@87: 
Chris@87:     def __new__(cls, data, mask=nomask):
Chris@87:         mat = np.matrix(data)
Chris@87:         _data = MaskedArray.__new__(cls, data=mat, mask=mask)
Chris@87:         return _data
Chris@87: 
Chris@87:     def __array_finalize__(self, obj):
Chris@87:         np.matrix.__array_finalize__(self, obj)
Chris@87:         MaskedArray.__array_finalize__(self, obj)
Chris@87:         return
Chris@87: 
Chris@87:     def _get_series(self):
Chris@87:         _view = self.view(MaskedArray)
Chris@87:         _view._sharedmask = False
Chris@87:         return _view
Chris@87:     _series = property(fget=_get_series)
Chris@87: 
Chris@87: mmatrix = MMatrix
Chris@87: 
Chris@87: 
Chris@87: # also a subclass that overrides __str__, __repr__ and __setitem__, disallowing
Chris@87: # setting to non-class values (and thus np.ma.core.masked_print_option)
Chris@87: class ComplicatedSubArray(SubArray):
Chris@87:     def __str__(self):
Chris@87:         return 'myprefix {0} mypostfix'.format(
Chris@87:             super(ComplicatedSubArray, self).__str__())
Chris@87: 
Chris@87:     def __repr__(self):
Chris@87:         # Return a repr that does not start with 'name('
Chris@87:         return '<{0} {1}>'.format(self.__class__.__name__, self)
Chris@87: 
Chris@87:     def __setitem__(self, item, value):
Chris@87:         # this ensures direct assignment to masked_print_option will fail
Chris@87:         if not isinstance(value, ComplicatedSubArray):
Chris@87:             raise ValueError("Can only set to MySubArray values")
Chris@87:         super(ComplicatedSubArray, self).__setitem__(item, value)
Chris@87: 
Chris@87: 
Chris@87: class TestSubclassing(TestCase):
Chris@87:     # Test suite for masked subclasses of ndarray.
Chris@87: 
Chris@87:     def setUp(self):
Chris@87:         x = np.arange(5)
Chris@87:         mx = mmatrix(x, mask=[0, 1, 0, 0, 0])
Chris@87:         self.data = (x, mx)
Chris@87: 
Chris@87:     def test_data_subclassing(self):
Chris@87:         # Tests whether the subclass is kept.
Chris@87:         x = np.arange(5)
Chris@87:         m = [0, 0, 1, 0, 0]
Chris@87:         xsub = SubArray(x)
Chris@87:         xmsub = masked_array(xsub, mask=m)
Chris@87:         self.assertTrue(isinstance(xmsub, MaskedArray))
Chris@87:         assert_equal(xmsub._data, xsub)
Chris@87:         self.assertTrue(isinstance(xmsub._data, SubArray))
Chris@87: 
Chris@87:     def test_maskedarray_subclassing(self):
Chris@87:         # Tests subclassing MaskedArray
Chris@87:         (x, mx) = self.data
Chris@87:         self.assertTrue(isinstance(mx._data, np.matrix))
Chris@87: 
Chris@87:     def test_masked_unary_operations(self):
Chris@87:         # Tests masked_unary_operation
Chris@87:         (x, mx) = self.data
Chris@87:         with np.errstate(divide='ignore'):
Chris@87:             self.assertTrue(isinstance(log(mx), mmatrix))
Chris@87:             assert_equal(log(x), np.log(x))
Chris@87: 
Chris@87:     def test_masked_binary_operations(self):
Chris@87:         # Tests masked_binary_operation
Chris@87:         (x, mx) = self.data
Chris@87:         # Result should be a mmatrix
Chris@87:         self.assertTrue(isinstance(add(mx, mx), mmatrix))
Chris@87:         self.assertTrue(isinstance(add(mx, x), mmatrix))
Chris@87:         # Result should work
Chris@87:         assert_equal(add(mx, x), mx+x)
Chris@87:         self.assertTrue(isinstance(add(mx, mx)._data, np.matrix))
Chris@87:         self.assertTrue(isinstance(add.outer(mx, mx), mmatrix))
Chris@87:         self.assertTrue(isinstance(hypot(mx, mx), mmatrix))
Chris@87:         self.assertTrue(isinstance(hypot(mx, x), mmatrix))
Chris@87: 
Chris@87:     def test_masked_binary_operations2(self):
Chris@87:         # Tests domained_masked_binary_operation
Chris@87:         (x, mx) = self.data
Chris@87:         xmx = masked_array(mx.data.__array__(), mask=mx.mask)
Chris@87:         self.assertTrue(isinstance(divide(mx, mx), mmatrix))
Chris@87:         self.assertTrue(isinstance(divide(mx, x), mmatrix))
Chris@87:         assert_equal(divide(mx, mx), divide(xmx, xmx))
Chris@87: 
Chris@87:     def test_attributepropagation(self):
Chris@87:         x = array(arange(5), mask=[0]+[1]*4)
Chris@87:         my = masked_array(subarray(x))
Chris@87:         ym = msubarray(x)
Chris@87:         #
Chris@87:         z = (my+1)
Chris@87:         self.assertTrue(isinstance(z, MaskedArray))
Chris@87:         self.assertTrue(not isinstance(z, MSubArray))
Chris@87:         self.assertTrue(isinstance(z._data, SubArray))
Chris@87:         assert_equal(z._data.info, {})
Chris@87:         #
Chris@87:         z = (ym+1)
Chris@87:         self.assertTrue(isinstance(z, MaskedArray))
Chris@87:         self.assertTrue(isinstance(z, MSubArray))
Chris@87:         self.assertTrue(isinstance(z._data, SubArray))
Chris@87:         self.assertTrue(z._data.info['added'] > 0)
Chris@87:         #
Chris@87:         ym._set_mask([1, 0, 0, 0, 1])
Chris@87:         assert_equal(ym._mask, [1, 0, 0, 0, 1])
Chris@87:         ym._series._set_mask([0, 0, 0, 0, 1])
Chris@87:         assert_equal(ym._mask, [0, 0, 0, 0, 1])
Chris@87:         #
Chris@87:         xsub = subarray(x, info={'name':'x'})
Chris@87:         mxsub = masked_array(xsub)
Chris@87:         self.assertTrue(hasattr(mxsub, 'info'))
Chris@87:         assert_equal(mxsub.info, xsub.info)
Chris@87: 
Chris@87:     def test_subclasspreservation(self):
Chris@87:         # Checks that masked_array(...,subok=True) preserves the class.
Chris@87:         x = np.arange(5)
Chris@87:         m = [0, 0, 1, 0, 0]
Chris@87:         xinfo = [(i, j) for (i, j) in zip(x, m)]
Chris@87:         xsub = MSubArray(x, mask=m, info={'xsub':xinfo})
Chris@87:         #
Chris@87:         mxsub = masked_array(xsub, subok=False)
Chris@87:         self.assertTrue(not isinstance(mxsub, MSubArray))
Chris@87:         self.assertTrue(isinstance(mxsub, MaskedArray))
Chris@87:         assert_equal(mxsub._mask, m)
Chris@87:         #
Chris@87:         mxsub = asarray(xsub)
Chris@87:         self.assertTrue(not isinstance(mxsub, MSubArray))
Chris@87:         self.assertTrue(isinstance(mxsub, MaskedArray))
Chris@87:         assert_equal(mxsub._mask, m)
Chris@87:         #
Chris@87:         mxsub = masked_array(xsub, subok=True)
Chris@87:         self.assertTrue(isinstance(mxsub, MSubArray))
Chris@87:         assert_equal(mxsub.info, xsub.info)
Chris@87:         assert_equal(mxsub._mask, xsub._mask)
Chris@87:         #
Chris@87:         mxsub = asanyarray(xsub)
Chris@87:         self.assertTrue(isinstance(mxsub, MSubArray))
Chris@87:         assert_equal(mxsub.info, xsub.info)
Chris@87:         assert_equal(mxsub._mask, m)
Chris@87: 
Chris@87:     def test_subclass_repr(self):
Chris@87:         """test that repr uses the name of the subclass
Chris@87:         and 'array' for np.ndarray"""
Chris@87:         x = np.arange(5)
Chris@87:         mx = masked_array(x, mask=[True, False, True, False, False])
Chris@87:         self.assertTrue(repr(mx).startswith('masked_array'))
Chris@87:         xsub = SubArray(x)
Chris@87:         mxsub = masked_array(xsub, mask=[True, False, True, False, False])
Chris@87:         self.assertTrue(repr(mxsub).startswith(
Chris@87:             'masked_{0}(data = [-- 1 -- 3 4]'.format(SubArray.__name__)))
Chris@87: 
Chris@87:     def test_subclass_str(self):
Chris@87:         """test str with subclass that has overridden str, setitem"""
Chris@87:         # first without override
Chris@87:         x = np.arange(5)
Chris@87:         xsub = SubArray(x)
Chris@87:         mxsub = masked_array(xsub, mask=[True, False, True, False, False])
Chris@87:         self.assertTrue(str(mxsub) == '[-- 1 -- 3 4]')
Chris@87: 
Chris@87:         xcsub = ComplicatedSubArray(x)
Chris@87:         assert_raises(ValueError, xcsub.__setitem__, 0,
Chris@87:                       np.ma.core.masked_print_option)
Chris@87:         mxcsub = masked_array(xcsub, mask=[True, False, True, False, False])
Chris@87:         self.assertTrue(str(mxcsub) == 'myprefix [-- 1 -- 3 4] mypostfix')
Chris@87: 
Chris@87: 
Chris@87: ###############################################################################
Chris@87: if __name__ == '__main__':
Chris@87:     run_module_suite()