Add implementation of dpnp.real_if_close (#2002)

* Add implementation of dpnp.real_if_close

* Updated third party tests

* Added CFD tests

* Add more tests

* State default value in the description

* Add negative test for 'tol' keyword

* Add proper link description per review comment

Author: antonwolfy

dpnp/dpnp_iface_mathematical.py

@@ -114,6 +114,7 @@
     "prod",
     "proj",
     "real",
+    "real_if_close",
     "remainder",
     "rint",
     "round",
@@ -505,6 +506,10 @@ def _process_ediff1d_args(arg, arg_name, ary_dtype, ary_sycl_queue, usm_type):
 :obj:`dpnp.arctan2` : Element-wise arc tangent of `x1/x2` choosing the quadrant correctly.
 :obj:`dpnp.arctan` : Trigonometric inverse tangent, element-wise.
 :obj:`dpnp.absolute` : Calculate the absolute value element-wise.
+:obj:`dpnp.real` : Return the real part of the complex argument.
+:obj:`dpnp.imag` : Return the imaginary part of the complex argument.
+:obj:`dpnp.real_if_close` : Return the real part of the input is complex
+                            with all imaginary parts close to zero.
 
 Examples
 --------
@@ -2201,6 +2206,9 @@ def gradient(f, *varargs, axis=None, edge_order=1):
 See Also
 --------
 :obj:`dpnp.real` : Return the real part of the complex argument.
+:obj:`dpnp.angle` : Return the angle of the complex argument.
+:obj:`dpnp.real_if_close` : Return the real part of the input is complex
+                            with all imaginary parts close to zero.
 :obj:`dpnp.conj` : Return the complex conjugate, element-wise.
 :obj:`dpnp.conjugate` : Return the complex conjugate, element-wise.
 
@@ -3054,6 +3062,28 @@ def prod(
     the same data type. If the input is a complex floating-point
     data type, the returned array has a floating-point data type
     with the same floating-point precision as complex input.
+
+See Also
+--------
+:obj:`dpnp.real_if_close` : Return the real part of the input is complex
+                            with all imaginary parts close to zero.
+:obj:`dpnp.imag` : Return the imaginary part of the complex argument.
+:obj:`dpnp.angle` : Return the angle of the complex argument.
+
+Examples
+--------
+>>> import dpnp as np
+>>> a = np.array([1+2j, 3+4j, 5+6j])
+>>> a.real
+array([1., 3., 5.])
+>>> a.real = 9
+>>> a
+array([9.+2.j, 9.+4.j, 9.+6.j])
+>>> a.real = np.array([9, 8, 7])
+>>> a
+array([9.+2.j, 8.+4.j, 7.+6.j])
+>>> np.real(np.array(1 + 1j))
+array(1.)
 """
 
 real = DPNPReal(
@@ -3064,6 +3094,69 @@ def prod(
 )
 
 
+def real_if_close(a, tol=100):
+    """
+    If input is complex with all imaginary parts close to zero, return real
+    parts.
+
+    "Close to zero" is defined as `tol` * (machine epsilon of the type for `a`).
+
+    For full documentation refer to :obj:`numpy.real_if_close`.
+
+    Parameters
+    ----------
+    a : {dpnp.ndarray, usm_ndarray}
+        Input array.
+    tol : scalar, optional
+        Tolerance in machine epsilons for the complex part of the elements in
+        the array. If the tolerance is <=1, then the absolute tolerance is used.
+        Default: ``100``.
+
+    Returns
+    -------
+    out : dpnp.ndarray
+        If `a` is real, the type of `a` is used for the output. If `a` has
+        complex elements, the returned type is float.
+
+    See Also
+    --------
+    :obj:`dpnp.real` : Return the real part of the complex argument.
+    :obj:`dpnp.imag` : Return the imaginary part of the complex argument.
+    :obj:`dpnp.angle` : Return the angle of the complex argument.
+
+    Examples
+    --------
+    >>> import dpnp as np
+    >>> np.finfo(np.float64).eps
+    2.220446049250313e-16 # may vary
+
+    >>> a = np.array([2.1 + 4e-14j, 5.2 + 3e-15j])
+    >>> np.real_if_close(a, tol=1000)
+    array([2.1, 5.2])
+
+    >>> a = np.array([2.1 + 4e-13j, 5.2 + 3e-15j])
+    >>> np.real_if_close(a, tol=1000)
+    array([2.1+4.e-13j, 5.2+3.e-15j])
+
+    """
+
+    dpnp.check_supported_arrays_type(a)
+
+    if not dpnp.issubdtype(a.dtype, dpnp.complexfloating):
+        return a
+
+    if not dpnp.isscalar(tol):
+        raise TypeError(f"Tolerance must be a scalar, but got {type(tol)}")
+
+    if tol > 1:
+        f = dpnp.finfo(a.dtype.type)
+        tol = f.eps * tol
+
+    if dpnp.all(dpnp.abs(a.imag) < tol):
+        return a.real
+    return a
+
+
 _REMAINDER_DOCSTRING = """
 Calculates the remainder of division for each element `x1_i` of the input array
 `x1` with the respective element `x2_i` of the input array `x2`.

tests/skipped_tests.tbl

@@ -182,14 +182,6 @@ tests/third_party/cupy/manipulation_tests/test_dims.py::TestInvalidBroadcast_par
 tests/third_party/cupy/manipulation_tests/test_dims.py::TestInvalidBroadcast_param_2_{shapes=[(3, 2), (3, 4)]}::test_invalid_broadcast
 tests/third_party/cupy/manipulation_tests/test_dims.py::TestInvalidBroadcast_param_3_{shapes=[(0,), (2,)]}::test_invalid_broadcast
 
-tests/third_party/cupy/math_tests/test_misc.py::TestMisc::test_real_if_close_real_dtypes
-tests/third_party/cupy/math_tests/test_misc.py::TestMisc::test_real_if_close_with_tol_real_dtypes
-tests/third_party/cupy/math_tests/test_misc.py::TestMisc::test_real_if_close_true
-tests/third_party/cupy/math_tests/test_misc.py::TestMisc::test_real_if_close_false
-tests/third_party/cupy/math_tests/test_misc.py::TestMisc::test_real_if_close_with_integer_tol_true
-tests/third_party/cupy/math_tests/test_misc.py::TestMisc::test_real_if_close_with_integer_tol_false
-tests/third_party/cupy/math_tests/test_misc.py::TestMisc::test_real_if_close_with_float_tol_true
-tests/third_party/cupy/math_tests/test_misc.py::TestMisc::test_real_if_close_with_float_tol_false
 tests/third_party/cupy/math_tests/test_misc.py::TestMisc::test_interp
 tests/third_party/cupy/math_tests/test_misc.py::TestMisc::test_interp_period
 tests/third_party/cupy/math_tests/test_misc.py::TestMisc::test_interp_left_right

tests/skipped_tests_gpu.tbl

@@ -236,14 +236,6 @@ tests/third_party/cupy/manipulation_tests/test_dims.py::TestInvalidBroadcast_par
 tests/third_party/cupy/manipulation_tests/test_dims.py::TestInvalidBroadcast_param_2_{shapes=[(3, 2), (3, 4)]}::test_invalid_broadcast
 tests/third_party/cupy/manipulation_tests/test_dims.py::TestInvalidBroadcast_param_3_{shapes=[(0,), (2,)]}::test_invalid_broadcast
 
-tests/third_party/cupy/math_tests/test_misc.py::TestMisc::test_real_if_close_real_dtypes
-tests/third_party/cupy/math_tests/test_misc.py::TestMisc::test_real_if_close_with_tol_real_dtypes
-tests/third_party/cupy/math_tests/test_misc.py::TestMisc::test_real_if_close_true
-tests/third_party/cupy/math_tests/test_misc.py::TestMisc::test_real_if_close_false
-tests/third_party/cupy/math_tests/test_misc.py::TestMisc::test_real_if_close_with_integer_tol_true
-tests/third_party/cupy/math_tests/test_misc.py::TestMisc::test_real_if_close_with_integer_tol_false
-tests/third_party/cupy/math_tests/test_misc.py::TestMisc::test_real_if_close_with_float_tol_true
-tests/third_party/cupy/math_tests/test_misc.py::TestMisc::test_real_if_close_with_float_tol_false
 tests/third_party/cupy/math_tests/test_misc.py::TestMisc::test_interp
 tests/third_party/cupy/math_tests/test_misc.py::TestMisc::test_interp_period
 tests/third_party/cupy/math_tests/test_misc.py::TestMisc::test_interp_left_right

tests/test_mathematical.py

@@ -1377,6 +1377,41 @@ def test_f16_corner_values_with_scalar(self, val, scalar):
         assert_equal(result, expected)
 
 
+class TestRealIfClose:
+    @pytest.mark.parametrize("dt", get_all_dtypes(no_none=True))
+    def test_basic(self, dt):
+        a = numpy.random.rand(10).astype(dt)
+        ia = dpnp.array(a)
+
+        result = dpnp.real_if_close(ia + 1e-15j)
+        expected = numpy.real_if_close(a + 1e-15j)
+        assert_equal(result, expected)
+
+    @pytest.mark.parametrize("dt", get_float_dtypes())
+    def test_singlecomplex(self, dt):
+        a = numpy.random.rand(10).astype(dt)
+        ia = dpnp.array(a)
+
+        result = dpnp.real_if_close(ia + 1e-7j)
+        expected = numpy.real_if_close(a + 1e-7j)
+        assert_equal(result, expected)
+
+    @pytest.mark.parametrize("dt", get_float_dtypes())
+    def test_tol(self, dt):
+        a = numpy.random.rand(10).astype(dt)
+        ia = dpnp.array(a)
+
+        result = dpnp.real_if_close(ia + 1e-7j, tol=1e-6)
+        expected = numpy.real_if_close(a + 1e-7j, tol=1e-6)
+        assert_equal(result, expected)
+
+    @pytest.mark.parametrize("xp", [numpy, dpnp])
+    @pytest.mark.parametrize("tol_val", [[10], (1, 2)], ids=["list", "tuple"])
+    def test_wrong_tol_type(self, xp, tol_val):
+        a = xp.array([2.1 + 4e-14j, 5.2 + 3e-15j])
+        assert_raises(TypeError, xp.real_if_close, a, tol=tol_val)
+
+
 class TestUnwrap:
     @pytest.mark.parametrize("dt", get_float_dtypes())
     def test_basic(self, dt):

tests/test_sycl_queue.py

@@ -493,6 +493,7 @@ def test_meshgrid(device):
         pytest.param(
             "real", [complex(1.0, 2.0), complex(3.0, 4.0), complex(5.0, 6.0)]
         ),
+        pytest.param("real_if_close", [2.1 + 4e-15j, 5.2 + 3e-16j]),
         pytest.param("reciprocal", [1.0, 2.0, 4.0, 7.0]),
         pytest.param("sign", [-5.0, 0.0, 4.5]),
         pytest.param("signbit", [-5.0, 0.0, 4.5]),

tests/test_usm_type.py

@@ -618,6 +618,7 @@ def test_norm(usm_type, ord, axis):
         pytest.param(
             "real", [complex(1.0, 2.0), complex(3.0, 4.0), complex(5.0, 6.0)]
         ),
+        pytest.param("real_if_close", [2.1 + 4e-15j, 5.2 + 3e-16j]),
         pytest.param("reciprocal", [1.0, 2.0, 4.0, 7.0]),
         pytest.param("reduce_hypot", [1.0, 2.0, 4.0, 7.0]),
         pytest.param("rsqrt", [1, 8, 27]),

tests/third_party/cupy/math_tests/test_misc.py

@@ -245,7 +245,7 @@ def test_nan_to_num_inf(self):
     def test_nan_to_num_nan(self):
         self.check_unary_nan("nan_to_num")
 
-    @pytest.mark.skip(reason="Scalar input is not supported")
+    @pytest.mark.skip(reason="scalar input is not supported")
     @testing.numpy_cupy_allclose(atol=1e-5)
     def test_nan_to_num_scalar_nan(self, xp):
         return xp.nan_to_num(xp.nan)
@@ -301,7 +301,8 @@ def test_real_if_close_with_tol_real_dtypes(self, xp, dtype):
     def test_real_if_close_true(self, xp, dtype):
         dtype = numpy.dtype(dtype).char.lower()
         tol = numpy.finfo(dtype).eps * 90
-        x = testing.shaped_random((10,), xp, dtype) + tol * 1j
+        x = testing.shaped_random((10,), xp, dtype)
+        x = xp.add(x, tol * 1j, dtype=xp.result_type(x, 1j))
         out = xp.real_if_close(x)
         assert x.dtype != out.dtype
         return out
@@ -311,7 +312,8 @@ def test_real_if_close_true(self, xp, dtype):
     def test_real_if_close_false(self, xp, dtype):
         dtype = numpy.dtype(dtype).char.lower()
         tol = numpy.finfo(dtype).eps * 110
-        x = testing.shaped_random((10,), xp, dtype) + tol * 1j
+        x = testing.shaped_random((10,), xp, dtype)
+        x = xp.add(x, tol * 1j, dtype=xp.result_type(x, 1j))
         out = xp.real_if_close(x)
         assert x.dtype == out.dtype
         return out
@@ -321,7 +323,8 @@ def test_real_if_close_false(self, xp, dtype):
     def test_real_if_close_with_integer_tol_true(self, xp, dtype):
         dtype = numpy.dtype(dtype).char.lower()
         tol = numpy.finfo(dtype).eps * 140
-        x = testing.shaped_random((10,), xp, dtype) + tol * 1j
+        x = testing.shaped_random((10,), xp, dtype)
+        x = xp.add(x, tol * 1j, dtype=xp.result_type(x, 1j))
         out = xp.real_if_close(x, tol=150)
         assert x.dtype != out.dtype
         return out
@@ -331,7 +334,8 @@ def test_real_if_close_with_integer_tol_true(self, xp, dtype):
     def test_real_if_close_with_integer_tol_false(self, xp, dtype):
         dtype = numpy.dtype(dtype).char.lower()
         tol = numpy.finfo(dtype).eps * 50
-        x = testing.shaped_random((10,), xp, dtype) + tol * 1j
+        x = testing.shaped_random((10,), xp, dtype)
+        x = xp.add(x, tol * 1j, dtype=xp.result_type(x, 1j))
         out = xp.real_if_close(x, tol=30)
         assert x.dtype == out.dtype
         return out