Fix `DataTree.coords.setitem` by adding `DataTreeCoordinates` class #9451

TomNicholas · 2024-09-08T01:31:54Z

Doesn't work yet - pushing to get feedback on the approach.

Closes DataTree.coords.__setitem__ is broken #9204
Tests added
User visible changes (including notable bug fixes) are documented in whats-new.rst
New functions/methods are listed in api.rst

for more information, see https://pre-commit.ci

TomNicholas · 2024-09-08T01:34:44Z

xarray/core/coordinates.py

+    # TODO should inherited coordinates be here? It would be very hard to allow updating them...
+    # But actually maybe the ChainMap approach would make this work okay??


I can't see an easy way to make assignment of coordinates to parent nodes work. I think it would be better just to have dt.coords constructed from inherited coords, but not allow setting coordinates further up the tree. I believe that's consistent with the behaviour we currently have for dt.__setitem__ anyway.

shoyer

Looking great! Here are some ideas

xarray/core/coordinates.py

shoyer · 2024-09-08T01:37:52Z

xarray/core/coordinates.py

+        # TODO is there a potential bug here? What happens if a dim is only present on data variables?
+        return self._data.dims


I think this should be fine. Dimensions on DatasetCoordinates also include dimensions that are only present on data variables.

Really? This (on main) looks wrong to me:

In [1]: import xarray as xr In [2]: ds = xr.Dataset({'a': ('x', [0, 1])}) In [3]: ds Out[3]: <xarray.Dataset> Size: 16B Dimensions: (x: 2) Dimensions without coordinates: x Data variables: a (x) int64 16B 0 1 In [4]: ds.coords Out[4]: Coordinates: *empty* In [5]: ds.coords.dims Out[5]: FrozenMappingWarningOnValuesAccess({'x': 2})

I mean the fact no-one has raised this before means it probably isn't of much consequence, but it does seem incorrect / misleading.

Raised #9466 to track this

xarray/core/coordinates.py

shoyer · 2024-09-08T01:46:01Z

xarray/core/coordinates.py

+        self, coords: dict[Hashable, Variable], indexes: Mapping[Any, Index]
+    ) -> None:
+
+        # TODO I don't know how to update coordinates that live in parent nodes


I think it is OK for this to be an error. The user can replace those coordinates on the parent nodes.

xarray/core/coordinates.py

Co-authored-by: Stephan Hoyer <shoyer@google.com>

TomNicholas · 2024-09-09T17:28:01Z

xarray/core/coordinates.py

+    def _drop_indexed_coords(self, coords_to_drop: set[Hashable]) -> None:
+        assert self._data.xindexes is not None
+        new_coords = drop_indexed_coords(coords_to_drop, self)
+        for name in self._data._node_coord_variables - new_coords._names:
+            del self._data._node_coord_variables[name]
+        self._data._node_indexes = dict(new_coords.xindexes)


I'm pretty sure that this ._drop_indexed_coords method, and its equivalent in DatasetCoordinates, is never being called... I can't find any references to it, and if I raise in both methods then the whole test suite still passes locally.

@dcherian @benbovy you both touched this method last, do you know why it still exists?

Removed in 9dc845a

…las/xarray into datatree_coords_setitem

TomNicholas · 2024-09-10T02:00:55Z

xarray/core/coordinates.py

+    # TODO should inherited coordinates be here? It would be very hard to allow updating them...
+    # But actually maybe the ChainMap approach would make this work okay??
+
+    _data: DataTree


Inheriting from Coordinates and using DataTree instead of DataWithCoords here both expose a few legitimate typing issues. Particularly:

DataTree doesn't directly inherit from DataWithCoords, though maybe it could inherit from a future version of it (see Refactor Dataset internals to store data variables and coordinate variables as separate dicts #9203 (comment))

Hashable vs str annoyances (DataTree should support Hashable names. #8836)

Mostly we have got away with these discrepancies so far but inheritance here is probably going to mean I need to either ignore or cast all over the place?

cc @headtr1ck

actually this wasn't very difficult to work around, it mostly just required adding a type: ignore[assignment] here :)

TomNicholas · 2024-09-11T01:26:54Z

I think this is basically ready now @shoyer

xarray/core/coordinates.py

xarray/core/datatree.py

xarray/tests/test_datatree.py

shoyer

Looking great, thanks Tom!

Co-authored-by: Stephan Hoyer <shoyer@google.com>

* main: (26 commits) Forbid modifying names of DataTree objects with parents (pydata#9494) DAS-2155 - Merge datatree documentation into main docs. (pydata#9033) Make illegal path-like variable names when constructing a DataTree from a Dataset (pydata#9378) Ensure TreeNode doesn't copy in-place (pydata#9482) `open_groups` for zarr backends (pydata#9469) Update pyproject.toml (pydata#9484) New whatsnew section (pydata#9483) Release notes for v2024.09.0 (pydata#9480) Fix `DataTree.coords.__setitem__` by adding `DataTreeCoordinates` class (pydata#9451) Rename DataTree's "ds" and "data" to "dataset" (pydata#9476) Update DataTree repr to indicate inheritance (pydata#9470) Bump pypa/gh-action-pypi-publish in the actions group (pydata#9460) Repo checker (pydata#9450) Add days_in_year and decimal_year to dt accessor (pydata#9105) remove parent argument from DataTree.__init__ (pydata#9465) Fix inheritance in DataTree.copy() (pydata#9457) Implement `DataTree.__delitem__` (pydata#9453) Add ASV for datatree.from_dict (pydata#9459) Make the first argument in DataTree.from_dict positional only (pydata#9446) Fix typos across the code, doc and comments (pydata#9443) ...

* main: (29 commits) Release notes for v2024.09.0 (pydata#9480) Fix `DataTree.coords.__setitem__` by adding `DataTreeCoordinates` class (pydata#9451) Rename DataTree's "ds" and "data" to "dataset" (pydata#9476) Update DataTree repr to indicate inheritance (pydata#9470) Bump pypa/gh-action-pypi-publish in the actions group (pydata#9460) Repo checker (pydata#9450) Add days_in_year and decimal_year to dt accessor (pydata#9105) remove parent argument from DataTree.__init__ (pydata#9465) Fix inheritance in DataTree.copy() (pydata#9457) Implement `DataTree.__delitem__` (pydata#9453) Add ASV for datatree.from_dict (pydata#9459) Make the first argument in DataTree.from_dict positional only (pydata#9446) Fix typos across the code, doc and comments (pydata#9443) DataTree should not be "Generic" (pydata#9445) Disallow passing a DataArray as data into the DataTree constructor (pydata#9444) Support additional dtypes in `resample` (pydata#9413) Shallow copy parent and children in DataTree constructor (pydata#9297) Bump minimum versions for dependencies (pydata#9434) Always include at least one category in random test data (pydata#9436) Avoid deep-copy when constructing groupby codes (pydata#9429) ...

…ss (pydata#9451) * add a DataTreeCoordinates class * passing read-only properties tests * tests for modifying in-place * WIP making the modification test pass * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * get to the delete tests * test * improve error message * implement delitem * test KeyError * subclass Coordinates instead of DatasetCoordinates * use Frozen(self._data._coord_variables) * Simplify when to raise KeyError Co-authored-by: Stephan Hoyer <shoyer@google.com> * correct bug in suggestion * Update xarray/core/coordinates.py Co-authored-by: Stephan Hoyer <shoyer@google.com> * simplify _update_coords by creating new node data first * update indexes correctly * passes test * update ._drop_indexed_coords * some mypy fixes * remove the apparently-unused _drop_indexed_coords method * fix import error * test that Dataset and DataArray constructors can handle being passed a DataTreeCoordinates object * test dt.coords can be passed to DataTree constructor * improve readability of inline comment * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * initial tests with inherited coords * ignore typeerror indicating dodgy inheritance * try to avoid Unbound type error * cast return value correctly * cehck that .coords works with inherited coords * fix data->dataset * fix return type of __getitem__ * Use .dataset instead of .to_dataset() Co-authored-by: Stephan Hoyer <shoyer@google.com> * _check_alignment -> check_alignment * remove dict comprehension Co-authored-by: Stephan Hoyer <shoyer@google.com> * KeyError message formatting Co-authored-by: Stephan Hoyer <shoyer@google.com> * keep generic types for .dims and .sizes * test verifying you cant delete inherited coord * fix mypy complaint * type hint as accepting objects * update note about .dims returning all dims --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Stephan Hoyer <shoyer@google.com>

TomNicholas added 4 commits September 7, 2024 20:51

add a DataTreeCoordinates class

704db79

passing read-only properties tests

417e3e9

tests for modifying in-place

9562e92

WIP making the modification test pass

0e7de82

TomNicholas added the topic-DataTree Related to the implementation of a DataTree class label Sep 8, 2024

[pre-commit.ci] auto fixes from pre-commit.com hooks

839858f

for more information, see https://pre-commit.ci

TomNicholas commented Sep 8, 2024

View reviewed changes

get to the delete tests

9370b9b

shoyer reviewed Sep 8, 2024

View reviewed changes

TomNicholas added 5 commits September 7, 2024 23:36

test

9b50567

improve error message

c466f8d

implement delitem

0397eca

test KeyError

85bb221

Merge branch 'delitem' into datatree_coords_setitem

7802c63

This comment was marked as outdated.

Sign in to view

TomNicholas and others added 11 commits September 8, 2024 00:36

subclass Coordinates instead of DatasetCoordinates

1bf5082

use Frozen(self._data._coord_variables)

e8620cf

Simplify when to raise KeyError

1108504

Co-authored-by: Stephan Hoyer <shoyer@google.com>

correct bug in suggestion

0a7201b

Update xarray/core/coordinates.py

51e11bc

Co-authored-by: Stephan Hoyer <shoyer@google.com>

simplify _update_coords by creating new node data first

7ecdd16

Merge branch 'main' into datatree_coords_setitem

dfcdb6d

Merge branch 'main' into datatree_coords_setitem

3278153

update indexes correctly

f672c5e

passes test

7fb1622

update ._drop_indexed_coords

897b7c4

TomNicholas commented Sep 9, 2024

View reviewed changes

TomNicholas marked this pull request as ready for review September 9, 2024 17:40

Merge branch 'main' into datatree_coords_setitem

b5a56f4

TomNicholas mentioned this pull request Sep 9, 2024

ds.coords.dims can display dims not present in the coordinate vars #9466

Open

TomNicholas added 2 commits September 9, 2024 21:56

initial tests with inherited coords

8486227

Merge branch 'datatree_coords_setitem' of https://github.com/TomNicho…

12f24df

…las/xarray into datatree_coords_setitem

TomNicholas commented Sep 10, 2024

View reviewed changes

TomNicholas added 5 commits September 10, 2024 20:38

ignore typeerror indicating dodgy inheritance

8f09c93

try to avoid Unbound type error

d23105f

cast return value correctly

978e05e

cehck that .coords works with inherited coords

bd47575

Merge branch 'main' into datatree_coords_setitem

8ef94df

Merge branch 'main' into datatree_coords_setitem

10b8a78

shoyer reviewed Sep 11, 2024

View reviewed changes

fix data->dataset

b9ede22

shoyer approved these changes Sep 11, 2024

View reviewed changes

TomNicholas and others added 10 commits September 10, 2024 23:22

fix return type of __getitem__

540a825

Use .dataset instead of .to_dataset()

b30d5e0

Co-authored-by: Stephan Hoyer <shoyer@google.com>

_check_alignment -> check_alignment

639ad07

remove dict comprehension

0a9a328

Co-authored-by: Stephan Hoyer <shoyer@google.com>

KeyError message formatting

80bc0bd

Co-authored-by: Stephan Hoyer <shoyer@google.com>

keep generic types for .dims and .sizes

a366bf6

test verifying you cant delete inherited coord

4d352bd

fix mypy complaint

4626fa8

type hint as accepting objects

ea8a195

update note about .dims returning all dims

af94af4

TomNicholas enabled auto-merge (squash) September 11, 2024 03:48

TomNicholas merged commit 781877c into pydata:main Sep 11, 2024
27 checks passed

TomNicholas deleted the datatree_coords_setitem branch September 11, 2024 04:11

TomNicholas mentioned this pull request Sep 12, 2024

Track merging datatree into xarray #8572

Closed

27 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fix `DataTree.coords.setitem` by adding `DataTreeCoordinates` class #9451

Fix `DataTree.coords.setitem` by adding `DataTreeCoordinates` class #9451

TomNicholas commented Sep 8, 2024

TomNicholas Sep 8, 2024

shoyer Sep 8, 2024

shoyer left a comment

shoyer Sep 8, 2024

TomNicholas Sep 8, 2024

TomNicholas Sep 10, 2024

shoyer Sep 8, 2024

This comment was marked as outdated.

TomNicholas Sep 9, 2024 •

edited

Loading

TomNicholas Sep 10, 2024

TomNicholas Sep 10, 2024 •

edited

Loading

TomNicholas Sep 11, 2024

TomNicholas commented Sep 11, 2024

shoyer left a comment

		# TODO should inherited coordinates be here? It would be very hard to allow updating them...
		# But actually maybe the ChainMap approach would make this work okay??

		# TODO is there a potential bug here? What happens if a dim is only present on data variables?
		return self._data.dims

Fix DataTree.coords.__setitem__ by adding DataTreeCoordinates class #9451

Fix DataTree.coords.__setitem__ by adding DataTreeCoordinates class #9451

Conversation

TomNicholas commented Sep 8, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

shoyer left a 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

This comment was marked as outdated.

TomNicholas Sep 9, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

TomNicholas Sep 10, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

TomNicholas commented Sep 11, 2024

shoyer left a comment

Choose a reason for hiding this comment

Fix `DataTree.coords.setitem` by adding `DataTreeCoordinates` class #9451

Fix `DataTree.coords.setitem` by adding `DataTreeCoordinates` class #9451

TomNicholas Sep 9, 2024 •

edited

Loading

TomNicholas Sep 10, 2024 •

edited

Loading