Fix swap_dims() index names (issue #3748) (#3752)

* Added test for GH3748

* Rename newly created index in swap_dims() to dim name if not multiindex

Fixes GH3748

* Updated whats-new.rst with pull request information for swap_dims fix

* Move tests for GH3748 into existing swap_dims tests

+ integrated new tests for GH3748 for DataArray into existing swap_dims
  tests
+ added similar tests for Dataset
+ added test for multiindex case

Co-authored-by: Deepak Cherian <[email protected]>

Author: jaicher

doc/whats-new.rst

@@ -34,10 +34,14 @@ New Features
 
 Bug fixes
 ~~~~~~~~~
+
+- Fix :py:meth:`Dataset.swap_dims` and :py:meth:`DataArray.swap_dims` producing
+  index with name reflecting the previous dimension name instead of the new one
+  (:issue:`3748`, :pull:`3752`). By `Joseph K Aicher
+  <https://github.com/jaicher>`_.
 - Use ``dask_array_type`` instead of ``dask_array.Array`` for type
   checking. (:issue:`3779`, :pull:`3787`)
   By `Justus Magin <https://github.com/keewis>`_.
-
 - :py:func:`concat` can now handle coordinate variables only present in one of
   the objects to be concatenated when ``coords="different"``.
   By `Deepak Cherian <https://github.com/dcherian>`_.

xarray/core/dataset.py

@@ -2948,7 +2948,11 @@ def swap_dims(
                 if k in self.indexes:
                     indexes[k] = self.indexes[k]
                 else:
-                    indexes[k] = var.to_index()
+                    new_index = var.to_index()
+                    if new_index.nlevels == 1:
+                        # make sure index name matches dimension name
+                        new_index = new_index.rename(k)
+                    indexes[k] = new_index
             else:
                 var = v.to_base_variable()
             var.dims = dims

xarray/tests/test_dataarray.py

@@ -1536,11 +1536,30 @@ def test_swap_dims(self):
         expected = DataArray(array.values, {"y": list("abc")}, dims="y")
         actual = array.swap_dims({"x": "y"})
         assert_identical(expected, actual)
+        for dim_name in set().union(expected.indexes.keys(), actual.indexes.keys()):
+            pd.testing.assert_index_equal(
+                expected.indexes[dim_name], actual.indexes[dim_name]
+            )
 
         array = DataArray(np.random.randn(3), {"x": list("abc")}, "x")
         expected = DataArray(array.values, {"x": ("y", list("abc"))}, dims="y")
         actual = array.swap_dims({"x": "y"})
         assert_identical(expected, actual)
+        for dim_name in set().union(expected.indexes.keys(), actual.indexes.keys()):
+            pd.testing.assert_index_equal(
+                expected.indexes[dim_name], actual.indexes[dim_name]
+            )
+
+        # multiindex case
+        idx = pd.MultiIndex.from_arrays([list("aab"), list("yzz")], names=["y1", "y2"])
+        array = DataArray(np.random.randn(3), {"y": ("x", idx)}, "x")
+        expected = DataArray(array.values, {"y": idx}, "y")
+        actual = array.swap_dims({"x": "y"})
+        assert_identical(expected, actual)
+        for dim_name in set().union(expected.indexes.keys(), actual.indexes.keys()):
+            pd.testing.assert_index_equal(
+                expected.indexes[dim_name], actual.indexes[dim_name]
+            )
 
     def test_expand_dims_error(self):
         array = DataArray(

xarray/tests/test_dataset.py

@@ -2596,7 +2596,7 @@ def test_swap_dims(self):
         assert_identical(expected, actual)
         assert isinstance(actual.variables["y"], IndexVariable)
         assert isinstance(actual.variables["x"], Variable)
-        assert actual.indexes["y"].equals(pd.Index(list("abc")))
+        pd.testing.assert_index_equal(actual.indexes["y"], expected.indexes["y"])
 
         roundtripped = actual.swap_dims({"y": "x"})
         assert_identical(original.set_coords("y"), roundtripped)
@@ -2612,6 +2612,16 @@ def test_swap_dims(self):
         actual = original.swap_dims({"x": "u"})
         assert_identical(expected, actual)
 
+        # handle multiindex case
+        idx = pd.MultiIndex.from_arrays([list("aab"), list("yzz")], names=["y1", "y2"])
+        original = Dataset({"x": [1, 2, 3], "y": ("x", idx), "z": 42})
+        expected = Dataset({"z": 42}, {"x": ("y", [1, 2, 3]), "y": idx})
+        actual = original.swap_dims({"x": "y"})
+        assert_identical(expected, actual)
+        assert isinstance(actual.variables["y"], IndexVariable)
+        assert isinstance(actual.variables["x"], Variable)
+        pd.testing.assert_index_equal(actual.indexes["y"], expected.indexes["y"])
+
     def test_expand_dims_error(self):
         original = Dataset(
             {