Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add Dataset.drop_dims #2767

Merged
merged 2 commits into from
Mar 3, 2019
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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
3 changes: 3 additions & 0 deletions doc/whats-new.rst
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,9 @@ Enhancements
serialization warnings raised if dates from a standard calendar are found to
be outside the :py:class:`pandas.Timestamp`-valid range (:issue:`2754`). By
`Spencer Clark <https://github.com/spencerkclark>`_.

- 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 @@ -2797,6 +2797,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