Enable many complex number tests

doc/source/whatsnew/v2.2.0.rst

@@ -351,6 +351,7 @@ Styler
 
 Other
 ^^^^^
+- Add complex128 to the types of numerical data we test across the test suite (:issue:`54761`)
 - Bug in :func:`cut` incorrectly allowing cutting of timezone-aware datetimes with timezone-naive bins (:issue:`54964`)
 
 .. ***DO NOT USE THIS SECTION***

pandas/core/dtypes/astype.py

@@ -100,6 +100,11 @@ def _astype_nansafe(
     elif np.issubdtype(arr.dtype, np.floating) and dtype.kind in "iu":
         return _astype_float_to_int_nansafe(arr, dtype, copy)
 
+    elif np.issubdtype(arr.dtype, np.complexfloating) and is_object_dtype(dtype):
+        res = arr.astype(dtype, copy=copy)
+        res[np.isnan(arr)] = np.nan
+        return res
+
     elif arr.dtype == object:
         # if we have a datetime/timedelta array of objects
         # then coerce to datetime64[ns] and use DatetimeArray.astype

pandas/core/nanops.py

@@ -1003,11 +1003,25 @@ def nanvar(
     # cancellation errors and relatively accurate for small numbers of
     # observations.
     #
-    # See https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance
-    avg = _ensure_numeric(values.sum(axis=axis, dtype=np.float64)) / count
+    # See https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance...
+    if values.dtype.kind == "c":
+        avg = _ensure_numeric(values.sum(axis=axis, dtype=values.dtype)) / count
+    else:
+        avg = _ensure_numeric(values.sum(axis=axis, dtype=np.float64)) / count
     if axis is not None:
         avg = np.expand_dims(avg, axis)
-    sqr = _ensure_numeric((avg - values) ** 2)
+    # ...but also,
+    # see https://numpy.org/doc/stable/reference/generated/numpy.nanvar.html#numpy-nanvar
+    # which explains why computing the variance of complex numbers
+    # requires first normalizing the complex differences to magnitudes
+    if values.dtype.kind == "c":
+        deltas = _ensure_numeric(avg - values)
+        avg_re = np.real(deltas)
+        avg_im = np.imag(deltas)
+        sqr = avg_re**2 + avg_im**2
+    else:
+        sqr = _ensure_numeric((avg - values) ** 2)
+
     if mask is not None:
         np.putmask(sqr, mask, 0)
     result = sqr.sum(axis=axis, dtype=np.float64) / d

pandas/tests/arithmetic/test_numeric.py

@@ -977,7 +977,7 @@ def test_frame_operators_none_to_nan(self):
         df = pd.DataFrame({"a": ["a", None, "b"]})
         tm.assert_frame_equal(df + df, pd.DataFrame({"a": ["aa", np.nan, "bb"]}))
 
-    @pytest.mark.parametrize("dtype", ("float", "int64"))
+    @pytest.mark.parametrize("dtype", ("float", "int64", "complex128"))
     def test_frame_operators_empty_like(self, dtype):
         # Test for issue #10181
         frames = [
@@ -1059,7 +1059,7 @@ def test_series_divmod_zero(self):
 class TestUFuncCompat:
     # TODO: add more dtypes
     @pytest.mark.parametrize("holder", [Index, RangeIndex, Series])
-    @pytest.mark.parametrize("dtype", [np.int64, np.uint64, np.float64])
+    @pytest.mark.parametrize("dtype", [np.int64, np.uint64, np.float64, np.complex128])
     def test_ufunc_compat(self, holder, dtype):
         box = Series if holder is Series else Index
 
@@ -1075,45 +1075,75 @@ def test_ufunc_compat(self, holder, dtype):
 
     # TODO: add more dtypes
     @pytest.mark.parametrize("holder", [Index, Series])
-    @pytest.mark.parametrize("dtype", [np.int64, np.uint64, np.float64])
+    @pytest.mark.parametrize("dtype", [np.int64, np.uint64, np.float64, np.complex128])
     def test_ufunc_coercions(self, holder, dtype):
         idx = holder([1, 2, 3, 4, 5], dtype=dtype, name="x")
         box = Series if holder is Series else Index
 
         result = np.sqrt(idx)
-        assert result.dtype == "f8" and isinstance(result, box)
-        exp = Index(np.sqrt(np.array([1, 2, 3, 4, 5], dtype=np.float64)), name="x")
+        if result.dtype.kind == "c":
+            assert result.dtype == dtype and isinstance(result, box)
+            exp_dtype = dtype
+        else:
+            assert result.dtype == "f8" and isinstance(result, box)
+            exp_dtype = np.float64
+        exp = Index(np.sqrt(np.array([1, 2, 3, 4, 5], dtype=exp_dtype)), name="x")
         exp = tm.box_expected(exp, box)
         tm.assert_equal(result, exp)
 
         result = np.divide(idx, 2.0)
-        assert result.dtype == "f8" and isinstance(result, box)
-        exp = Index([0.5, 1.0, 1.5, 2.0, 2.5], dtype=np.float64, name="x")
+        if result.dtype.kind == "c":
+            assert result.dtype == dtype and isinstance(result, box)
+            exp_dtype = dtype
+        else:
+            assert result.dtype == "f8" and isinstance(result, box)
+            exp_dtype = np.float64
+        exp = Index([0.5, 1.0, 1.5, 2.0, 2.5], dtype=exp_dtype, name="x")
         exp = tm.box_expected(exp, box)
         tm.assert_equal(result, exp)
 
         # _evaluate_numeric_binop
         result = idx + 2.0
-        assert result.dtype == "f8" and isinstance(result, box)
-        exp = Index([3.0, 4.0, 5.0, 6.0, 7.0], dtype=np.float64, name="x")
+        if result.dtype.kind == "c":
+            assert result.dtype == dtype and isinstance(result, box)
+            exp_dtype = dtype
+        else:
+            assert result.dtype == "f8" and isinstance(result, box)
+            exp_dtype = np.float64
+        exp = Index([3.0, 4.0, 5.0, 6.0, 7.0], dtype=exp_dtype, name="x")
         exp = tm.box_expected(exp, box)
         tm.assert_equal(result, exp)
 
         result = idx - 2.0
-        assert result.dtype == "f8" and isinstance(result, box)
-        exp = Index([-1.0, 0.0, 1.0, 2.0, 3.0], dtype=np.float64, name="x")
+        if result.dtype.kind == "c":
+            assert result.dtype == dtype and isinstance(result, box)
+            exp_dtype = dtype
+        else:
+            assert result.dtype == "f8" and isinstance(result, box)
+            exp_dtype = np.float64
+        exp = Index([-1.0, 0.0, 1.0, 2.0, 3.0], dtype=exp_dtype, name="x")
         exp = tm.box_expected(exp, box)
         tm.assert_equal(result, exp)
 
         result = idx * 1.0
-        assert result.dtype == "f8" and isinstance(result, box)
-        exp = Index([1.0, 2.0, 3.0, 4.0, 5.0], dtype=np.float64, name="x")
+        if result.dtype.kind == "c":
+            assert result.dtype == dtype and isinstance(result, box)
+            exp_dtype = dtype
+        else:
+            assert result.dtype == "f8" and isinstance(result, box)
+            exp_dtype = np.float64
+        exp = Index([1.0, 2.0, 3.0, 4.0, 5.0], dtype=exp_dtype, name="x")
         exp = tm.box_expected(exp, box)
         tm.assert_equal(result, exp)
 
         result = idx / 2.0
-        assert result.dtype == "f8" and isinstance(result, box)
-        exp = Index([0.5, 1.0, 1.5, 2.0, 2.5], dtype=np.float64, name="x")
+        if result.dtype.kind == "c":
+            assert result.dtype == dtype and isinstance(result, box)
+            exp_dtype = dtype
+        else:
+            assert result.dtype == "f8" and isinstance(result, box)
+            exp_dtype = np.float64
+        exp = Index([0.5, 1.0, 1.5, 2.0, 2.5], dtype=exp_dtype, name="x")
         exp = tm.box_expected(exp, box)
         tm.assert_equal(result, exp)
 
@@ -1367,7 +1397,7 @@ def test_numeric_compat2_floordiv(self, idx, div, expected):
         # __floordiv__
         tm.assert_index_equal(idx // div, expected, exact=True)
 
-    @pytest.mark.parametrize("dtype", [np.int64, np.float64])
+    @pytest.mark.parametrize("dtype", [np.int64, np.float64, np.complex128])
     @pytest.mark.parametrize("delta", [1, 0, -1])
     def test_addsub_arithmetic(self, dtype, delta):
         # GH#8142

pandas/tests/extension/base/dim2.py

@@ -8,6 +8,7 @@
 
 from pandas.core.dtypes.common import (
     is_bool_dtype,
+    is_complex_dtype,
     is_integer_dtype,
 )
 
@@ -272,6 +273,9 @@ def get_reduction_result_dtype(dtype):
                 data = data.astype("Float64")
             if method == "mean":
                 tm.assert_extension_array_equal(result, data)
+            elif is_complex_dtype(data) and method in ["std", "var"]:
+                # std and var produce real-only results
+                tm.assert_extension_array_equal(result, data - data, check_dtype=False)
             else:
                 tm.assert_extension_array_equal(result, data - data)
 

pandas/tests/extension/base/io.py

@@ -9,7 +9,13 @@
 
 class BaseParsingTests:
     @pytest.mark.parametrize("engine", ["c", "python"])
-    def test_EA_types(self, engine, data):
+    def test_EA_types(self, engine, data, request):
+        if engine == "c" and data.dtype.kind == "c":
+            request.node.add_marker(
+                pytest.mark.xfail(
+                    reason=f"engine '{engine}' cannot parse the dtype {data.dtype.name}"
+                )
+            )
         df = pd.DataFrame({"with_dtype": pd.Series(data, dtype=str(data.dtype))})
         csv_output = df.to_csv(index=False, na_rep=np.nan)
         result = pd.read_csv(

pandas/tests/extension/base/ops.py

@@ -19,7 +19,7 @@ class BaseOpsUtil:
     divmod_exc: type[Exception] | None = TypeError
 
     def _get_expected_exception(
-        self, op_name: str, obj, other
+        self, op_name: str, obj, other, request
     ) -> type[Exception] | None:
         # Find the Exception, if any we expect to raise calling
         #  obj.__op_name__(other)
@@ -54,8 +54,8 @@ def get_op_from_name(self, op_name: str):
     #  case that still requires overriding _check_op or _combine, please let
     #  us know at github.com/pandas-dev/pandas/issues
     @final
-    def check_opname(self, ser: pd.Series, op_name: str, other):
-        exc = self._get_expected_exception(op_name, ser, other)
+    def check_opname(self, ser: pd.Series, op_name: str, other, request):
+        exc = self._get_expected_exception(op_name, ser, other, request)
         op = self.get_op_from_name(op_name)
 
         self._check_op(ser, op, other, op_name, exc)
@@ -91,12 +91,12 @@ def _check_op(
 
     # see comment on check_opname
     @final
-    def _check_divmod_op(self, ser: pd.Series, op, other):
+    def _check_divmod_op(self, ser: pd.Series, op, other, request):
         # check that divmod behavior matches behavior of floordiv+mod
         if op is divmod:
-            exc = self._get_expected_exception("__divmod__", ser, other)
+            exc = self._get_expected_exception("__divmod__", ser, other, request)
         else:
-            exc = self._get_expected_exception("__rdivmod__", ser, other)
+            exc = self._get_expected_exception("__rdivmod__", ser, other, request)
         if exc is None:
             result_div, result_mod = op(ser, other)
             if op is divmod:
@@ -128,53 +128,53 @@ class BaseArithmeticOpsTests(BaseOpsUtil):
     series_array_exc: type[Exception] | None = TypeError
     divmod_exc: type[Exception] | None = TypeError
 
-    def test_arith_series_with_scalar(self, data, all_arithmetic_operators):
+    def test_arith_series_with_scalar(self, data, all_arithmetic_operators, request):
         # series & scalar
         if all_arithmetic_operators == "__rmod__" and is_string_dtype(data.dtype):
             pytest.skip("Skip testing Python string formatting")
 
         op_name = all_arithmetic_operators
         ser = pd.Series(data)
-        self.check_opname(ser, op_name, ser.iloc[0])
+        self.check_opname(ser, op_name, ser.iloc[0], request)
 
-    def test_arith_frame_with_scalar(self, data, all_arithmetic_operators):
+    def test_arith_frame_with_scalar(self, data, all_arithmetic_operators, request):
         # frame & scalar
         if all_arithmetic_operators == "__rmod__" and is_string_dtype(data.dtype):
             pytest.skip("Skip testing Python string formatting")
 
         op_name = all_arithmetic_operators
         df = pd.DataFrame({"A": data})
-        self.check_opname(df, op_name, data[0])
+        self.check_opname(df, op_name, data[0], request)
 
-    def test_arith_series_with_array(self, data, all_arithmetic_operators):
+    def test_arith_series_with_array(self, data, all_arithmetic_operators, request):
         # ndarray & other series
         op_name = all_arithmetic_operators
         ser = pd.Series(data)
-        self.check_opname(ser, op_name, pd.Series([ser.iloc[0]] * len(ser)))
+        self.check_opname(ser, op_name, pd.Series([ser.iloc[0]] * len(ser)), request)
 
-    def test_divmod(self, data):
+    def test_divmod(self, data, request):
         ser = pd.Series(data)
-        self._check_divmod_op(ser, divmod, 1)
-        self._check_divmod_op(1, ops.rdivmod, ser)
+        self._check_divmod_op(ser, divmod, 1, request)
+        self._check_divmod_op(1, ops.rdivmod, ser, request)
 
-    def test_divmod_series_array(self, data, data_for_twos):
+    def test_divmod_series_array(self, data, data_for_twos, request):
         ser = pd.Series(data)
-        self._check_divmod_op(ser, divmod, data)
+        self._check_divmod_op(ser, divmod, data, request)
 
         other = data_for_twos
-        self._check_divmod_op(other, ops.rdivmod, ser)
+        self._check_divmod_op(other, ops.rdivmod, ser, request)
 
         other = pd.Series(other)
-        self._check_divmod_op(other, ops.rdivmod, ser)
+        self._check_divmod_op(other, ops.rdivmod, ser, request)
 
-    def test_add_series_with_extension_array(self, data):
+    def test_add_series_with_extension_array(self, data, request):
         # Check adding an ExtensionArray to a Series of the same dtype matches
         # the behavior of adding the arrays directly and then wrapping in a
         # Series.
 
         ser = pd.Series(data)
 
-        exc = self._get_expected_exception("__add__", ser, data)
+        exc = self._get_expected_exception("__add__", ser, data, request)
         if exc is not None:
             with pytest.raises(exc):
                 ser + data

pandas/tests/extension/base/setitem.py

@@ -337,7 +337,8 @@ def test_setitem_slice_array(self, data):
 
     def test_setitem_scalar_key_sequence_raise(self, data):
         arr = data[:5].copy()
-        with pytest.raises(ValueError):
+        # complex128 data raises TypeError; other numeric types raise ValueError
+        with pytest.raises((ValueError, TypeError)):
             arr[0] = arr[[0, 1]]
 
     def test_setitem_preserves_views(self, data):
@@ -438,7 +439,7 @@ def test_setitem_invalid(self, data, invalid_scalar):
             data[:] = invalid_scalar
 
     def test_setitem_2d_values(self, data):
-        # GH50085
+        # GH54445
         original = data.copy()
         df = pd.DataFrame({"a": data, "b": data})
         df.loc[[0, 1], :] = df.loc[[1, 0], :].values