Add Dataset.drop_dims (pydata#2767)

* ENH: Add Dataset.drop_dims()

* Drops full dimensions and any corresponding variables in a
  Dataset
* Fixes GH1949

* DOC: Add Dataset.drop_dims() documentation

Author: kmsquire

doc/api.rst

@@ -87,6 +87,7 @@ Dataset contents
    Dataset.swap_dims
    Dataset.expand_dims
    Dataset.drop
+   Dataset.drop_dims
    Dataset.set_coords
    Dataset.reset_coords
 

doc/data-structures.rst

@@ -408,6 +408,13 @@ operations keep around coordinates:
     list(ds[['x']])
     list(ds.drop('temperature'))
 
+To remove a dimension, you can use :py:meth:`~xarray.Dataset.drop_dims` method.
+Any variables using that dimension are dropped:
+
+.. ipython:: python
+
+    ds.drop_dims('time')
+
 As an alternate to dictionary-like modifications, you can use
 :py:meth:`~xarray.Dataset.assign` and :py:meth:`~xarray.Dataset.assign_coords`.
 These methods return a new dataset with additional (or replaced) or values:

doc/indexing.rst

@@ -229,8 +229,8 @@ arrays). However, you can do normal indexing with dimension names:
 Using indexing to *assign* values to a subset of dataset (e.g.,
 ``ds[dict(space=0)] = 1``) is not yet supported.
 
-Dropping labels
----------------
+Dropping labels and dimensions
+------------------------------
 
 The :py:meth:`~xarray.Dataset.drop` method returns a new object with the listed
 index labels along a dimension dropped:
@@ -241,6 +241,13 @@ index labels along a dimension dropped:
 
 ``drop`` is both a ``Dataset`` and ``DataArray`` method.
 
+Use :py:meth:`~xarray.Dataset.drop_dims` to drop a full dimension from a Dataset.
+Any variables with these dimensions are also dropped:
+
+.. ipython:: python
+
+    ds.drop_dims('time')
+
 
 .. _masking with where:
 

doc/whats-new.rst

@@ -83,6 +83,8 @@ Enhancements
   with size > 1. (:issue:`2710`)
   By `Martin Pletcher <https://github.com/pletchm>`_.
 
+- Added :py:meth:`~xarray.Dataset.drop_dims` (:issue:`1949`).
+  By `Kevin Squire <https://github.com/kmsquire>`_.
 
 Bug fixes
 ~~~~~~~~~

xarray/core/dataset.py

@@ -2839,6 +2839,37 @@ def _drop_vars(self, names):
         coord_names = set(k for k in self._coord_names if k in variables)
         return self._replace_vars_and_dims(variables, coord_names)
 
+    def drop_dims(self, drop_dims):
+        """Drop dimensions and associated variables from this dataset.
+
+        Parameters
+        ----------
+        drop_dims : str or list
+            Dimension or dimensions to drop.
+
+        Returns
+        -------
+        obj : Dataset
+            The dataset without the given dimensions (or any variables
+            containing those dimensions)
+        """
+        if utils.is_scalar(drop_dims):
+            drop_dims = [drop_dims]
+
+        missing_dimensions = [d for d in drop_dims if d not in self.dims]
+        if missing_dimensions:
+            raise ValueError('Dataset does not contain the dimensions: %s'
+                             % missing_dimensions)
+
+        drop_vars = set(k for k, v in self._variables.items()
+                        for d in v.dims if d in drop_dims)
+
+        variables = OrderedDict((k, v) for k, v in self._variables.items()
+                                if k not in drop_vars)
+        coord_names = set(k for k in self._coord_names if k in variables)
+
+        return self._replace_with_new_dims(variables, coord_names)
+
     def transpose(self, *dims):
         """Return a new Dataset object with all array dimensions transposed.
 

xarray/tests/test_dataset.py

@@ -1863,6 +1863,26 @@ def test_drop_index_labels(self):
                 ValueError, 'does not have coordinate labels'):
             data.drop(1, 'y')
 
+    def test_drop_dims(self):
+        data = xr.Dataset({'A': (['x', 'y'], np.random.randn(2, 3)),
+                           'B': ('x', np.random.randn(2)),
+                           'x': ['a', 'b'], 'z': np.pi})
+
+        actual = data.drop_dims('x')
+        expected = data.drop(['A', 'B', 'x'])
+        assert_identical(expected, actual)
+
+        actual = data.drop_dims('y')
+        expected = data.drop('A')
+        assert_identical(expected, actual)
+
+        actual = data.drop_dims(['x', 'y'])
+        expected = data.drop(['A', 'B', 'x'])
+        assert_identical(expected, actual)
+
+        with pytest.raises((ValueError, KeyError)):
+            data.drop_dims('z')  # not a dimension
+
     def test_copy(self):
         data = create_test_data()