Skip to content

Commit 2956067

Browse files
authored
Add drop_isel (#4819)
* Closes #4658 - Use get_index(dim) in drop_sel - Add drop_isel * address issues in PR * extract dict creation out of the loop
1 parent ba42c08 commit 2956067

File tree

6 files changed

+131
-3
lines changed

6 files changed

+131
-3
lines changed

doc/api.rst

+2
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
@@ -307,6 +308,7 @@ Indexing
307308
DataArray.isel
308309
DataArray.sel
309310
DataArray.drop_sel
311+
DataArray.drop_isel
310312
DataArray.head
311313
DataArray.tail
312314
DataArray.thin

doc/whats-new.rst

+1
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` and :py:meth:`DataArray.drop_isel` (:issue:`4658`, :pull:`4819`). By `Daniel Mesejo <https://github.com/mesejo>`_.
8384

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

xarray/core/dataarray.py

+22
Original file line numberDiff line numberDiff line change
@@ -2248,6 +2248,28 @@ def drop_sel(
22482248
ds = self._to_temp_dataset().drop_sel(labels, errors=errors)
22492249
return self._from_temp_dataset(ds)
22502250

2251+
def drop_isel(self, indexers=None, **indexers_kwargs):
2252+
"""Drop index positions from this DataArray.
2253+
2254+
Parameters
2255+
----------
2256+
indexers : mapping of hashable to Any
2257+
Index locations to drop
2258+
**indexers_kwargs : {dim: position, ...}, optional
2259+
The keyword arguments form of ``dim`` and ``positions``
2260+
2261+
Returns
2262+
-------
2263+
dropped : DataArray
2264+
2265+
Raises
2266+
------
2267+
IndexError
2268+
"""
2269+
dataset = self._to_temp_dataset()
2270+
dataset = dataset.drop_isel(indexers=indexers, **indexers_kwargs)
2271+
return self._from_temp_dataset(dataset)
2272+
22512273
def dropna(
22522274
self, dim: Hashable, how: str = "any", thresh: int = None
22532275
) -> "DataArray":

xarray/core/dataset.py

+66-1
Original file line numberDiff line numberDiff line change
@@ -4054,13 +4054,78 @@ def drop_sel(self, labels=None, *, errors="raise", **labels_kwargs):
40544054
labels_for_dim = [labels_for_dim]
40554055
labels_for_dim = np.asarray(labels_for_dim)
40564056
try:
4057-
index = self.indexes[dim]
4057+
index = self.get_index(dim)
40584058
except KeyError:
40594059
raise ValueError("dimension %r does not have coordinate labels" % dim)
40604060
new_index = index.drop(labels_for_dim, errors=errors)
40614061
ds = ds.loc[{dim: new_index}]
40624062
return ds
40634063

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

xarray/tests/test_dataarray.py

+6
Original file line numberDiff line numberDiff line change
@@ -2327,6 +2327,12 @@ def test_drop_index_labels(self):
23272327
with pytest.warns(DeprecationWarning):
23282328
arr.drop([0, 1, 3], dim="y", errors="ignore")
23292329

2330+
def test_drop_index_positions(self):
2331+
arr = DataArray(np.random.randn(2, 3), dims=["x", "y"])
2332+
actual = arr.drop_sel(y=[0, 1])
2333+
expected = arr[:, 2:]
2334+
assert_identical(actual, expected)
2335+
23302336
def test_dropna(self):
23312337
x = np.random.randn(4, 4)
23322338
x[::2, 0] = np.nan

xarray/tests/test_dataset.py

+34-2
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,34 @@ 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+
actual = data.drop_isel(x=0)
2425+
expected = data.drop_sel(x="a")
2426+
assert_identical(expected, actual)
2427+
2428+
actual = data.drop_isel(x=[0])
2429+
expected = data.drop_sel(x=["a"])
2430+
assert_identical(expected, actual)
2431+
2432+
actual = data.drop_isel(x=[0, 1])
2433+
expected = data.drop_sel(x=["a", "b"])
2434+
assert_identical(expected, actual)
2435+
assert actual.coords["x"].size == 0
2436+
2437+
actual = data.drop_isel(x=[0, 1], y=range(0, 6, 2))
2438+
expected = data.drop_sel(x=["a", "b"], y=range(0, 6, 2))
2439+
assert_identical(expected, actual)
2440+
assert actual.coords["x"].size == 0
2441+
2442+
with pytest.raises(KeyError):
2443+
data.drop_isel(z=1)
2444+
24132445
def test_drop_dims(self):
24142446
data = xr.Dataset(
24152447
{

0 commit comments

Comments
 (0)