ENH: Index[complex] (#45256)

jbrockmendel · web-flow · commit b77797cde579 · 2022-01-08T17:18:59.000-05:00
diff --git a/pandas/_libs/index.pyi b/pandas/_libs/index.pyi
@@ -29,6 +29,8 @@ class IndexEngine:
 
 class Float64Engine(IndexEngine): ...
 class Float32Engine(IndexEngine): ...
+class Complex128Engine(IndexEngine): ...
+class Complex64Engine(IndexEngine): ...
 class Int64Engine(IndexEngine): ...
 class Int32Engine(IndexEngine): ...
 class Int16Engine(IndexEngine): ...
diff --git a/pandas/_libs/index_class_helper.pxi.in b/pandas/_libs/index_class_helper.pxi.in
@@ -21,6 +21,8 @@ dtypes = [('Float64', 'float64'),
           ('UInt32', 'uint32'),
           ('UInt16', 'uint16'),
           ('UInt8', 'uint8'),
+          ('Complex64', 'complex64'),
+          ('Complex128', 'complex128'),
           ]
 }}
 
@@ -33,18 +35,25 @@ cdef class {{name}}Engine(IndexEngine):
         return _hash.{{name}}HashTable(n)
 
     cdef _check_type(self, object val):
-    {{if name not in {'Float64', 'Float32'} }}
+    {{if name not in {'Float64', 'Float32', 'Complex64', 'Complex128'} }}
         if not util.is_integer_object(val):
             raise KeyError(val)
         {{if name.startswith("U")}}
         if val < 0:
             # cannot have negative values with unsigned int dtype
             raise KeyError(val)
         {{endif}}
-    {{else}}
+    {{elif name not in {'Complex64', 'Complex128'} }}
         if not util.is_integer_object(val) and not util.is_float_object(val):
             # in particular catch bool and avoid casting True -> 1.0
             raise KeyError(val)
+    {{else}}
+        if (not util.is_integer_object(val)
+            and not util.is_float_object(val)
+            and not util.is_complex_object(val)
+        ):
+            # in particular catch bool and avoid casting True -> 1.0
+            raise KeyError(val)
     {{endif}}
 
 
diff --git a/pandas/conftest.py b/pandas/conftest.py
@@ -539,6 +539,8 @@ def _create_mi_with_dt64tz_level():
     "uint": tm.makeUIntIndex(100),
     "range": tm.makeRangeIndex(100),
     "float": tm.makeFloatIndex(100),
+    "complex64": tm.makeFloatIndex(100).astype("complex64"),
+    "complex128": tm.makeFloatIndex(100).astype("complex128"),
     "num_int64": tm.makeNumericIndex(100, dtype="int64"),
     "num_int32": tm.makeNumericIndex(100, dtype="int32"),
     "num_int16": tm.makeNumericIndex(100, dtype="int16"),
diff --git a/pandas/core/indexes/base.py b/pandas/core/indexes/base.py
@@ -487,6 +487,8 @@ def __new__(
             if data.dtype.kind in ["i", "u", "f"]:
                 # maybe coerce to a sub-class
                 arr = data
+            elif data.dtype.kind in ["c"]:
+                arr = np.asarray(data)
             else:
                 arr = com.asarray_tuplesafe(data, dtype=_dtype_obj)
 
@@ -614,7 +616,9 @@ def _dtype_to_subclass(cls, dtype: DtypeObj):
             # NB: assuming away MultiIndex
             return Index
 
-        elif issubclass(dtype.type, (str, bool, np.bool_)):
+        elif issubclass(
+            dtype.type, (str, bool, np.bool_, complex, np.complex64, np.complex128)
+        ):
             return Index
 
         raise NotImplementedError(dtype)
@@ -858,6 +862,11 @@ def _engine(
             # TODO(ExtensionIndex): use libindex.ExtensionEngine(self._values)
             return libindex.ObjectEngine(self._get_engine_target())
 
+        elif self.values.dtype == np.complex64:
+            return libindex.Complex64Engine(self._get_engine_target())
+        elif self.values.dtype == np.complex128:
+            return libindex.Complex128Engine(self._get_engine_target())
+
         # to avoid a reference cycle, bind `target_values` to a local variable, so
         # `self` is not passed into the lambda.
         target_values = self._get_engine_target()
@@ -5980,8 +5989,6 @@ def _find_common_type_compat(self, target) -> DtypeObj:
                 # FIXME: find_common_type incorrect with Categorical GH#38240
                 # FIXME: some cases where float64 cast can be lossy?
                 dtype = np.dtype(np.float64)
-        if dtype.kind == "c":
-            dtype = _dtype_obj
         return dtype
 
     @final
@@ -7120,7 +7127,7 @@ def _maybe_cast_data_without_dtype(
             FutureWarning,
             stacklevel=3,
         )
-    if result.dtype.kind in ["b", "c"]:
+    if result.dtype.kind in ["b"]:
         return subarr
     result = ensure_wrapped_if_datetimelike(result)
     return result
diff --git a/pandas/core/indexes/numeric.py b/pandas/core/indexes/numeric.py
@@ -114,6 +114,8 @@ def _can_hold_na(self) -> bool:  # type: ignore[override]
         np.dtype(np.uint64): libindex.UInt64Engine,
         np.dtype(np.float32): libindex.Float32Engine,
         np.dtype(np.float64): libindex.Float64Engine,
+        np.dtype(np.complex64): libindex.Complex64Engine,
+        np.dtype(np.complex128): libindex.Complex128Engine,
     }
 
     @property
@@ -128,6 +130,7 @@ def inferred_type(self) -> str:
             "i": "integer",
             "u": "integer",
             "f": "floating",
+            "c": "complex",
         }[self.dtype.kind]
 
     def __new__(cls, data=None, dtype: Dtype | None = None, copy=False, name=None):
diff --git a/pandas/tests/arrays/categorical/test_constructors.py b/pandas/tests/arrays/categorical/test_constructors.py
@@ -676,7 +676,6 @@ def test_construction_with_ordered(self, ordered):
         cat = Categorical([0, 1, 2], ordered=ordered)
         assert cat.ordered == bool(ordered)
 
-    @pytest.mark.xfail(reason="Imaginary values not supported in Categorical")
     def test_constructor_imaginary(self):
         values = [1, 2, 3 + 1j]
         c1 = Categorical(values)
diff --git a/pandas/tests/base/test_misc.py b/pandas/tests/base/test_misc.py
@@ -137,14 +137,19 @@ def test_memory_usage_components_narrow_series(dtype):
     assert total_usage == non_index_usage + index_usage
 
 
-def test_searchsorted(index_or_series_obj):
+def test_searchsorted(index_or_series_obj, request):
     # numpy.searchsorted calls obj.searchsorted under the hood.
     # See gh-12238
     obj = index_or_series_obj
 
     if isinstance(obj, pd.MultiIndex):
         # See gh-14833
         pytest.skip("np.searchsorted doesn't work on pd.MultiIndex")
+    if obj.dtype.kind == "c" and isinstance(obj, Index):
+        # TODO: Should Series cases also raise? Looks like they use numpy
+        #  comparison semantics https://github.com/numpy/numpy/issues/15981
+        mark = pytest.mark.xfail(reason="complex objects are not comparable")
+        request.node.add_marker(mark)
 
     max_obj = max(obj, default=0)
     index = np.searchsorted(obj, max_obj)
diff --git a/pandas/tests/groupby/test_groupby.py b/pandas/tests/groupby/test_groupby.py
@@ -1054,15 +1054,14 @@ def test_groupby_complex_numbers():
     )
     expected = DataFrame(
         np.array([1, 1, 1], dtype=np.int64),
-        index=Index([(1 + 1j), (1 + 2j), (1 + 0j)], dtype="object", name="b"),
+        index=Index([(1 + 1j), (1 + 2j), (1 + 0j)], name="b"),
         columns=Index(["a"], dtype="object"),
     )
     result = df.groupby("b", sort=False).count()
     tm.assert_frame_equal(result, expected)
 
     # Sorted by the magnitude of the complex numbers
-    # Complex Index dtype is cast to object
-    expected.index = Index([(1 + 0j), (1 + 1j), (1 + 2j)], dtype="object", name="b")
+    expected.index = Index([(1 + 0j), (1 + 1j), (1 + 2j)], name="b")
     result = df.groupby("b", sort=True).count()
     tm.assert_frame_equal(result, expected)
 
diff --git a/pandas/tests/indexes/multi/test_setops.py b/pandas/tests/indexes/multi/test_setops.py
@@ -525,11 +525,17 @@ def test_union_nan_got_duplicated():
     tm.assert_index_equal(result, mi2)
 
 
-def test_union_duplicates(index):
+def test_union_duplicates(index, request):
     # GH#38977
     if index.empty or isinstance(index, (IntervalIndex, CategoricalIndex)):
         # No duplicates in empty indexes
         return
+    if index.dtype.kind == "c":
+        mark = pytest.mark.xfail(
+            reason="sort_values() call raises bc complex objects are not comparable"
+        )
+        request.node.add_marker(mark)
+
     values = index.unique().values.tolist()
     mi1 = MultiIndex.from_arrays([values, [1] * len(values)])
     mi2 = MultiIndex.from_arrays([[values[0]] + values, [1] * (len(values) + 1)])
diff --git a/pandas/tests/indexes/test_any_index.py b/pandas/tests/indexes/test_any_index.py
@@ -46,8 +46,14 @@ def test_mutability(index):
         index[0] = index[0]
 
 
-def test_map_identity_mapping(index):
+def test_map_identity_mapping(index, request):
     # GH#12766
+    if index.dtype == np.complex64:
+        mark = pytest.mark.xfail(
+            reason="maybe_downcast_to_dtype doesn't handle complex"
+        )
+        request.node.add_marker(mark)
+
     result = index.map(lambda x: x)
     tm.assert_index_equal(result, index, exact="equiv")
 
diff --git a/pandas/tests/indexes/test_base.py b/pandas/tests/indexes/test_base.py
@@ -526,14 +526,19 @@ def test_map_dictlike_simple(self, mapper):
             lambda values, index: Series(values, index),
         ],
     )
-    def test_map_dictlike(self, index, mapper):
+    def test_map_dictlike(self, index, mapper, request):
         # GH 12756
         if isinstance(index, CategoricalIndex):
             # Tested in test_categorical
             return
         elif not index.is_unique:
             # Cannot map duplicated index
             return
+        if index.dtype == np.complex64 and not isinstance(mapper(index, index), Series):
+            mark = pytest.mark.xfail(
+                reason="maybe_downcast_to_dtype doesn't handle complex"
+            )
+            request.node.add_marker(mark)
 
         rng = np.arange(len(index), 0, -1)
 
@@ -655,7 +660,8 @@ def test_format_missing(self, vals, nulls_fixture):
         # 2845
         vals = list(vals)  # Copy for each iteration
         vals.append(nulls_fixture)
-        index = Index(vals)
+        index = Index(vals, dtype=object)
+        # TODO: case with complex dtype?
 
         formatted = index.format()
         expected = [str(index[0]), str(index[1]), str(index[2]), "NaN"]
diff --git a/pandas/tests/indexes/test_common.py b/pandas/tests/indexes/test_common.py
@@ -386,13 +386,20 @@ def test_astype_preserves_name(self, index, dtype):
         if dtype in ["int64", "uint64"]:
             if needs_i8_conversion(index.dtype):
                 warn = FutureWarning
+            elif index.dtype.kind == "c":
+                # imaginary components discarded
+                warn = np.ComplexWarning
         elif (
             isinstance(index, DatetimeIndex)
             and index.tz is not None
             and dtype == "datetime64[ns]"
         ):
             # This astype is deprecated in favor of tz_localize
             warn = FutureWarning
+        elif index.dtype.kind == "c" and dtype == "float64":
+            # imaginary components discarded
+            warn = np.ComplexWarning
+
         try:
             # Some of these conversions cannot succeed so we use a try / except
             with tm.assert_produces_warning(warn):
diff --git a/pandas/tests/indexes/test_numpy_compat.py b/pandas/tests/indexes/test_numpy_compat.py
@@ -54,8 +54,10 @@ def test_numpy_ufuncs_basic(index, func):
         with tm.external_error_raised((TypeError, AttributeError)):
             with np.errstate(all="ignore"):
                 func(index)
-    elif isinstance(index, NumericIndex) or (
-        not isinstance(index.dtype, np.dtype) and index.dtype._is_numeric
+    elif (
+        isinstance(index, NumericIndex)
+        or (not isinstance(index.dtype, np.dtype) and index.dtype._is_numeric)
+        or (index.dtype.kind == "c" and func not in [np.deg2rad, np.rad2deg])
     ):
         # coerces to float (e.g. np.sin)
         with np.errstate(all="ignore"):
@@ -99,8 +101,10 @@ def test_numpy_ufuncs_other(index, func, request):
         with tm.external_error_raised(TypeError):
             func(index)
 
-    elif isinstance(index, NumericIndex) or (
-        not isinstance(index.dtype, np.dtype) and index.dtype._is_numeric
+    elif (
+        isinstance(index, NumericIndex)
+        or (not isinstance(index.dtype, np.dtype) and index.dtype._is_numeric)
+        or (index.dtype.kind == "c" and func is not np.signbit)
     ):
         # Results in bool array
         result = func(index)
diff --git a/pandas/tests/indexes/test_setops.py b/pandas/tests/indexes/test_setops.py
@@ -67,6 +67,25 @@ def test_union_different_types(index_flat, index_flat2, request):
 
     common_dtype = find_common_type([idx1.dtype, idx2.dtype])
 
+    warn = None
+    if not len(idx1) or not len(idx2):
+        pass
+    elif (
+        idx1.dtype.kind == "c"
+        and (
+            idx2.dtype.kind not in ["i", "u", "f", "c"]
+            or not isinstance(idx2.dtype, np.dtype)
+        )
+    ) or (
+        idx2.dtype.kind == "c"
+        and (
+            idx1.dtype.kind not in ["i", "u", "f", "c"]
+            or not isinstance(idx1.dtype, np.dtype)
+        )
+    ):
+        # complex objects non-sortable
+        warn = RuntimeWarning
+
     any_uint64 = idx1.dtype == np.uint64 or idx2.dtype == np.uint64
     idx1_signed = is_signed_integer_dtype(idx1.dtype)
     idx2_signed = is_signed_integer_dtype(idx2.dtype)
@@ -76,8 +95,9 @@ def test_union_different_types(index_flat, index_flat2, request):
     idx1 = idx1.sort_values()
     idx2 = idx2.sort_values()
 
-    res1 = idx1.union(idx2)
-    res2 = idx2.union(idx1)
+    with tm.assert_produces_warning(warn, match="'<' not supported between"):
+        res1 = idx1.union(idx2)
+        res2 = idx2.union(idx1)
 
     if any_uint64 and (idx1_signed or idx2_signed):
         assert res1.dtype == np.dtype("O")
diff --git a/pandas/tests/indexing/test_coercion.py b/pandas/tests/indexing/test_coercion.py
@@ -433,9 +433,6 @@ def test_where_object(self, index_or_series, fill_val, exp_dtype):
     )
     def test_where_int64(self, index_or_series, fill_val, exp_dtype, request):
         klass = index_or_series
-        if klass is pd.Index and exp_dtype is np.complex128:
-            mark = pytest.mark.xfail(reason="Complex Index not supported")
-            request.node.add_marker(mark)
 
         obj = klass([1, 2, 3, 4])
         assert obj.dtype == np.int64
@@ -447,9 +444,6 @@ def test_where_int64(self, index_or_series, fill_val, exp_dtype, request):
     )
     def test_where_float64(self, index_or_series, fill_val, exp_dtype, request):
         klass = index_or_series
-        if klass is pd.Index and exp_dtype is np.complex128:
-            mark = pytest.mark.xfail(reason="Complex Index not supported")
-            request.node.add_marker(mark)
 
         obj = klass([1.1, 2.2, 3.3, 4.4])
         assert obj.dtype == np.float64
@@ -464,8 +458,8 @@ def test_where_float64(self, index_or_series, fill_val, exp_dtype, request):
             (True, object),
         ],
     )
-    def test_where_series_complex128(self, fill_val, exp_dtype):
-        klass = pd.Series  # TODO: use index_or_series once we have Index[complex]
+    def test_where_series_complex128(self, index_or_series, fill_val, exp_dtype):
+        klass = index_or_series
         obj = klass([1 + 1j, 2 + 2j, 3 + 3j, 4 + 4j])
         assert obj.dtype == np.complex128
         self._run_test(obj, fill_val, klass, exp_dtype)
@@ -624,11 +618,6 @@ def test_fillna_float64(self, index_or_series, fill_val, fill_dtype):
         assert obj.dtype == np.float64
 
         exp = klass([1.1, fill_val, 3.3, 4.4])
-        # float + complex -> we don't support a complex Index
-        # complex for Series,
-        # object for Index
-        if fill_dtype == np.complex128 and klass == pd.Index:
-            fill_dtype = object
         self._assert_fillna_conversion(obj, fill_val, exp, fill_dtype)
 
     @pytest.mark.parametrize(
diff --git a/pandas/tests/series/methods/test_value_counts.py b/pandas/tests/series/methods/test_value_counts.py
@@ -216,13 +216,12 @@ def test_value_counts_bool_with_nan(self, ser, dropna, exp):
                 Series([3, 2, 1], index=pd.Index([3j, 1 + 1j, 1], dtype=np.complex128)),
             ),
             (
-                [1 + 1j, 1 + 1j, 1, 3j, 3j, 3j],
+                np.array([1 + 1j, 1 + 1j, 1, 3j, 3j, 3j], dtype=np.complex64),
                 Series([3, 2, 1], index=pd.Index([3j, 1 + 1j, 1], dtype=np.complex64)),
             ),
         ],
     )
     def test_value_counts_complex_numbers(self, input_array, expected):
         # GH 17927
-        # Complex Index dtype is cast to object
         result = Series(input_array).value_counts()
         tm.assert_series_equal(result, expected)
