Skip to content

Commit

Permalink
Add Dataset.drop_dims (pydata#2767)
Browse files Browse the repository at this point in the history
* ENH: Add Dataset.drop_dims()

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

* DOC: Add Dataset.drop_dims() documentation
  • Loading branch information
kmsquire authored and pletchm committed Mar 21, 2019
1 parent fa342c7 commit 93185ef
Show file tree
Hide file tree
Showing 6 changed files with 70 additions and 2 deletions.
1 change: 1 addition & 0 deletions doc/api.rst
Original file line number Diff line number Diff line change
Expand Up @@ -87,6 +87,7 @@ Dataset contents
Dataset.swap_dims
Dataset.expand_dims
Dataset.drop
Dataset.drop_dims
Dataset.set_coords
Dataset.reset_coords

Expand Down
7 changes: 7 additions & 0 deletions doc/data-structures.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
11 changes: 9 additions & 2 deletions doc/indexing.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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:

Expand Down
2 changes: 2 additions & 0 deletions doc/whats-new.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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
~~~~~~~~~
Expand Down
31 changes: 31 additions & 0 deletions xarray/core/dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
20 changes: 20 additions & 0 deletions xarray/tests/test_dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand Down

0 comments on commit 93185ef

Please sign in to comment.