ENH: Index[complex] (#45845)

* ENH: Index[complex]

* mypy fixup

* whatsnew

Author: jbrockmendel

doc/source/whatsnew/v1.5.0.rst

@@ -38,6 +38,8 @@ Other enhancements
 - :meth:`Series.reset_index` and :meth:`DataFrame.reset_index` now support the argument ``allow_duplicates`` (:issue:`44410`)
 - :meth:`.GroupBy.min` and :meth:`.GroupBy.max` now supports `Numba <https://numba.pydata.org/>`_ execution with the ``engine`` keyword (:issue:`45428`)
 - Implemented a ``bool``-dtype :class:`Index`, passing a bool-dtype arraylike to ``pd.Index`` will now retain ``bool`` dtype instead of casting to ``object`` (:issue:`45061`)
+- Implemented a complex-dtype :class:`Index`, passing a complex-dtype arraylike to ``pd.Index`` will now retain complex dtype instead of casting to ``object`` (:issue:`45845`)
+
 -
 
 .. ---------------------------------------------------------------------------

pandas/_libs/index.pyi

@@ -30,6 +30,8 @@ class IndexEngine:
 
 class Float64Engine(IndexEngine): ...
 class Float32Engine(IndexEngine): ...
+class Complex128Engine(IndexEngine): ...
+class Complex64Engine(IndexEngine): ...
 class Int64Engine(IndexEngine): ...
 class Int32Engine(IndexEngine): ...
 class Int16Engine(IndexEngine): ...

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,7 +35,7 @@ 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):
             if util.is_float_object(val):
                 # Make sure Int64Index.get_loc(2.0) works
@@ -45,10 +47,17 @@ cdef class {{name}}Engine(IndexEngine):
             # 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}}
         return val
 

pandas/conftest.py

@@ -547,6 +547,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"),

pandas/core/indexes/base.py

@@ -505,7 +505,7 @@ def __new__(
             if data.dtype.kind in ["i", "u", "f"]:
                 # maybe coerce to a sub-class
                 arr = data
-            elif data.dtype.kind == "b":
+            elif data.dtype.kind in ["b", "c"]:
                 # No special subclass, and Index._ensure_array won't do this
                 #  for us.
                 arr = np.asarray(data)
@@ -636,7 +636,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)
@@ -881,6 +883,10 @@ def _engine(
         # `self` is not passed into the lambda.
         if target_values.dtype == bool:
             return libindex.BoolEngine(target_values)
+        elif target_values.dtype == np.complex64:
+            return libindex.Complex64Engine(target_values)
+        elif target_values.dtype == np.complex128:
+            return libindex.Complex128Engine(target_values)
 
         # error: Argument 1 to "ExtensionEngine" has incompatible type
         # "ndarray[Any, Any]"; expected "ExtensionArray"
@@ -6162,9 +6168,6 @@ def _find_common_type_compat(self, target) -> DtypeObj:
 
         dtype = find_common_type([self.dtype, target_dtype])
         dtype = common_dtype_categorical_compat([self, target], dtype)
-
-        if dtype.kind == "c":
-            dtype = _dtype_obj
         return dtype
 
     @final
@@ -7308,8 +7311,6 @@ def _maybe_cast_data_without_dtype(
             FutureWarning,
             stacklevel=3,
         )
-    if result.dtype.kind in ["c"]:
-        return subarr
     result = ensure_wrapped_if_datetimelike(result)
     return result
 

pandas/core/indexes/numeric.py

@@ -101,6 +101,8 @@ class NumericIndex(Index):
         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
@@ -115,6 +117,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):

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)

pandas/tests/base/test_misc.py

@@ -160,6 +160,11 @@ def test_searchsorted(request, index_or_series_obj):
                 reason="np.searchsorted doesn't work on pd.MultiIndex: GH 14833"
             )
         )
+    elif 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)

pandas/tests/groupby/test_groupby.py

@@ -1047,15 +1047,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)
 

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)])

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)
     if index.dtype == object and result.dtype == bool:
         assert (index == result).all()

pandas/tests/indexes/test_base.py

@@ -532,14 +532,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)
 
@@ -664,7 +669,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()
         null_repr = "NaN" if isinstance(nulls_fixture, float) else str(nulls_fixture)

pandas/tests/indexes/test_common.py

@@ -389,6 +389,10 @@ def test_astype_preserves_name(self, index, dtype):
         ):
             # This astype is deprecated in favor of tz_localize
             warn = FutureWarning
+        elif index.dtype.kind == "c" and dtype in ["float64", "int64", "uint64"]:
+            # 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):

pandas/tests/indexes/test_numpy_compat.py

@@ -71,6 +71,7 @@ def test_numpy_ufuncs_basic(index, func):
     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])
         or index.dtype == bool
     ):
         # coerces to float (e.g. np.sin)
@@ -122,6 +123,7 @@ def test_numpy_ufuncs_other(index, func):
     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)
         or index.dtype == bool
     ):
         # Results in bool array

pandas/tests/indexes/test_setops.py

@@ -72,6 +72,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)
@@ -81,8 +100,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")

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,9 +458,9 @@ 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]
-        obj = klass([1 + 1j, 2 + 2j, 3 + 3j, 4 + 4j])
+    def test_where_complex128(self, index_or_series, fill_val, exp_dtype):
+        klass = index_or_series
+        obj = klass([1 + 1j, 2 + 2j, 3 + 3j, 4 + 4j], dtype=np.complex128)
         assert obj.dtype == np.complex128
         self._run_test(obj, fill_val, klass, exp_dtype)
 
@@ -608,11 +602,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(
@@ -624,16 +613,12 @@ def test_fillna_float64(self, index_or_series, fill_val, fill_dtype):
             (True, object),
         ],
     )
-    def test_fillna_complex128(self, index_or_series, fill_val, fill_dtype, request):
+    def test_fillna_complex128(self, index_or_series, fill_val, fill_dtype):
         klass = index_or_series
-        if klass is pd.Index:
-            mark = pytest.mark.xfail(reason="No Index[complex]")
-            request.node.add_marker(mark)
-
         obj = klass([1 + 1j, np.nan, 3 + 3j, 4 + 4j], dtype=np.complex128)
         assert obj.dtype == np.complex128
 
-        exp = klass([1 + 1j, fill_val, 3 + 3j, 4 + 4j], dtype=fill_dtype)
+        exp = klass([1 + 1j, fill_val, 3 + 3j, 4 + 4j])
         self._assert_fillna_conversion(obj, fill_val, exp, fill_dtype)
 
     @pytest.mark.parametrize(

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)