Skip to content

modification methods on Coordinates #10318

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

Open
wants to merge 6 commits into
base: main
Choose a base branch
from
Open
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
101 changes: 99 additions & 2 deletions xarray/core/coordinates.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
from __future__ import annotations

from collections.abc import Hashable, Iterator, Mapping, Sequence
from collections.abc import Callable, Hashable, Iterable, Iterator, Mapping, Sequence
from contextlib import contextmanager
from typing import (
TYPE_CHECKING,
Expand All @@ -21,7 +21,7 @@
assert_no_index_corrupted,
create_default_index_implicit,
)
from xarray.core.types import DataVars, Self, T_DataArray, T_Xarray
from xarray.core.types import DataVars, ErrorOptions, Self, T_DataArray, T_Xarray
from xarray.core.utils import (
Frozen,
ReprObject,
Expand Down Expand Up @@ -561,6 +561,31 @@ def merge(self, other: Mapping[Any, Any] | None) -> Dataset:
variables=coords, coord_names=coord_names, indexes=indexes
)

def __or__(self, other: Self) -> Self:
"""Merge two sets of coordinates to create a new Coordinates object
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we need such docstrings for an operator?


The method implements the logic used for joining coordinates in the
result of a binary operation performed on xarray objects:

- If two index coordinates conflict (are not equal), an exception is
raised. You must align your data before passing it to this method.
- If an index coordinate and a non-index coordinate conflict, the non-
index coordinate is dropped.
- If two non-index coordinates conflict, both are dropped.

Parameters
----------
other : dict-like, optional
A :py:class:`Coordinates` object or any mapping that can be turned
into coordinates.

Returns
-------
merged : Coordinates
A new Coordinates object with merged coordinates.
"""
return self.merge(other).coords

def __setitem__(self, key: Hashable, value: Any) -> None:
self.update({key: value})

Expand Down Expand Up @@ -719,6 +744,78 @@ def copy(
),
)

def drop_vars(
self,
names: str | Iterable[Hashable] | Callable[[Self], str | Iterable[Hashable]],
*,
errors: ErrorOptions = "raise",
) -> Self:
"""Drop variables from this Coordinates object.

Note that indexes that depend on these variables will also be dropped.

Parameters
----------
names : hashable or iterable or callable
Name(s) of variables to drop. If a callable, this is object is passed as its
only argument and its result is used.
errors : {"raise", "ignore"}, default: "raise"
Error treatment.

- ``'raise'``: raises a :py:class:`ValueError` error if any of the variable
passed are not in the dataset
- ``'ignore'``: any given names that are in the dataset are dropped and no
error is raised.
"""
return self.to_dataset().drop_vars(names, errors=errors).coords

def rename_dims(
self,
dims_dict: Mapping[Any, Hashable] | None = None,
**dims: Hashable,
) -> Self:
"""Returns a new object with renamed dimensions only.

Parameters
----------
dims_dict : dict-like, optional
Dictionary whose keys are current dimension names and
whose values are the desired names. The desired names must
not be the name of an existing dimension or Variable in the Dataset.
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
not be the name of an existing dimension or Variable in the Dataset.
not be the name of an existing dimension or Variable in the Coordinates.

**dims : optional
Keyword form of ``dims_dict``.
One of dims_dict or dims must be provided.

Returns
-------
renamed : Coordinates
Coordinates object with renamed dimensions.
"""
return self.to_dataset().rename_dims(dims_dict, **dims).coords
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

self._data would have prevented duplicated copy of the wrapped Dataset, but it won't work with the DataArrayCoordinates subclass so I guess self.to_dataset() is fine.


def rename_vars(
self,
name_dict: Mapping[Any, Hashable] | None = None,
**names: Hashable,
) -> Self:
"""Returns a new object with renamed variables including coordinates
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
"""Returns a new object with renamed variables including coordinates
"""Returns a new object with renamed variables.


Parameters
----------
name_dict : dict-like, optional
Dictionary whose keys are current variable or coordinate names and
whose values are the desired names.
**names : optional
Keyword form of ``name_dict``.
One of name_dict or names must be provided.

Returns
-------
renamed : Dataset
Dataset with renamed variables including coordinates
Comment on lines +814 to +815
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
renamed : Dataset
Dataset with renamed variables including coordinates
renamed : Coordinates
Coordinates object with renamed variables

"""
return self.to_dataset().rename_vars(name_dict, **names)


class DatasetCoordinates(Coordinates):
"""Dictionary like container for Dataset coordinates (variables + indexes).
Expand Down
60 changes: 60 additions & 0 deletions xarray/tests/test_coordinates.py
Original file line number Diff line number Diff line change
Expand Up @@ -208,3 +208,63 @@ def test_dataset_from_coords_with_multidim_var_same_name(self):
coords = Coordinates(coords={"x": var}, indexes={})
ds = Dataset(coords=coords)
assert ds.coords["x"].dims == ("x", "y")

def test_drop_vars(self):
coords = Coordinates(
coords={
"x": Variable("x", range(3)),
"y": Variable("y", list("ab")),
"a": Variable(["x", "y"], np.arange(6).reshape(3, 2)),
},
indexes={},
)

actual = coords.drop_vars("x")
assert set(actual.variables) == {"a", "y"}

actual = coords.drop_vars(["x", "y"])
assert set(actual.variables) == {"a"}

def test_rename_dims(self):
coords = Coordinates(
coords={
"x": Variable("x", range(3)),
"y": Variable("y", list("ab")),
"a": Variable(["x", "y"], np.arange(6).reshape(3, 2)),
},
indexes={},
)

actual = coords.rename_dims({"x": "X"})
assert set(actual.dims) == {"X", "y"}
assert set(actual.variables) == {"a", "x", "y"}

actual = coords.rename_dims({"x": "u", "y": "v"})
assert set(actual.dims) == {"u", "v"}
assert set(actual.variables) == {"a", "x", "y"}

def test_rename_vars(self):
coords = Coordinates(
coords={
"x": Variable("x", range(3)),
"y": Variable("y", list("ab")),
"a": Variable(["x", "y"], np.arange(6).reshape(3, 2)),
},
indexes={},
)

actual = coords.rename_vars({"x": "X"})
assert set(actual.dims) == {"x", "y"}
assert set(actual.variables) == {"a", "X", "y"}

actual = coords.rename_vars({"x": "u", "y": "v"})
assert set(actual.dims) == {"x", "y"}
assert set(actual.variables) == {"a", "u", "v"}

def test_operator_merge(self):
coords1 = Coordinates({"x": ("x", [0, 1, 2])})
coords2 = Coordinates({"y": ("y", [3, 4, 5])})
expected = Dataset(coords={"x": [0, 1, 2], "y": [3, 4, 5]})

actual = coords1 | coords2
assert_identical(Dataset(coords=actual), expected)
Loading