Chris@87: from __future__ import division, absolute_import, print_function
Chris@87: 
Chris@87: import numpy.core as nx
Chris@87: import numpy.lib.ufunclike as ufl
Chris@87: from numpy.testing import (
Chris@87:     run_module_suite, TestCase, assert_, assert_equal, assert_array_equal
Chris@87:     )
Chris@87: 
Chris@87: 
Chris@87: class TestUfunclike(TestCase):
Chris@87: 
Chris@87:     def test_isposinf(self):
Chris@87:         a = nx.array([nx.inf, -nx.inf, nx.nan, 0.0, 3.0, -3.0])
Chris@87:         out = nx.zeros(a.shape, bool)
Chris@87:         tgt = nx.array([True, False, False, False, False, False])
Chris@87: 
Chris@87:         res = ufl.isposinf(a)
Chris@87:         assert_equal(res, tgt)
Chris@87:         res = ufl.isposinf(a, out)
Chris@87:         assert_equal(res, tgt)
Chris@87:         assert_equal(out, tgt)
Chris@87: 
Chris@87:     def test_isneginf(self):
Chris@87:         a = nx.array([nx.inf, -nx.inf, nx.nan, 0.0, 3.0, -3.0])
Chris@87:         out = nx.zeros(a.shape, bool)
Chris@87:         tgt = nx.array([False, True, False, False, False, False])
Chris@87: 
Chris@87:         res = ufl.isneginf(a)
Chris@87:         assert_equal(res, tgt)
Chris@87:         res = ufl.isneginf(a, out)
Chris@87:         assert_equal(res, tgt)
Chris@87:         assert_equal(out, tgt)
Chris@87: 
Chris@87:     def test_fix(self):
Chris@87:         a = nx.array([[1.0, 1.1, 1.5, 1.8], [-1.0, -1.1, -1.5, -1.8]])
Chris@87:         out = nx.zeros(a.shape, float)
Chris@87:         tgt = nx.array([[1., 1., 1., 1.], [-1., -1., -1., -1.]])
Chris@87: 
Chris@87:         res = ufl.fix(a)
Chris@87:         assert_equal(res, tgt)
Chris@87:         res = ufl.fix(a, out)
Chris@87:         assert_equal(res, tgt)
Chris@87:         assert_equal(out, tgt)
Chris@87:         assert_equal(ufl.fix(3.14), 3)
Chris@87: 
Chris@87:     def test_fix_with_subclass(self):
Chris@87:         class MyArray(nx.ndarray):
Chris@87:             def __new__(cls, data, metadata=None):
Chris@87:                 res = nx.array(data, copy=True).view(cls)
Chris@87:                 res.metadata = metadata
Chris@87:                 return res
Chris@87: 
Chris@87:             def __array_wrap__(self, obj, context=None):
Chris@87:                 obj.metadata = self.metadata
Chris@87:                 return obj
Chris@87: 
Chris@87:         a = nx.array([1.1, -1.1])
Chris@87:         m = MyArray(a, metadata='foo')
Chris@87:         f = ufl.fix(m)
Chris@87:         assert_array_equal(f, nx.array([1, -1]))
Chris@87:         assert_(isinstance(f, MyArray))
Chris@87:         assert_equal(f.metadata, 'foo')
Chris@87: 
Chris@87: if __name__ == "__main__":
Chris@87:     run_module_suite()