diff options
author | gfyoung <gfyoung17@gmail.com> | 2016-12-17 04:06:16 -0500 |
---|---|---|
committer | gfyoung <gfyoung17@gmail.com> | 2017-03-25 19:12:56 -0700 |
commit | d268966200adbe0e2afdd40aa92bb83c5a931e30 (patch) | |
tree | 17d52548fa9a0110406c44321ed2b84a3411bcb6 /numpy/lib | |
parent | b7dc5193f7136f5b09f45c3fb10118556e961e74 (diff) | |
download | numpy-d268966200adbe0e2afdd40aa92bb83c5a931e30.tar.gz |
API: Return scalars for scalar inputs to np.real/imag
Closes gh-8386.
Diffstat (limited to 'numpy/lib')
-rw-r--r-- | numpy/lib/tests/test_type_check.py | 40 | ||||
-rw-r--r-- | numpy/lib/type_check.py | 38 |
2 files changed, 66 insertions, 12 deletions
diff --git a/numpy/lib/tests/test_type_check.py b/numpy/lib/tests/test_type_check.py index 473b558be..383ffa55c 100644 --- a/numpy/lib/tests/test_type_check.py +++ b/numpy/lib/tests/test_type_check.py @@ -98,10 +98,30 @@ class TestReal(TestCase): y = np.random.rand(10,) assert_array_equal(y, np.real(y)) + y = np.array(1) + out = np.real(y) + assert_array_equal(y, out) + assert_(isinstance(out, np.ndarray)) + + y = 1 + out = np.real(y) + assert_equal(y, out) + assert_(not isinstance(out, np.ndarray)) + def test_cmplx(self): y = np.random.rand(10,)+1j*np.random.rand(10,) assert_array_equal(y.real, np.real(y)) + y = np.array(1 + 1j) + out = np.real(y) + assert_array_equal(y.real, out) + assert_(isinstance(out, np.ndarray)) + + y = 1 + 1j + out = np.real(y) + assert_equal(1.0, out) + assert_(not isinstance(out, np.ndarray)) + class TestImag(TestCase): @@ -109,10 +129,30 @@ class TestImag(TestCase): y = np.random.rand(10,) assert_array_equal(0, np.imag(y)) + y = np.array(1) + out = np.imag(y) + assert_array_equal(0, out) + assert_(isinstance(out, np.ndarray)) + + y = 1 + out = np.imag(y) + assert_equal(0, out) + assert_(not isinstance(out, np.ndarray)) + def test_cmplx(self): y = np.random.rand(10,)+1j*np.random.rand(10,) assert_array_equal(y.imag, np.imag(y)) + y = np.array(1 + 1j) + out = np.imag(y) + assert_array_equal(y.imag, out) + assert_(isinstance(out, np.ndarray)) + + y = 1 + 1j + out = np.imag(y) + assert_equal(1.0, out) + assert_(not isinstance(out, np.ndarray)) + class TestIscomplex(TestCase): diff --git a/numpy/lib/type_check.py b/numpy/lib/type_check.py index a59fe3cc4..5202cebde 100644 --- a/numpy/lib/type_check.py +++ b/numpy/lib/type_check.py @@ -9,8 +9,7 @@ __all__ = ['iscomplexobj', 'isrealobj', 'imag', 'iscomplex', 'common_type'] import numpy.core.numeric as _nx -from numpy.core.numeric import asarray, asanyarray, array, isnan, \ - obj2sctype, zeros +from numpy.core.numeric import asarray, asanyarray, array, isnan, zeros from .ufunclike import isneginf, isposinf _typecodes_by_elsize = 'GDFgdfQqLlIiHhBb?' @@ -104,9 +103,10 @@ def asfarray(a, dtype=_nx.float_): dtype = _nx.float_ return asarray(a, dtype=dtype) + def real(val): """ - Return the real part of the elements of the array. + Return the real part of the complex argument. Parameters ---------- @@ -115,9 +115,10 @@ def real(val): Returns ------- - out : ndarray - Output array. If `val` is real, the type of `val` is used for the - output. If `val` has complex elements, the returned type is float. + out : ndarray or scalar + The real component of the complex argument. If `val` is real, the type + of `val` is used for the output. If `val` has complex elements, the + returned type is float. See Also -------- @@ -134,13 +135,19 @@ def real(val): >>> a.real = np.array([9, 8, 7]) >>> a array([ 9.+2.j, 8.+4.j, 7.+6.j]) + >>> np.real(1 + 1j) + 1.0 """ - return asanyarray(val).real + try: + return val.real + except AttributeError: + return asanyarray(val).real + def imag(val): """ - Return the imaginary part of the elements of the array. + Return the imaginary part of the complex argument. Parameters ---------- @@ -149,9 +156,10 @@ def imag(val): Returns ------- - out : ndarray - Output array. If `val` is real, the type of `val` is used for the - output. If `val` has complex elements, the returned type is float. + out : ndarray or scalar + The imaginary component of the complex argument. If `val` is real, + the type of `val` is used for the output. If `val` has complex + elements, the returned type is float. See Also -------- @@ -165,9 +173,15 @@ def imag(val): >>> a.imag = np.array([8, 10, 12]) >>> a array([ 1. +8.j, 3.+10.j, 5.+12.j]) + >>> np.imag(1 + 1j) + 1.0 """ - return asanyarray(val).imag + try: + return val.imag + except AttributeError: + return asanyarray(val).imag + def iscomplex(x): """ |