diff --git a/doc/api.rst b/doc/api.rst index 552582a553f..00b33959eed 100644 --- a/doc/api.rst +++ b/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 diff --git a/doc/data-structures.rst b/doc/data-structures.rst index 618ccccff3e..a8887471ec7 100644 --- a/doc/data-structures.rst +++ b/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: diff --git a/doc/indexing.rst b/doc/indexing.rst index 77ec7428991..9af9db227bc 100644 --- a/doc/indexing.rst +++ b/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: diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 57b6d197ce0..b1c74aca740 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -83,6 +83,8 @@ Enhancements with size > 1. (:issue:`2710`) By `Martin Pletcher `_. +- Added :py:meth:`~xarray.Dataset.drop_dims` (:issue:`1949`). + By `Kevin Squire `_. Bug fixes ~~~~~~~~~ diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index 4373d986733..5296e9834e9 100644 --- a/xarray/core/dataset.py +++ b/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. diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py index e258e799461..918725f8be6 100644 --- a/xarray/tests/test_dataset.py +++ b/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()