Skip to content

Commit

Permalink
Closes #4658
Browse files Browse the repository at this point in the history
- Use get_index(dim) in drop_sel
- Add drop_isel
  • Loading branch information
mesejo committed Jan 17, 2021
1 parent a2b1712 commit 236872e
Show file tree
Hide file tree
Showing 7 changed files with 91 additions and 8 deletions.
2 changes: 1 addition & 1 deletion doc/groupby.rst
Original file line number Diff line number Diff line change
Expand Up @@ -226,4 +226,4 @@ applying your function, and then unstacking the result:
.. ipython:: python
stacked = da.stack(gridcell=["ny", "nx"])
stacked.groupby("gridcell").sum(...).unstack("gridcell")
stacked.groupby("gridcell").sum(...).unstack("gridcell")
2 changes: 1 addition & 1 deletion doc/internals.rst
Original file line number Diff line number Diff line change
Expand Up @@ -230,4 +230,4 @@ re-open it directly with Zarr:
zgroup = zarr.open("rasm.zarr")
print(zgroup.tree())
dict(zgroup["Tair"].attrs)
dict(zgroup["Tair"].attrs)
2 changes: 1 addition & 1 deletion doc/interpolation.rst
Original file line number Diff line number Diff line change
Expand Up @@ -317,4 +317,4 @@ The remapping can be done as follows
dsi = ds.interp(lon=lon, lat=lat)
dsi.air.plot(ax=axes[1])
@savefig interpolation_sample4.png width=8in
axes[1].set_title("Remapped data")
axes[1].set_title("Remapped data")
2 changes: 1 addition & 1 deletion doc/plotting.rst
Original file line number Diff line number Diff line change
Expand Up @@ -955,4 +955,4 @@ One can also make line plots with multidimensional coordinates. In this case, ``
f, ax = plt.subplots(2, 1)
da.plot.line(x="lon", hue="y", ax=ax[0])
@savefig plotting_example_2d_hue_xy.png
da.plot.line(x="lon", hue="x", ax=ax[1])
da.plot.line(x="lon", hue="x", ax=ax[1])
2 changes: 1 addition & 1 deletion doc/reshaping.rst
Original file line number Diff line number Diff line change
Expand Up @@ -288,4 +288,4 @@ As a shortcut, you can refer to existing coordinates by name:
ds.sortby("x")
ds.sortby(["y", "x"])
ds.sortby(["y", "x"], ascending=False)
ds.sortby(["y", "x"], ascending=False)
57 changes: 56 additions & 1 deletion xarray/core/dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -4053,13 +4053,68 @@ def drop_sel(self, labels=None, *, errors="raise", **labels_kwargs):
labels_for_dim = [labels_for_dim]
labels_for_dim = np.asarray(labels_for_dim)
try:
index = self.indexes[dim]
index = self.get_index(dim)
except KeyError:
raise ValueError("dimension %r does not have coordinate labels" % dim)
new_index = index.drop(labels_for_dim, errors=errors)
ds = ds.loc[{dim: new_index}]
return ds

def drop_isel(self, indexers=None, **indexers_kwargs):
"""Drop index positions from this dataset.
Parameters
----------
indexers : mapping of hashable to Any
Index locations to drop
**indexers_kwargs : {dim: position, ...}, optional
The keyword arguments form of ``dim`` and ``positions``
Returns
-------
dropped : Dataset
Raises
------
IndexError
Examples
--------
>>> data = np.random.randn(2, 3)
>>> labels = ["a", "b", "c"]
>>> ds = xr.Dataset({"A": (["x", "y"], data), "y": labels})
>>> ds.drop_isel(y=[0, 2])
<xarray.Dataset>
Dimensions: (x: 2, y: 1)
Coordinates:
* y (y) <U1 'b'
Dimensions without coordinates: x
Data variables:
A (x, y) float64 0.4002 1.868
>>> ds.drop_sel(y=1)
<xarray.Dataset>
Dimensions: (x: 2, y: 2)
Coordinates:
* y (y) <U1 'a' 'c'
Dimensions without coordinates: x
Data variables:
A (x, y) float64 1.764 0.9787 2.241 -0.9773
"""

indexers = either_dict_or_kwargs(indexers, indexers_kwargs, "drop")

ds = self
for dim, pos_for_dim in indexers.items():
# Don't cast to set, as it would harm performance when labels
# is a large numpy array
if utils.is_scalar(pos_for_dim):
pos_for_dim = [pos_for_dim]
pos_for_dim = np.asarray(pos_for_dim).tolist()
index = self.get_index(dim)
new_index = index.delete(pos_for_dim)
ds = ds.loc[{dim: new_index}]
return ds

def drop_dims(
self, drop_dims: Union[Hashable, Iterable[Hashable]], *, errors: str = "raise"
) -> "Dataset":
Expand Down
32 changes: 30 additions & 2 deletions xarray/tests/test_dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -2371,8 +2371,12 @@ def test_drop_index_labels(self):
data.drop(DataArray(["a", "b", "c"]), dim="x", errors="ignore")
assert_identical(expected, actual)

with raises_regex(ValueError, "does not have coordinate labels"):
data.drop_sel(y=1)
actual = data.drop_sel(y=[1])
expected = data.isel(y=[0, 2])
assert_identical(expected, actual)

with raises_regex(KeyError, "not found in axis"):
data.drop_sel(x=0)

def test_drop_labels_by_keyword(self):
data = Dataset(
Expand Down Expand Up @@ -2410,6 +2414,30 @@ def test_drop_labels_by_keyword(self):
with pytest.raises(ValueError):
data.drop(dim="x", x="a")

def test_drop_labels_by_position(self):
data = Dataset(
{"A": (["x", "y"], np.random.randn(2, 6)), "x": ["a", "b"], "y": range(6)}
)
# Basic functionality.
assert len(data.coords["x"]) == 2

ds2 = data.drop_isel(x=0)
ds3 = data.drop_isel(x=[0])
ds4 = data.drop_isel(x=[0, 1])
ds5 = data.drop_isel(x=[0, 1], y=range(0, 6, 2))

assert_array_equal(ds2.coords["x"], ["b"])
assert_array_equal(ds3.coords["x"], ["b"])
assert ds4.coords["x"].size == 0
assert ds5.coords["x"].size == 0
assert_array_equal(ds5.coords["y"], [1, 3, 5])

with pytest.raises(KeyError):
data.drop_isel(z=1)

with pytest.raises(IndexError):
data.drop_isel(x=5)

def test_drop_dims(self):
data = xr.Dataset(
{
Expand Down

0 comments on commit 236872e

Please sign in to comment.