Expose "Coordinates" as part of Xarray's public API #7368

benbovy · 2022-12-08T16:59:29Z

Closes Pass indexes directly to the DataArray and Dataset constructors #7214
Closes Pass indexes to the Dataset and DataArray constructors #6392
xref Opening dataset without loading any indexes? #6633
Tests added
User visible changes (including notable bug fixes) are documented in whats-new.rst
New functions/methods are listed in api.rst

This is a rework of #7214. It follows the suggestions made in #7214 (comment), #7214 (comment) and #7214 (comment):

No indexes argument is added to Dataset.__init__, and the indexes argument of DataArray.__init__ is kept private (i.e., valid only if fastpath=True)
When a Coordinates object is passed to a new Dataset or DataArray via the coords argument, both coordinate variables and indexes are copied/extracted and added to the new object
This PR also adds ~~an IndexedCoordinates subclass~~ Coordinates public constructors used to create Xarray coordinates and indexes from non-Xarray objects. For example, the Coordinates.from_pandas_multiindex() class method creates a new set of index and coordinates from an existing pd.MultiIndex.

EDIT: IndexCoordinates has been merged with Coordinates

EDIT2: it ended up as a pretty big refactor with the promotion of Coordinates has a 2nd-class Xarray container that supports alignment like Dataset and DataArray. It is still quite advanced API, useful for passing coordinate variables and indexes around. Internally, Coordinates objects are still "virtual" containers (i.e., proxies for coordinate variables and indexes stored in their corresponding DataArray or Dataset objects). For now, a "stand-alone" Coordinates object created from scratch wraps a Dataset with no data variables.

Some examples of usage:

import pandas as pd
import xarray as xr

midx = pd.MultiIndex.from_product([["a", "b"], [1, 2]], names=("one", "two"))

coords = xr.Coordinates.from_pandas_multiindex(midx, "x")
# Coordinates:
#   * x        (x) object MultiIndex
#   * one      (x) object 'a' 'a' 'b' 'b'
#   * two      (x) int64 1 2 1 2

ds = xr.Dataset(coords=coords)
# <xarray.Dataset>
# Dimensions:  (x: 4)
# Coordinates:
#   * x        (x) object MultiIndex
#   * one      (x) object 'a' 'a' 'b' 'b'
#   * two      (x) int64 1 2 1 2
# Data variables:
#     *empty*

ds_to_be_deprecated = xr.Dataset(coords={"x": midx})
ds_to_be_deprecated.identical(ds)
# True

da = xr.DataArray([1, 2, 3, 4], dims="x", coords=ds.coords)
# <xarray.DataArray (x: 4)>
# array([1, 2, 3, 4])
# Coordinates:
#   * x        (x) object MultiIndex
#   * one      (x) object 'a' 'a' 'b' 'b'
#   * two      (x) int64 1 2 1 2

TODO:

update assign_coords too so it has the same behavior if a Coordinates object is passed?
How to avoid building any default index? It seems silly to add or use the indexes argument just for that purpose? ~~We could address that later.~~ Solution: wrap the coordinates dict in a Coordinates objects, e.g., ds = xr.Dataset(coords=xr.Coordinates(coords_dict)).

@shoyer, @dcherian, anyone -- what do you think about the approach proposed here? I'd like to check that with you before going further with tests, docs, etc.

- easily create an empty Indexes collection - check consistency between indexes and variables

TODO: check indexes shapes / dims for DataArray

Since its constructor can now be used publicly. Copy input mappings and check the type of input indexes.

+ add `IndexedCoordinates.from_pandas_multiindex` helper.

Drop the `indexes` argument or keep it as private API. When a `Coordinates` object is passed as `coords` argument, extract both coordinate variables and indexes and add them to the new Dataset or DataArray.

benbovy · 2022-12-09T08:32:28Z

IndexedCoordinates and Indexes have a lot of overlap. At some point we might consider merging the two classes, like @shoyer suggests in #7214 (comment). The main difference is that one is a mapping of coordinates and the other is a mapping of indexes. IndexedCoordinates is mostly reusing Indexes and Dataset under the hood, it is only a facade.

Alternatively to an IndexedCoordinates subclass I was wondering if we could reuse the Coordinates base class? There's some benefit of providing a subclass:

besides specific constructors like .from_pandas_multiindex() it has a generic __init__ for advanced use cases. Not sure it is a good idea to add this constructor to the base class?
unlike Coordinates, IndexedCoordinates is immutable.

What if the Indexes class was a facade based on IndexedCoordinates instead of the other way around? It would probably make more sense but it would also be a bigger refactor. I've chosen the easy way :).

benbovy · 2022-12-09T09:13:24Z

I added IndexedCoordinates.merge_coords so that it is easier to combine different coordinates to pass to a new Dataset / DataArray, e.g.,

midx = pd.MultiIndex.from_product([["a", "b"], [1, 2]], names=("one", "two"))

coords = xr.IndexedCoordinates.from_pandas_multiindex(midx, "x")

coords = coords.merge_coords({"y": [0, 1, 2]})
# Coordinates:
#   * x        (x) object MultiIndex
#   * one      (x) object 'a' 'a' 'b' 'b'
#   * two      (x) int64 1 2 1 2
#   * y        (y) int64 0 1 2

ds = xr.Dataset(coords=coords)
# <xarray.Dataset>
# Dimensions:  (x: 4)
# Coordinates:
#   * x        (x) object MultiIndex
#   * one      (x) object 'a' 'a' 'b' 'b'
#   * two      (x) int64 1 2 1 2
#   * y        (y) int64 0 1 2
# Data variables:
#     *empty*

IndexedCoordinates.merge_coords is very much like Coordinates.merge except that it returns a new Coordinates object instead of a Dataset.

Or should we just use merge? It would require that:

Coordinates.merge accepts Mapping[Any, Any] for its other argument. Only changing the type hint is enough here since the implementation already accepts any input passed to Dataset.
When a Dataset is passed as coords argument to a new Dataset and DataArray, both variables and indexes should be extracted. It is already the case for Dataset but I think it only works for PandasIndex and PandasMultiIndex (default indexes & backwards compatibility).

shoyer · 2022-12-10T00:31:46Z

what do you think about the approach proposed here? I'd like to check that with you before going further with tests, docs, etc.

Generally this looks great to me!

How to avoid building any default index? It seems silly to add or use the indexes argument just for that purpose? We could address that later.

My suggestion would be:

coords passed as a dict: create default indexes
coords passed as IndexedCoordinates: do not create defaults

Alternatively to an IndexedCoordinates subclass I was wondering if we could reuse the Coordinates base class?

Yes, this makes more sense to me!

What if the Indexes class was a facade based on IndexedCoordinates instead of the other way around?

Yes, I also agree! This makes more sense.

shoyer · 2022-12-10T01:37:35Z

Long term, do you think it would make sense to merge together Indexes, Coordinates and IndexedCoordinates? They are sort of all containers for the same thing.

benbovy · 2022-12-10T16:59:44Z

Long term, do you think it would make sense to merge together Indexes, Coordinates and IndexedCoordinates? They are sort of all containers for the same thing.

Yes I think so.

I'm actually trying to merge IndexedCoordinates with Coordinates but I'm stuck: the latter is abstract and I don't really see how I could refactor it together with DatasetCoordinates and DataArrayCoordinates. Do you have any idea on how best to proceed?

Ideally, I'd see Coordinates be exposed in Xarray's main namespace with at least the two following constructors:

class Coordinates:

    def __init__(
        self,
        coords: Mapping[Any, Any] | None = None,
        indexes: Mapping[Any, Index] | None = None,
    ):
        # Similar to Dataset.__init__ but without the need
        # to merge coords and data vars...
        # Probably ok to allow more flexibility / less safety here?
        ...

    @classmethod
    from_pandas_multiindex(cls, index: pd.MultiIndex, dim: str):
        ...

shoyer · 2022-12-11T20:17:15Z

I'm actually trying to merge IndexedCoordinates with Coordinates but I'm stuck: the latter is abstract and I don't really see how I could refactor it together with DatasetCoordinates and DataArrayCoordinates

Coordinates is abstract because in the (current) Xarray data model, it doesn't actually store any data -- coordinates are stored in private attributes of the original Dataset (._variables and ._coord_names) or DataArray (._coords). So Coordinates needs to serve as a proxy for the data.

In the long term, I think we should refactor Dataset/DataArray to actually store data (coordinate variables, indexes and dimension sizes) on Coordinates, but that's a bigger refactor.

For now, it's worth noting that the current Coordinates class isn't actually exposed in Xarray's public API, just the DatasetCoordinates and DataArrayCoordinates classes (and not even their constructors). So the intermdiate step I would try is:

Rename the current Coordinates baseclass to AbstractCoordinates.
Rename your IndexedCoordinates class to Coordinates. Expose it in the public API. Make sure it can handle DatasetCoordinates and DataArrayCoordinates in the constructor.
Maybe: use some Python magic to make DatasetCorodinates/DataArrayCoordinates subclasses of the new Coordinates. Or maybe make them actual subclassses, overriding many of the methods (including the constructor).

benbovy · 2022-12-12T08:36:09Z

Thanks @shoyer, I've been thinking about similar short/long term plans although so far I haven't figured out how to implement your point 3. I'll give it another try.

headtr1ck

I totally forgot that you can use Protocols to solve this typing issue, nice!

Obviously need to fix the issues found by the tests :)

headtr1ck · 2023-07-17T19:32:12Z

xarray/core/types.py

+    import sys
+
+    if sys.version_info >= (3, 11):
+        from typing import Self


Nice! Finally this is merged.

This means we should be able to remove lots of the code that annotated self with a typevar (in another PR ofc).

headtr1ck · 2023-07-17T19:34:55Z

xarray/core/types.py

@@ -93,10 +94,63 @@
    DTypeLikeSave: Any = None


+class Alignable(Protocol):


benbovy · 2023-07-17T19:49:26Z

Thanks for your feedback @headtr1ck, your previous comments have been helpful. Yeah I gave it another try using Protocol, I've been struggling with it until trying with Self. This should be good now, hopefully. Another advantage is that the Alignable protocol class makes it pretty clear what is the required interface for an Xarray object to support alignment, i.e., it provides some self-documentation of the internals.

headtr1ck · 2023-07-17T19:51:40Z

xarray/core/coordinates.py

+
+    @classmethod
+    def _construct_direct(
+        cls: type[T_Coordinates],


Can't you use Self here as well?

Yes 👍 I'll go trough the coordinate classes and see if I can replace type variables by Self.

headtr1ck · 2023-07-17T19:51:51Z

xarray/core/coordinates.py

+
+    @classmethod
+    def from_pandas_multiindex(
+        cls: type[T_Coordinates], midx: pd.MultiIndex, dim: str


headtr1ck · 2023-07-17T19:52:27Z

xarray/core/coordinates.py

+        return cast(T_Coordinates, results.coords)
+
+    def _reindex_callback(
+        self: T_Coordinates,


And here (not sure if it works with cast though..., Can always use a type: ignore)

headtr1ck · 2023-07-17T19:56:54Z

xarray/core/coordinates.py

+            exclude_dims,
+            exclude_vars,
+        )
+        return cast(T_Coordinates, aligned.coords)


I have not checked this, but maybe mypy is right to complain?

If you start with a DataArrayCoordinates, do you not end up with a DatasetCoordinates?

If you start with a DataArrayCoordinates, do you not end up with a DatasetCoordinates?

This case shouldn't happen because Dataset and DataArray have their own _reindex_callback() called internally in Aligner. The callback here is called when aligning a Coordinates object, so we need to cast DatasetCoordinates into a Coordinates.

I should probably add a comment here to explain this.

dcherian · 2023-07-17T21:10:43Z

RTD failure is real:

home/docs/checkouts/readthedocs.org/user_builds/xray/checkouts/7368/doc/api-hidden.rst:9: WARNING: autosummary: failed to import Coordinates.merge_coords.
Possible hints:
* ModuleNotFoundError: No module named 'xarray.Coordinates'
* PycodeError: no source found for module 'builtins'
* AttributeError: type object 'Coordinates' has no attribute 'merge_coords'
* ModuleNotFoundError: No module named 'Coordinates'
* KeyError: 'Coordinates'
/home/docs/checkouts/readthedocs.org/user_builds/xray/checkouts/7368/doc/api-hidden.rst:9: WARNING: autosummary: failed to import core.coordinates.DatasetCoordinates.merge_coords.
Possible hints:
* ModuleNotFoundError: No module named 'xarray.core.coordinates.DatasetCoordinates'; 'xarray.core.coordinates' is not a package
* PycodeError: no source found for module 'builtins'
* AttributeError: type object 'DatasetCoordinates' has no attribute 'merge_coords'
* KeyError: 'core'
* ModuleNotFoundError: No module named 'core'
/home/docs/checkouts/readthedocs.org/user_builds/xray/checkouts/7368/doc/api-hidden.rst:9: WARNING: autosummary: failed to import core.coordinates.DataArrayCoordinates.merge_coords.
Possible hints:
* PycodeError: no source found for module 'builtins'
* KeyError: 'core'
* AttributeError: type object 'DataArrayCoordinates' has no attribute 'merge_coords'
* ModuleNotFoundError: No module named 'xarray.core.coordinates.DataArrayCoordinates'; 'xarray.core.coordinates' is not a package
* ModuleNotFoundError: No module named 'core'

dcherian · 2023-07-17T21:17:22Z

@benbovy does this actually close #6633? It seems like we need a way to pass a Coordinates object to open_dataset and down to the backend code. We can totally address this in a followup PR.

benbovy · 2023-07-18T09:02:42Z

@dcherian yes this PR would enable #6633 by passing a Coordinates object to the Dataset constructor in the backend code. It could be done in another PR indeed.

This doesn't create a default index for "x" (since no index is passed explicitly to the Coordinates object):

import xarray as xr

ds = xr.Dataset(
    data_vars={"foo": ("x", [0.0, 1.0])},
    coords=xr.Coordinates({"x": ("x", [1, 2])}),
)

Note that if "x" is passed as a data variable it will still be promoted as a coordinate with a default index, even when a Coordinates object is passed explicitly:

xr.Dataset(
    data_vars={"foo": ("x", [1.0, 2.0]), "x": [1, 2]},
    coords=xr.Coordinates(),
)

This might be slightly confusing but I don't know if it is really worth handling this special case (i.e., dimension coordinate variables in data_vars + empty Coordinates). I think we should keep promoting data variables as coordinates since data_vars is often used as a unique positional argument of Dataset.__init__().

benbovy · 2023-07-18T11:35:58Z

I updated the docstrings of Dataset.__init__ and DataArray.__init__ to clarify that we can pass a Coordinates object to bypass the creation of default indexes for dimension coordinates.

I took this opportunity to improve the glossary in the documentation and make the concepts of (non)-dimension / (non)-indexed coordinate more in line with the current state of Xarray since the index refactor. I'm sure there are still other inconsistent parts in the documentation, though. I'll review and address that in a follow-up PR.

Pyright displays an info message "Self is not valid in this context" but most important this should avoid runtime errors with python < 3.11.

TomNicholas · 2023-07-18T13:56:04Z

I just want to say this is another amazing piece of work @benbovy 👏 👏

dcherian · 2023-07-21T15:11:12Z

@benbovy shall we merge?

benbovy · 2023-07-21T19:30:33Z

@dcherian yes!

…lazy-array * upstream/main: (153 commits) Add HDF5 Section to read/write docs page (pydata#8012) [pre-commit.ci] pre-commit autoupdate (pydata#8014) Update interpolate_na in dataset.py (pydata#7974) improved docstring of to_netcdf (issue pydata#7127) (pydata#7947) Expose "Coordinates" as part of Xarray's public API (pydata#7368) Core team member guide (pydata#7999) join together duplicate entries in the text `repr` (pydata#7225) Update copyright year in README (pydata#8007) Allow opening datasets with nD dimenson coordinate variables. (pydata#7989) Move whats-new entry [pre-commit.ci] pre-commit autoupdate (pydata#7997) Add documentation on custom indexes (pydata#6975) Use variable name in all exceptions raised in `as_variable` (pydata#7995) Bump pypa/gh-action-pypi-publish from 1.8.7 to 1.8.8 (pydata#7994) New whatsnew section Remove future release notes before this release Update whats-new.rst for new release (pydata#7993) Remove hue_style from plot1d docstring (pydata#7925) Add new what's new section (pydata#7986) Release summary for v2023.07.0 (pydata#7979) ...

benbovy added 16 commits October 25, 2022 16:11

add indexes argument to Dataset.__init__

41f4fd8

make indexes arg public for DataArray.__init__

4baa8af

Indexes constructor updates

dbc058a

- easily create an empty Indexes collection - check consistency between indexes and variables

use the generic Mapping[Any, Index] for indexes

16a9983

add wrap_pandas_multiindex function

3c076d5

do not create default indexes when not desired

70e7a5d

fix Dataset dimensions

00e1766

TODO: check indexes shapes / dims for DataArray

copy the coordinate variables of passed indexes

3bf92cd

DataArray: check dimensions/shape of index coords

c9b6363

docstrings tweaks

82dc5cc

more Indexes safety

a58c9d0

Since its constructor can now be used publicly. Copy input mappings and check the type of input indexes.

ensure input indexes are Xarray indexes

9beeea7

add .assign_indexes() method

c6e94b4

Merge branch 'main' into indexes-arg-constructors

ddd505e

add IndexedCoordinates subclass

f97adb5

+ add `IndexedCoordinates.from_pandas_multiindex` helper.

rollback/update Dataset and DataArray constructors

45709ef

Drop the `indexes` argument or keep it as private API. When a `Coordinates` object is passed as `coords` argument, extract both coordinate variables and indexes and add them to the new Dataset or DataArray.

github-actions bot added the topic-indexing label Dec 8, 2022

update docstrings

4c559f1

fix Dataset creation internal error

1192948

add IndexedCoordinates.merge_coords

a877a74

benbovy added 2 commits December 12, 2022 12:26

drop IndexedCoordinates and reuse Coordinates

9d6d2ae

update api docs

3ee26ef

headtr1ck approved these changes Jul 17, 2023

View reviewed changes

fix groupby test

139b13a

headtr1ck reviewed Jul 17, 2023

View reviewed changes

benbovy added 4 commits July 18, 2023 11:18

doc: remove merge_coords in api reference

7ed6279

doc: improve docstrings and glossary

3d94357

use Self type annotation in Coordinate class

4a6e915

better comment

31f66b4

fix Self undefined error with python < 3.11

4cb70d0

Pyright displays an info message "Self is not valid in this context" but most important this should avoid runtime errors with python < 3.11.

This was referenced Jul 18, 2023

Pass indexes directly to the DataArray and Dataset constructors #7214

Closed

Explicit indexes: next steps #6293

Open

Merge branch 'main' into indexes-arg-constructors-2

4ef5f17

dcherian merged commit 4441f99 into pydata:main Jul 21, 2023
27 checks passed

bzah mentioned this pull request Jul 26, 2023

⚠️ Nightly upstream-dev CI failed ⚠️ Ouranosinc/xclim#1417

Closed

abarciauskas-bgse mentioned this pull request Jul 28, 2023

Investigate if it is possible to avoid reading all coordinate chunks when opening a dataset with xarray developmentseed/tile-benchmarking#34

Closed

aulemahal mentioned this pull request Aug 1, 2023

Update assign_coords with a MultiIndex to match new Coordinates API #8039

Closed

headtr1ck mentioned this pull request Aug 3, 2023

Docs page on internal design #7991

Merged

4 tasks

Zeitsperre mentioned this pull request Aug 21, 2023

Address typing hints for xarray 2023.08 compatibility, Update ReadTheDocs roocs/clisops#295

Merged

5 tasks

benbovy mentioned this pull request Aug 23, 2023

Confusing terminologies and some errors in the official documentation #6866

Closed

4 tasks

benbovy deleted the indexes-arg-constructors-2 branch August 30, 2023 09:11

jklenzing mentioned this pull request Sep 8, 2023

ENH: allow custom concat_data methods pysat/pysat#1131

Merged

10 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Expose "Coordinates" as part of Xarray's public API #7368

Expose "Coordinates" as part of Xarray's public API #7368

benbovy commented Dec 8, 2022 •

edited by dcherian

Loading

benbovy commented Dec 9, 2022 •

edited

Loading

benbovy commented Dec 9, 2022 •

edited

Loading

shoyer commented Dec 10, 2022

shoyer commented Dec 10, 2022

benbovy commented Dec 10, 2022

shoyer commented Dec 11, 2022

benbovy commented Dec 12, 2022

headtr1ck left a comment •

edited

Loading

headtr1ck Jul 17, 2023

headtr1ck Jul 17, 2023

benbovy commented Jul 17, 2023

headtr1ck Jul 17, 2023

benbovy Jul 17, 2023

headtr1ck Jul 17, 2023

headtr1ck Jul 17, 2023

headtr1ck Jul 17, 2023

benbovy Jul 17, 2023

dcherian commented Jul 17, 2023

dcherian commented Jul 17, 2023

benbovy commented Jul 18, 2023 •

edited

Loading

benbovy commented Jul 18, 2023

TomNicholas commented Jul 18, 2023

dcherian commented Jul 21, 2023

benbovy commented Jul 21, 2023

		@@ -93,10 +94,63 @@
		DTypeLikeSave: Any = None


		class Alignable(Protocol):

Expose "Coordinates" as part of Xarray's public API #7368

Expose "Coordinates" as part of Xarray's public API #7368

Conversation

benbovy commented Dec 8, 2022 • edited by dcherian Loading

benbovy commented Dec 9, 2022 • edited Loading

benbovy commented Dec 9, 2022 • edited Loading

shoyer commented Dec 10, 2022

shoyer commented Dec 10, 2022

benbovy commented Dec 10, 2022

shoyer commented Dec 11, 2022

benbovy commented Dec 12, 2022

headtr1ck left a comment • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

benbovy commented Jul 17, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

dcherian commented Jul 17, 2023

dcherian commented Jul 17, 2023

benbovy commented Jul 18, 2023 • edited Loading

benbovy commented Jul 18, 2023

TomNicholas commented Jul 18, 2023

dcherian commented Jul 21, 2023

benbovy commented Jul 21, 2023

benbovy commented Dec 8, 2022 •

edited by dcherian

Loading

benbovy commented Dec 9, 2022 •

edited

Loading

benbovy commented Dec 9, 2022 •

edited

Loading

headtr1ck left a comment •

edited

Loading

benbovy commented Jul 18, 2023 •

edited

Loading