DataTree.assign_coords and similar should only set node coords
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 4.2k
- Forks
- 1.4k
- Avg merge
- 2d 15h
- Merged PRs (30d)
- 14
Description
What is your issue?
Currently DataTree.assign_coords tries to call Dataset.assign_coords on every node in the subtree. This was always a little weird, but now that we have implemented coordinate inheritance (#9077) it's really silly, because we only need to assign to the root node and the coordinates will be accessible on every child.
There's also some kind of bug right now too:
In [6]: ds = Dataset(
...: data_vars={
...: "foo": (["x", "y"], np.random.randn(2, 3)),
...: },
...: coords={
...: "x": ("x", np.array([-1, -2], "int64")),
...: "y": ("y", np.array([0, 1, 2], "int64")),
...: "a": ("x", np.array([4, 5], "int64")),
...: "b": np.int64(-10),
...: },
...: )
...: dt = DataTree(data=ds)
...: dt["child"] = DataTree()
In [7]: child = dt["child"]
In [8]: child.assign_coords({'c': 11})
---------------------------------------------------------------------------
TypeError Traceback (most recent call last)
Cell In[8], line 1
----> 1 child.assign_coords({'c': 11})
File ~/Documents/Work/Code/xarray/xarray/core/datatree_mapping.py:193, in map_over_subtree.<locals>._map_over_subtree(*args, **kwargs)
190 out_data_objects[node_of_first_tree.path] = results
192 # Find out how many return values we received
--> 193 num_return_values = _check_all_return_values(out_data_objects)
195 # Reconstruct 1+ subtrees from the dict of results, by filling in all nodes of all result trees
196 original_root_path = first_tree.path
File ~/Documents/Work/Code/xarray/xarray/core/datatree_mapping.py:282, in _check_all_return_values(returned_objects)
279 """Walk through all values returned by mapping func over subtrees, raising on any invalid or inconsistent types."""
281 if all(r is None for r in returned_objects.values()):
--> 282 raise TypeError(
283 "Called supplied function on all nodes but found a return value of None for"
284 "all of them."
285 )
287 result_data_objects = [
288 (path_to_node, r)
289 for path_to_node, r in returned_objects.items()
290 if r is not None
291 ]
293 if len(result_data_objects) == 1:
294 # Only one node in the tree: no need to check consistency of results between nodes
TypeError: Called supplied function on all nodes but found a return value of None forall of them.
We should change assign_coords (and possibly any similar functions like set_coords to just operate on one node, and not try to map over the subtree.
This would also help make the coordinate-related behaviour of DataTree more similar to that of Dataset, which relates to https://github.com/pydata/xarray/issues/9203#issuecomment-2338399208.
cc @shoyer
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start at DataTree.assign_coords and the mapping logic shown in xarray/core/datatree_mapping.py, then inspect existing DataTree coordinate-related tests. Reproduce the child.assign_coords({'c': 11}) example first. Done means coordinate assignment works on the intended node without mapping over every subtree node, with related coordinate methods checked for the same behavior.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- data
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 68/100