Skip to content

Commit ab268de

Browse files
kmsquirepletchm
authored andcommitted
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
1 parent b9663b2 commit ab268de

6 files changed

Lines changed: 75 additions & 2 deletions

File tree

doc/api.rst

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,7 @@ Dataset contents
8787
Dataset.swap_dims
8888
Dataset.expand_dims
8989
Dataset.drop
90+
Dataset.drop_dims
9091
Dataset.set_coords
9192
Dataset.reset_coords
9293

doc/data-structures.rst

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -408,6 +408,13 @@ operations keep around coordinates:
408408
list(ds[['x']])
409409
list(ds.drop('temperature'))
410410
411+
To remove a dimension, you can use :py:meth:`~xarray.Dataset.drop_dims` method.
412+
Any variables using that dimension are dropped:
413+
414+
.. ipython:: python
415+
416+
ds.drop_dims('time')
417+
411418
As an alternate to dictionary-like modifications, you can use
412419
:py:meth:`~xarray.Dataset.assign` and :py:meth:`~xarray.Dataset.assign_coords`.
413420
These methods return a new dataset with additional (or replaced) or values:

doc/indexing.rst

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -229,8 +229,8 @@ arrays). However, you can do normal indexing with dimension names:
229229
Using indexing to *assign* values to a subset of dataset (e.g.,
230230
``ds[dict(space=0)] = 1``) is not yet supported.
231231

232-
Dropping labels
233-
---------------
232+
Dropping labels and dimensions
233+
------------------------------
234234

235235
The :py:meth:`~xarray.Dataset.drop` method returns a new object with the listed
236236
index labels along a dimension dropped:
@@ -241,6 +241,13 @@ index labels along a dimension dropped:
241241
242242
``drop`` is both a ``Dataset`` and ``DataArray`` method.
243243

244+
Use :py:meth:`~xarray.Dataset.drop_dims` to drop a full dimension from a Dataset.
245+
Any variables with these dimensions are also dropped:
246+
247+
.. ipython:: python
248+
249+
ds.drop_dims('time')
250+
244251
245252
.. _masking with where:
246253

doc/whats-new.rst

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,13 @@ Enhancements
7979
be outside the :py:class:`pandas.Timestamp`-valid range (:issue:`2754`). By
8080
`Spencer Clark <https://github.com/spencerkclark>`_.
8181

82+
- Allow ``expand_dims`` method to support inserting/broadcasting dimensions
83+
with size > 1. (:issue:`2710`)
84+
By `Martin Pletcher <https://github.com/pletchm>`_.
85+
86+
- Added :py:meth:`~xarray.Dataset.drop_dims` (:issue:`1949`).
87+
By `Kevin Squire <https://github.com/kmsquire>`_.
88+
8289
Bug fixes
8390
~~~~~~~~~
8491

xarray/core/dataset.py

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2797,6 +2797,37 @@ def _drop_vars(self, names):
27972797
coord_names = set(k for k in self._coord_names if k in variables)
27982798
return self._replace_vars_and_dims(variables, coord_names)
27992799

2800+
def drop_dims(self, drop_dims):
2801+
"""Drop dimensions and associated variables from this dataset.
2802+
2803+
Parameters
2804+
----------
2805+
drop_dims : str or list
2806+
Dimension or dimensions to drop.
2807+
2808+
Returns
2809+
-------
2810+
obj : Dataset
2811+
The dataset without the given dimensions (or any variables
2812+
containing those dimensions)
2813+
"""
2814+
if utils.is_scalar(drop_dims):
2815+
drop_dims = [drop_dims]
2816+
2817+
missing_dimensions = [d for d in drop_dims if d not in self.dims]
2818+
if missing_dimensions:
2819+
raise ValueError('Dataset does not contain the dimensions: %s'
2820+
% missing_dimensions)
2821+
2822+
drop_vars = set(k for k, v in self._variables.items()
2823+
for d in v.dims if d in drop_dims)
2824+
2825+
variables = OrderedDict((k, v) for k, v in self._variables.items()
2826+
if k not in drop_vars)
2827+
coord_names = set(k for k in self._coord_names if k in variables)
2828+
2829+
return self._replace_with_new_dims(variables, coord_names)
2830+
28002831
def transpose(self, *dims):
28012832
"""Return a new Dataset object with all array dimensions transposed.
28022833

xarray/tests/test_dataset.py

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1863,6 +1863,26 @@ def test_drop_index_labels(self):
18631863
ValueError, 'does not have coordinate labels'):
18641864
data.drop(1, 'y')
18651865

1866+
def test_drop_dims(self):
1867+
data = xr.Dataset({'A': (['x', 'y'], np.random.randn(2, 3)),
1868+
'B': ('x', np.random.randn(2)),
1869+
'x': ['a', 'b'], 'z': np.pi})
1870+
1871+
actual = data.drop_dims('x')
1872+
expected = data.drop(['A', 'B', 'x'])
1873+
assert_identical(expected, actual)
1874+
1875+
actual = data.drop_dims('y')
1876+
expected = data.drop('A')
1877+
assert_identical(expected, actual)
1878+
1879+
actual = data.drop_dims(['x', 'y'])
1880+
expected = data.drop(['A', 'B', 'x'])
1881+
assert_identical(expected, actual)
1882+
1883+
with pytest.raises((ValueError, KeyError)):
1884+
data.drop_dims('z') # not a dimension
1885+
18661886
def test_copy(self):
18671887
data = create_test_data()
18681888

0 commit comments

Comments
 (0)