Repository: pydata/xarray

Issue:
swap_dims should support dimension names that are not existing variables
#### Code Sample, a copy-pastable example if possible

```python
input_ds = xarray.Dataset({'foo': ('x', [1, 2])}, {'x': [0, 1]})
input_ds.swap_dims({'x': 'z'})
```
#### Problem description

Currently this results in the error `KeyError: 'z'`

#### Expected Output

We now support dimensions without associated coordinate variables. So `swap_dims()` should be able to create new dimensions (e.g., `z` in this example) even if there isn't already a coordinate variable.

Discussion:

Patch:
diff --git a/doc/whats-new.rst b/doc/whats-new.rst
index 00d1c507..351424fb 100644
--- a/doc/whats-new.rst
+++ b/doc/whats-new.rst
@@ -37,6 +37,9 @@ New Features
 - Added the ``count`` reduction method to both :py:class:`~core.rolling.DatasetCoarsen`
   and :py:class:`~core.rolling.DataArrayCoarsen` objects. (:pull:`3500`)
   By `Deepak Cherian <https://github.com/dcherian>`_
+- :py:meth:`Dataset.swap_dims` and :py:meth:`DataArray.swap_dims`
+  now allow swapping to dimension names that don't exist yet. (:pull:`3636`)
+  By `Justus Magin <https://github.com/keewis>`_.
 - Extend :py:class:`core.accessor_dt.DatetimeAccessor` properties 
   and support `.dt` accessor for timedelta 
   via :py:class:`core.accessor_dt.TimedeltaAccessor` (:pull:`3612`)
@@ -62,6 +65,8 @@ Bug fixes
   By `Tom Augspurger <https://github.com/TomAugspurger>`_.
 - Ensure :py:meth:`Dataset.quantile`, :py:meth:`DataArray.quantile` issue the correct error
   when ``q`` is out of bounds (:issue:`3634`) by `Mathias Hauser <https://github.com/mathause>`_.
+- :py:meth:`Dataset.rename`, :py:meth:`DataArray.rename` now check for conflicts with
+  MultiIndex level names.
 
 Documentation
 ~~~~~~~~~~~~~
diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py
index 31aa4da5..cbd8d243 100644
--- a/xarray/core/dataarray.py
+++ b/xarray/core/dataarray.py
@@ -1480,8 +1480,7 @@ class DataArray(AbstractArray, DataWithCoords):
         ----------
         dims_dict : dict-like
             Dictionary whose keys are current dimension names and whose values
-            are new names. Each value must already be a coordinate on this
-            array.
+            are new names.
 
         Returns
         -------
@@ -1504,6 +1503,13 @@ class DataArray(AbstractArray, DataWithCoords):
         Coordinates:
             x        (y) <U1 'a' 'b'
           * y        (y) int64 0 1
+        >>> arr.swap_dims({"x": "z"})
+        <xarray.DataArray (z: 2)>
+        array([0, 1])
+        Coordinates:
+            x        (z) <U1 'a' 'b'
+            y        (z) int64 0 1
+        Dimensions without coordinates: z
 
         See Also
         --------
diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py
index 6be06fed..607350c8 100644
--- a/xarray/core/dataset.py
+++ b/xarray/core/dataset.py
@@ -89,7 +89,13 @@ from .utils import (
     is_scalar,
     maybe_wrap_array,
 )
-from .variable import IndexVariable, Variable, as_variable, broadcast_variables
+from .variable import (
+    IndexVariable,
+    Variable,
+    as_variable,
+    broadcast_variables,
+    assert_unique_multiindex_level_names,
+)
 
 if TYPE_CHECKING:
     from ..backends import AbstractDataStore, ZarrStore
@@ -2780,6 +2786,7 @@ class Dataset(Mapping, ImplementsDatasetReduce, DataWithCoords):
         variables, coord_names, dims, indexes = self._rename_all(
             name_dict=name_dict, dims_dict=name_dict
         )
+        assert_unique_multiindex_level_names(variables)
         return self._replace(variables, coord_names, dims=dims, indexes=indexes)
 
     def rename_dims(
@@ -2868,8 +2875,7 @@ class Dataset(Mapping, ImplementsDatasetReduce, DataWithCoords):
         ----------
         dims_dict : dict-like
             Dictionary whose keys are current dimension names and whose values
-            are new names. Each value must already be a variable in the
-            dataset.
+            are new names.
 
         Returns
         -------
@@ -2898,6 +2904,16 @@ class Dataset(Mapping, ImplementsDatasetReduce, DataWithCoords):
         Data variables:
             a        (y) int64 5 7
             b        (y) float64 0.1 2.4
+        >>> ds.swap_dims({"x": "z"})
+        <xarray.Dataset>
+        Dimensions:  (z: 2)
+        Coordinates:
+            x        (z) <U1 'a' 'b'
+            y        (z) int64 0 1
+        Dimensions without coordinates: z
+        Data variables:
+            a        (z) int64 5 7
+            b        (z) float64 0.1 2.4
 
         See Also
         --------
@@ -2914,7 +2930,7 @@ class Dataset(Mapping, ImplementsDatasetReduce, DataWithCoords):
                     "cannot swap from dimension %r because it is "
                     "not an existing dimension" % k
                 )
-            if self.variables[v].dims != (k,):
+            if v in self.variables and self.variables[v].dims != (k,):
                 raise ValueError(
                     "replacement dimension %r is not a 1D "
                     "variable along the old dimension %r" % (v, k)
@@ -2923,7 +2939,7 @@ class Dataset(Mapping, ImplementsDatasetReduce, DataWithCoords):
         result_dims = {dims_dict.get(dim, dim) for dim in self.dims}
 
         coord_names = self._coord_names.copy()
-        coord_names.update(dims_dict.values())
+        coord_names.update({dim for dim in dims_dict.values() if dim in self.variables})
 
         variables: Dict[Hashable, Variable] = {}
         indexes: Dict[Hashable, pd.Index] = {}
diff --git a/xarray/core/parallel.py b/xarray/core/parallel.py
index dd6c6733..e4fb5803 100644
--- a/xarray/core/parallel.py
+++ b/xarray/core/parallel.py
@@ -154,6 +154,48 @@ def map_blocks(
     --------
     dask.array.map_blocks, xarray.apply_ufunc, xarray.Dataset.map_blocks,
     xarray.DataArray.map_blocks
+
+    Examples
+    --------
+
+    Calculate an anomaly from climatology using ``.groupby()``. Using
+    ``xr.map_blocks()`` allows for parallel operations with knowledge of ``xarray``,
+    its indices, and its methods like ``.groupby()``.
+
+    >>> def calculate_anomaly(da, groupby_type='time.month'):
+    ...     # Necessary workaround to xarray's check with zero dimensions
+    ...     # https://github.com/pydata/xarray/issues/3575
+    ...     if sum(da.shape) == 0:
+    ...         return da
+    ...     gb = da.groupby(groupby_type)
+    ...     clim = gb.mean(dim='time')
+    ...     return gb - clim
+    >>> time = xr.cftime_range('1990-01', '1992-01', freq='M')
+    >>> np.random.seed(123)
+    >>> array = xr.DataArray(np.random.rand(len(time)),
+    ...                      dims="time", coords=[time]).chunk()
+    >>> xr.map_blocks(calculate_anomaly, array).compute()
+    <xarray.DataArray (time: 24)>
+    array([ 0.12894847,  0.11323072, -0.0855964 , -0.09334032,  0.26848862,
+            0.12382735,  0.22460641,  0.07650108, -0.07673453, -0.22865714,
+           -0.19063865,  0.0590131 , -0.12894847, -0.11323072,  0.0855964 ,
+            0.09334032, -0.26848862, -0.12382735, -0.22460641, -0.07650108,
+            0.07673453,  0.22865714,  0.19063865, -0.0590131 ])
+    Coordinates:
+      * time     (time) object 1990-01-31 00:00:00 ... 1991-12-31 00:00:00
+
+    Note that one must explicitly use ``args=[]`` and ``kwargs={}`` to pass arguments
+    to the function being applied in ``xr.map_blocks()``:
+
+    >>> xr.map_blocks(calculate_anomaly, array, kwargs={'groupby_type': 'time.year'})
+    <xarray.DataArray (time: 24)>
+    array([ 0.15361741, -0.25671244, -0.31600032,  0.008463  ,  0.1766172 ,
+           -0.11974531,  0.43791243,  0.14197797, -0.06191987, -0.15073425,
+           -0.19967375,  0.18619794, -0.05100474, -0.42989909, -0.09153273,
+            0.24841842, -0.30708526, -0.31412523,  0.04197439,  0.0422506 ,
+            0.14482397,  0.35985481,  0.23487834,  0.12144652])
+    Coordinates:
+        * time     (time) object 1990-01-31 00:00:00 ... 1991-12-31 00:00:00
     """
 
     def _wrapper(func, obj, to_array, args, kwargs):


Test patch:
diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py
index f957316d..4189c3b5 100644
--- a/xarray/tests/test_dataarray.py
+++ b/xarray/tests/test_dataarray.py
@@ -1530,6 +1530,11 @@ class TestDataArray:
         actual = array.swap_dims({"x": "y"})
         assert_identical(expected, actual)
 
+        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)
+
     def test_expand_dims_error(self):
         array = DataArray(
             np.random.randn(3, 4),
diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py
index 7db19116..48d8c25b 100644
--- a/xarray/tests/test_dataset.py
+++ b/xarray/tests/test_dataset.py
@@ -2461,6 +2461,14 @@ class TestDataset:
         with pytest.raises(ValueError):
             original.rename_vars(names_dict_bad)
 
+    def test_rename_multiindex(self):
+        mindex = pd.MultiIndex.from_tuples(
+            [([1, 2]), ([3, 4])], names=["level0", "level1"]
+        )
+        data = Dataset({}, {"x": mindex})
+        with raises_regex(ValueError, "conflicting MultiIndex"):
+            data.rename({"x": "level0"})
+
     @requires_cftime
     def test_rename_does_not_change_CFTimeIndex_type(self):
         # make sure CFTimeIndex is not converted to DatetimeIndex #3522
@@ -2525,6 +2533,12 @@ class TestDataset:
         with raises_regex(ValueError, "replacement dimension"):
             original.swap_dims({"x": "z"})
 
+        expected = Dataset(
+            {"y": ("u", list("abc")), "z": 42}, coords={"x": ("u", [1, 2, 3])}
+        )
+        actual = original.swap_dims({"x": "u"})
+        assert_identical(expected, actual)
+
     def test_expand_dims_error(self):
         original = Dataset(
             {