Skip to content

Commit 1617485

Browse files
committed
Closes #4658
- Use get_index(dim) in drop_sel - Add drop_isel
1 parent a2b1712 commit 1617485

File tree

4 files changed

+85
-3
lines changed

4 files changed

+85
-3
lines changed

doc/api.rst

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -126,6 +126,7 @@ Indexing
126126
Dataset.isel
127127
Dataset.sel
128128
Dataset.drop_sel
129+
Dataset.drop_isel
129130
Dataset.head
130131
Dataset.tail
131132
Dataset.thin

doc/whats-new.rst

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,7 @@ Bug fixes
8080
- Expand user directory paths (e.g. ``~/``) in :py:func:`open_mfdataset` and
8181
:py:meth:`Dataset.to_zarr` (:issue:`4783`, :pull:`4795`).
8282
By `Julien Seguinot <https://github.com/juseg>`_.
83+
- Add :py:meth:`Dataset.drop_isel` (:issue:`4658`, :pull:`4819`). By `Daniel Mesejo <https://github.com/mesejo>`_.
8384

8485
Documentation
8586
~~~~~~~~~~~~~

xarray/core/dataset.py

Lines changed: 56 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4053,13 +4053,68 @@ def drop_sel(self, labels=None, *, errors="raise", **labels_kwargs):
40534053
labels_for_dim = [labels_for_dim]
40544054
labels_for_dim = np.asarray(labels_for_dim)
40554055
try:
4056-
index = self.indexes[dim]
4056+
index = self.get_index(dim)
40574057
except KeyError:
40584058
raise ValueError("dimension %r does not have coordinate labels" % dim)
40594059
new_index = index.drop(labels_for_dim, errors=errors)
40604060
ds = ds.loc[{dim: new_index}]
40614061
return ds
40624062

4063+
def drop_isel(self, indexers=None, **indexers_kwargs):
4064+
"""Drop index positions from this dataset.
4065+
4066+
Parameters
4067+
----------
4068+
indexers : mapping of hashable to Any
4069+
Index locations to drop
4070+
**indexers_kwargs : {dim: position, ...}, optional
4071+
The keyword arguments form of ``dim`` and ``positions``
4072+
4073+
Returns
4074+
-------
4075+
dropped : Dataset
4076+
4077+
Raises
4078+
------
4079+
IndexError
4080+
4081+
Examples
4082+
--------
4083+
>>> data = np.random.randn(2, 3)
4084+
>>> labels = ["a", "b", "c"]
4085+
>>> ds = xr.Dataset({"A": (["x", "y"], data), "y": labels})
4086+
>>> ds.drop_isel(y=[0, 2])
4087+
<xarray.Dataset>
4088+
Dimensions: (x: 2, y: 1)
4089+
Coordinates:
4090+
* y (y) <U1 'b'
4091+
Dimensions without coordinates: x
4092+
Data variables:
4093+
A (x, y) float64 0.4002 1.868
4094+
>>> ds.drop_isel(y=1)
4095+
<xarray.Dataset>
4096+
Dimensions: (x: 2, y: 2)
4097+
Coordinates:
4098+
* y (y) <U1 'a' 'c'
4099+
Dimensions without coordinates: x
4100+
Data variables:
4101+
A (x, y) float64 1.764 0.9787 2.241 -0.9773
4102+
"""
4103+
4104+
indexers = either_dict_or_kwargs(indexers, indexers_kwargs, "drop")
4105+
4106+
ds = self
4107+
for dim, pos_for_dim in indexers.items():
4108+
# Don't cast to set, as it would harm performance when labels
4109+
# is a large numpy array
4110+
if utils.is_scalar(pos_for_dim):
4111+
pos_for_dim = [pos_for_dim]
4112+
pos_for_dim = np.asarray(pos_for_dim).tolist()
4113+
index = self.get_index(dim)
4114+
new_index = index.delete(pos_for_dim)
4115+
ds = ds.loc[{dim: new_index}]
4116+
return ds
4117+
40634118
def drop_dims(
40644119
self, drop_dims: Union[Hashable, Iterable[Hashable]], *, errors: str = "raise"
40654120
) -> "Dataset":

xarray/tests/test_dataset.py

Lines changed: 27 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2371,8 +2371,12 @@ def test_drop_index_labels(self):
23712371
data.drop(DataArray(["a", "b", "c"]), dim="x", errors="ignore")
23722372
assert_identical(expected, actual)
23732373

2374-
with raises_regex(ValueError, "does not have coordinate labels"):
2375-
data.drop_sel(y=1)
2374+
actual = data.drop_sel(y=[1])
2375+
expected = data.isel(y=[0, 2])
2376+
assert_identical(expected, actual)
2377+
2378+
with raises_regex(KeyError, "not found in axis"):
2379+
data.drop_sel(x=0)
23762380

23772381
def test_drop_labels_by_keyword(self):
23782382
data = Dataset(
@@ -2410,6 +2414,27 @@ def test_drop_labels_by_keyword(self):
24102414
with pytest.raises(ValueError):
24112415
data.drop(dim="x", x="a")
24122416

2417+
def test_drop_labels_by_position(self):
2418+
data = Dataset(
2419+
{"A": (["x", "y"], np.random.randn(2, 6)), "x": ["a", "b"], "y": range(6)}
2420+
)
2421+
# Basic functionality.
2422+
assert len(data.coords["x"]) == 2
2423+
2424+
ds2 = data.drop_isel(x=0)
2425+
ds3 = data.drop_isel(x=[0])
2426+
ds4 = data.drop_isel(x=[0, 1])
2427+
ds5 = data.drop_isel(x=[0, 1], y=range(0, 6, 2))
2428+
2429+
assert_array_equal(ds2.coords["x"], ["b"])
2430+
assert_array_equal(ds3.coords["x"], ["b"])
2431+
assert ds4.coords["x"].size == 0
2432+
assert ds5.coords["x"].size == 0
2433+
assert_array_equal(ds5.coords["y"], [1, 3, 5])
2434+
2435+
with pytest.raises(KeyError):
2436+
data.drop_isel(z=1)
2437+
24132438
def test_drop_dims(self):
24142439
data = xr.Dataset(
24152440
{

0 commit comments

Comments
 (0)