test_ufunclike.py 3.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100
  1. import numpy as np
  2. from numpy import fix, isposinf, isneginf
  3. from numpy.testing import (
  4. assert_, assert_equal, assert_array_equal, assert_raises
  5. )
  6. class TestUfunclike:
  7. def test_isposinf(self):
  8. a = np.array([np.inf, -np.inf, np.nan, 0.0, 3.0, -3.0])
  9. out = np.zeros(a.shape, bool)
  10. tgt = np.array([True, False, False, False, False, False])
  11. res = isposinf(a)
  12. assert_equal(res, tgt)
  13. res = isposinf(a, out)
  14. assert_equal(res, tgt)
  15. assert_equal(out, tgt)
  16. a = a.astype(np.complex128)
  17. with assert_raises(TypeError):
  18. isposinf(a)
  19. def test_isneginf(self):
  20. a = np.array([np.inf, -np.inf, np.nan, 0.0, 3.0, -3.0])
  21. out = np.zeros(a.shape, bool)
  22. tgt = np.array([False, True, False, False, False, False])
  23. res = isneginf(a)
  24. assert_equal(res, tgt)
  25. res = isneginf(a, out)
  26. assert_equal(res, tgt)
  27. assert_equal(out, tgt)
  28. a = a.astype(np.complex128)
  29. with assert_raises(TypeError):
  30. isneginf(a)
  31. def test_fix(self):
  32. a = np.array([[1.0, 1.1, 1.5, 1.8], [-1.0, -1.1, -1.5, -1.8]])
  33. out = np.zeros(a.shape, float)
  34. tgt = np.array([[1., 1., 1., 1.], [-1., -1., -1., -1.]])
  35. res = fix(a)
  36. assert_equal(res, tgt)
  37. res = fix(a, out)
  38. assert_equal(res, tgt)
  39. assert_equal(out, tgt)
  40. assert_equal(fix(3.14), 3)
  41. def test_fix_with_subclass(self):
  42. class MyArray(np.ndarray):
  43. def __new__(cls, data, metadata=None):
  44. res = np.array(data, copy=True).view(cls)
  45. res.metadata = metadata
  46. return res
  47. def __array_wrap__(self, obj, context=None, return_scalar=False):
  48. if not isinstance(obj, MyArray):
  49. obj = obj.view(MyArray)
  50. if obj.metadata is None:
  51. obj.metadata = self.metadata
  52. return obj
  53. def __array_finalize__(self, obj):
  54. self.metadata = getattr(obj, 'metadata', None)
  55. return self
  56. a = np.array([1.1, -1.1])
  57. m = MyArray(a, metadata='foo')
  58. f = fix(m)
  59. assert_array_equal(f, np.array([1, -1]))
  60. assert_(isinstance(f, MyArray))
  61. assert_equal(f.metadata, 'foo')
  62. # check 0d arrays don't decay to scalars
  63. m0d = m[0,...]
  64. m0d.metadata = 'bar'
  65. f0d = fix(m0d)
  66. assert_(isinstance(f0d, MyArray))
  67. assert_equal(f0d.metadata, 'bar')
  68. def test_scalar(self):
  69. x = np.inf
  70. actual = np.isposinf(x)
  71. expected = np.True_
  72. assert_equal(actual, expected)
  73. assert_equal(type(actual), type(expected))
  74. x = -3.4
  75. actual = np.fix(x)
  76. expected = np.float64(-3.0)
  77. assert_equal(actual, expected)
  78. assert_equal(type(actual), type(expected))
  79. out = np.array(0.0)
  80. actual = np.fix(x, out=out)
  81. assert_(actual is out)