Skip to content

Commit

Permalink
map_over_subtree -> map_over_datasets (pydata#9622)
Browse files Browse the repository at this point in the history
  • Loading branch information
TomNicholas authored Oct 15, 2024
1 parent 017279c commit 33ead65
Show file tree
Hide file tree
Showing 9 changed files with 54 additions and 54 deletions.
4 changes: 2 additions & 2 deletions doc/api.rst
Original file line number Diff line number Diff line change
Expand Up @@ -687,7 +687,7 @@ For manipulating, traversing, navigating, or mapping over the tree structure.
DataTree.relative_to
DataTree.iter_lineage
DataTree.find_common_ancestor
DataTree.map_over_subtree
DataTree.map_over_datasets
DataTree.pipe
DataTree.match
DataTree.filter
Expand Down Expand Up @@ -954,7 +954,7 @@ DataTree methods

open_datatree
open_groups
map_over_subtree
map_over_datasets
DataTree.to_dict
DataTree.to_netcdf
DataTree.to_zarr
Expand Down
8 changes: 4 additions & 4 deletions doc/user-guide/hierarchical-data.rst
Original file line number Diff line number Diff line change
Expand Up @@ -549,13 +549,13 @@ See that the same change (fast-forwarding by adding 10 years to the age of each
Mapping Custom Functions Over Trees
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

You can map custom computation over each node in a tree using :py:meth:`xarray.DataTree.map_over_subtree`.
You can map custom computation over each node in a tree using :py:meth:`xarray.DataTree.map_over_datasets`.
You can map any function, so long as it takes :py:class:`xarray.Dataset` objects as one (or more) of the input arguments,
and returns one (or more) xarray datasets.

.. note::

Functions passed to :py:func:`~xarray.DataTree.map_over_subtree` cannot alter nodes in-place.
Functions passed to :py:func:`~xarray.DataTree.map_over_datasets` cannot alter nodes in-place.
Instead they must return new :py:class:`xarray.Dataset` objects.

For example, we can define a function to calculate the Root Mean Square of a timeseries
Expand All @@ -569,11 +569,11 @@ Then calculate the RMS value of these signals:

.. ipython:: python
voltages.map_over_subtree(rms)
voltages.map_over_datasets(rms)
.. _multiple trees:

We can also use the :py:meth:`~xarray.map_over_subtree` decorator to promote a function which accepts datasets into one which
We can also use the :py:meth:`~xarray.map_over_datasets` decorator to promote a function which accepts datasets into one which
accepts datatrees.

Operating on Multiple Trees
Expand Down
2 changes: 1 addition & 1 deletion doc/whats-new.rst
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ New Features
~~~~~~~~~~~~
- ``DataTree`` related functionality is now exposed in the main ``xarray`` public
API. This includes: ``xarray.DataTree``, ``xarray.open_datatree``, ``xarray.open_groups``,
``xarray.map_over_subtree``, ``xarray.register_datatree_accessor`` and
``xarray.map_over_datasets``, ``xarray.register_datatree_accessor`` and
``xarray.testing.assert_isomorphic``.
By `Owen Littlejohns <https://github.com/owenlittlejohns>`_,
`Eni Awowale <https://github.com/eni-awowale>`_,
Expand Down
4 changes: 2 additions & 2 deletions xarray/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@
from xarray.core.dataarray import DataArray
from xarray.core.dataset import Dataset
from xarray.core.datatree import DataTree
from xarray.core.datatree_mapping import TreeIsomorphismError, map_over_subtree
from xarray.core.datatree_mapping import TreeIsomorphismError, map_over_datasets
from xarray.core.extensions import (
register_dataarray_accessor,
register_dataset_accessor,
Expand Down Expand Up @@ -86,7 +86,7 @@
"load_dataarray",
"load_dataset",
"map_blocks",
"map_over_subtree",
"map_over_datasets",
"merge",
"ones_like",
"open_dataarray",
Expand Down
14 changes: 7 additions & 7 deletions xarray/core/datatree.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@
from xarray.core.datatree_mapping import (
TreeIsomorphismError,
check_isomorphic,
map_over_subtree,
map_over_datasets,
)
from xarray.core.formatting import datatree_repr, dims_and_coords_repr
from xarray.core.formatting_html import (
Expand Down Expand Up @@ -254,14 +254,14 @@ def _constructor(
def __setitem__(self, key, val) -> None:
raise AttributeError(
"Mutation of the DatasetView is not allowed, please use `.__setitem__` on the wrapping DataTree node, "
"or use `dt.to_dataset()` if you want a mutable dataset. If calling this from within `map_over_subtree`,"
"or use `dt.to_dataset()` if you want a mutable dataset. If calling this from within `map_over_datasets`,"
"use `.copy()` first to get a mutable version of the input dataset."
)

def update(self, other) -> NoReturn:
raise AttributeError(
"Mutation of the DatasetView is not allowed, please use `.update` on the wrapping DataTree node, "
"or use `dt.to_dataset()` if you want a mutable dataset. If calling this from within `map_over_subtree`,"
"or use `dt.to_dataset()` if you want a mutable dataset. If calling this from within `map_over_datasets`,"
"use `.copy()` first to get a mutable version of the input dataset."
)

Expand Down Expand Up @@ -1330,7 +1330,7 @@ def filter(self: DataTree, filterfunc: Callable[[DataTree], bool]) -> DataTree:
--------
match
pipe
map_over_subtree
map_over_datasets
"""
filtered_nodes = {
node.path: node.dataset for node in self.subtree if filterfunc(node)
Expand All @@ -1356,7 +1356,7 @@ def match(self, pattern: str) -> DataTree:
--------
filter
pipe
map_over_subtree
map_over_datasets
Examples
--------
Expand All @@ -1383,7 +1383,7 @@ def match(self, pattern: str) -> DataTree:
}
return DataTree.from_dict(matching_nodes, name=self.root.name)

def map_over_subtree(
def map_over_datasets(
self,
func: Callable,
*args: Iterable[Any],
Expand Down Expand Up @@ -1417,7 +1417,7 @@ def map_over_subtree(
# TODO this signature means that func has no way to know which node it is being called upon - change?

# TODO fix this typing error
return map_over_subtree(func)(self, *args, **kwargs)
return map_over_datasets(func)(self, *args, **kwargs)

def pipe(
self, func: Callable | tuple[Callable, str], *args: Any, **kwargs: Any
Expand Down
10 changes: 5 additions & 5 deletions xarray/core/datatree_mapping.py
Original file line number Diff line number Diff line change
Expand Up @@ -72,7 +72,7 @@ def check_isomorphic(
raise TreeIsomorphismError("DataTree objects are not isomorphic:\n" + diff)


def map_over_subtree(func: Callable) -> Callable:
def map_over_datasets(func: Callable) -> Callable:
"""
Decorator which turns a function which acts on (and returns) Datasets into one which acts on and returns DataTrees.
Expand Down Expand Up @@ -112,8 +112,8 @@ def map_over_subtree(func: Callable) -> Callable:
See also
--------
DataTree.map_over_subtree
DataTree.map_over_subtree_inplace
DataTree.map_over_datasets
DataTree.map_over_datasets_inplace
DataTree.subtree
"""

Expand All @@ -122,7 +122,7 @@ def map_over_subtree(func: Callable) -> Callable:
# TODO inspect function to work out immediately if the wrong number of arguments were passed for it?

@functools.wraps(func)
def _map_over_subtree(*args, **kwargs) -> DataTree | tuple[DataTree, ...]:
def _map_over_datasets(*args, **kwargs) -> DataTree | tuple[DataTree, ...]:
"""Internal function which maps func over every node in tree, returning a tree of the results."""
from xarray.core.datatree import DataTree

Expand Down Expand Up @@ -227,7 +227,7 @@ def _map_over_subtree(*args, **kwargs) -> DataTree | tuple[DataTree, ...]:
else:
return tuple(result_trees)

return _map_over_subtree
return _map_over_datasets


def _handle_errors_with_path_context(path: str):
Expand Down
14 changes: 7 additions & 7 deletions xarray/core/datatree_ops.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
import textwrap

from xarray.core.dataset import Dataset
from xarray.core.datatree_mapping import map_over_subtree
from xarray.core.datatree_mapping import map_over_datasets

"""
Module which specifies the subset of xarray.Dataset's API which we wish to copy onto DataTree.
Expand All @@ -17,7 +17,7 @@
_MAPPED_DOCSTRING_ADDENDUM = (
"This method was copied from :py:class:`xarray.Dataset`, but has been altered to "
"call the method on the Datasets stored in every node of the subtree. "
"See the `map_over_subtree` function for more details."
"See the `map_over_datasets` function for more details."
)

# TODO equals, broadcast_equals etc.
Expand Down Expand Up @@ -174,7 +174,7 @@ def _wrap_then_attach_to_cls(
target_cls_dict, source_cls, methods_to_set, wrap_func=None
):
"""
Attach given methods on a class, and optionally wrap each method first. (i.e. with map_over_subtree).
Attach given methods on a class, and optionally wrap each method first. (i.e. with map_over_datasets).
Result is like having written this in the classes' definition:
```
Expand Down Expand Up @@ -206,7 +206,7 @@ def method_name(self, *args, **kwargs):
)
target_cls_dict[method_name] = wrapped_method

if wrap_func is map_over_subtree:
if wrap_func is map_over_datasets:
# Add a paragraph to the method's docstring explaining how it's been mapped
orig_method_docstring = orig_method.__doc__

Expand Down Expand Up @@ -277,7 +277,7 @@ class MappedDatasetMethodsMixin:
target_cls_dict=vars(),
source_cls=Dataset,
methods_to_set=_ALL_DATASET_METHODS_TO_MAP,
wrap_func=map_over_subtree,
wrap_func=map_over_datasets,
)


Expand All @@ -291,7 +291,7 @@ class MappedDataWithCoords:
target_cls_dict=vars(),
source_cls=Dataset,
methods_to_set=_DATA_WITH_COORDS_METHODS_TO_MAP,
wrap_func=map_over_subtree,
wrap_func=map_over_datasets,
)


Expand All @@ -305,5 +305,5 @@ class DataTreeArithmeticMixin:
target_cls_dict=vars(),
source_cls=Dataset,
methods_to_set=_ARITHMETIC_METHODS_TO_MAP,
wrap_func=map_over_subtree,
wrap_func=map_over_datasets,
)
6 changes: 3 additions & 3 deletions xarray/tests/test_datatree.py
Original file line number Diff line number Diff line change
Expand Up @@ -1805,7 +1805,7 @@ def test_tree(self, create_test_datatree):


class TestDocInsertion:
"""Tests map_over_subtree docstring injection."""
"""Tests map_over_datasets docstring injection."""

def test_standard_doc(self):

Expand Down Expand Up @@ -1839,7 +1839,7 @@ def test_standard_doc(self):
.. note::
This method was copied from :py:class:`xarray.Dataset`, but has
been altered to call the method on the Datasets stored in every
node of the subtree. See the `map_over_subtree` function for more
node of the subtree. See the `map_over_datasets` function for more
details.
Normally, it should not be necessary to call this method in user code,
Expand Down Expand Up @@ -1870,7 +1870,7 @@ def test_one_liner(self):
This method was copied from :py:class:`xarray.Dataset`, but has been altered to
call the method on the Datasets stored in every node of the subtree. See
the `map_over_subtree` function for more details."""
the `map_over_datasets` function for more details."""
)

actual_doc = insert_doc_addendum(mixin_doc, _MAPPED_DOCSTRING_ADDENDUM)
Expand Down
Loading

0 comments on commit 33ead65

Please sign in to comment.