diff --git a/datatree/datatree.py b/datatree/datatree.py index e39b0c05..76fc1bf9 100644 --- a/datatree/datatree.py +++ b/datatree/datatree.py @@ -4,45 +4,25 @@ from typing import Any, Callable, Dict, Hashable, Iterable, List, Mapping, Union import anytree +from xarray import DataArray, Dataset, merge from xarray.core import dtypes, utils -from xarray.core.arithmetic import DatasetArithmetic -from xarray.core.combine import merge -from xarray.core.common import DataWithCoords -from xarray.core.dataarray import DataArray -from xarray.core.dataset import Dataset -from xarray.core.ops import NAN_CUM_METHODS, NAN_REDUCE_METHODS, REDUCE_METHODS from xarray.core.variable import Variable from .mapping import map_over_subtree +from .ops import ( + DataTreeArithmeticMixin, + MappedDatasetMethodsMixin, + MappedDataWithCoords, +) from .treenode import PathType, TreeNode, _init_single_treenode """ -The structure of a populated Datatree looks roughly like this: - -DataTree("root name") -|-- DataNode("weather") -| | Variable("wind_speed") -| | Variable("pressure") -| |-- DataNode("temperature") -| | Variable("sea_surface_temperature") -| | Variable("dew_point_temperature") -|-- DataNode("satellite image") -| | Variable("true_colour") -| |-- DataNode("infrared") -| | Variable("near_infrared") -| | Variable("far_infrared") -|-- DataNode("topography") -| |-- DataNode("elevation") -| | Variable("height_above_sea_level") -|-- DataNode("population") - - DEVELOPERS' NOTE ---------------- The idea of this module is to create a `DataTree` class which inherits the tree structure from TreeNode, and also copies the entire API of `xarray.Dataset`, but with certain methods decorated to instead map the dataset function over every node in the tree. As this API is copied without directly subclassing `xarray.Dataset` we instead create various Mixin -classes which each define part of `xarray.Dataset`'s extensive API. +classes (in ops.py) which each define part of `xarray.Dataset`'s extensive API. Some of these methods must be wrapped to map over all nodes in the subtree. Others are fine to inherit unaltered (normally because they (a) only call dataset properties and (b) don't return a dataset that should be nested into a new @@ -56,6 +36,8 @@ class DatasetPropertiesMixin: # TODO a neater way of setting all of these? # We wouldn't need this at all if we inherited directly from Dataset... + # TODO we could also just not define these at all, and require users to call e.g. dt.ds.dims ... + @property def dims(self): if self.has_data: @@ -159,202 +141,12 @@ def imag(self): chunks.__doc__ = Dataset.chunks.__doc__ -_MAPPED_DOCSTRING_ADDENDUM = textwrap.fill( - "This method was copied from 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.", - width=117, -) - -# TODO equals, broadcast_equals etc. -# TODO do dask-related private methods need to be exposed? -_DATASET_DASK_METHODS_TO_MAP = [ - "load", - "compute", - "persist", - "unify_chunks", - "chunk", - "map_blocks", -] -_DATASET_METHODS_TO_MAP = [ - "copy", - "as_numpy", - "__copy__", - "__deepcopy__", - "set_coords", - "reset_coords", - "info", - "isel", - "sel", - "head", - "tail", - "thin", - "broadcast_like", - "reindex_like", - "reindex", - "interp", - "interp_like", - "rename", - "rename_dims", - "rename_vars", - "swap_dims", - "expand_dims", - "set_index", - "reset_index", - "reorder_levels", - "stack", - "unstack", - "update", - "merge", - "drop_vars", - "drop_sel", - "drop_isel", - "drop_dims", - "transpose", - "dropna", - "fillna", - "interpolate_na", - "ffill", - "bfill", - "combine_first", - "reduce", - "map", - "assign", - "diff", - "shift", - "roll", - "sortby", - "quantile", - "rank", - "differentiate", - "integrate", - "cumulative_integrate", - "filter_by_attrs", - "polyfit", - "pad", - "idxmin", - "idxmax", - "argmin", - "argmax", - "query", - "curvefit", -] -# TODO unsure if these are called by external functions or not? -_DATASET_OPS_TO_MAP = ["_unary_op", "_binary_op", "_inplace_binary_op"] -_ALL_DATASET_METHODS_TO_MAP = ( - _DATASET_DASK_METHODS_TO_MAP + _DATASET_METHODS_TO_MAP + _DATASET_OPS_TO_MAP -) - -_DATA_WITH_COORDS_METHODS_TO_MAP = [ - "squeeze", - "clip", - "assign_coords", - "where", - "close", - "isnull", - "notnull", - "isin", - "astype", -] - -# TODO NUM_BINARY_OPS apparently aren't defined on DatasetArithmetic, and don't appear to be injected anywhere... -_ARITHMETIC_METHODS_TO_MAP = ( - REDUCE_METHODS + NAN_REDUCE_METHODS + NAN_CUM_METHODS + ["__array_ufunc__"] -) - - -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) - - Result is like having written this in the classes' definition: - ``` - @wrap_func - def method_name(self, *args, **kwargs): - return self.method(*args, **kwargs) - ``` - - Every method attached here needs to have a return value of Dataset or DataArray in order to construct a new tree. - - Parameters - ---------- - target_cls_dict : MappingProxy - The __dict__ attribute of the class which we want the methods to be added to. (The __dict__ attribute can also - be accessed by calling vars() from within that classes' definition.) This will be updated by this function. - source_cls : class - Class object from which we want to copy methods (and optionally wrap them). Should be the actual class object - (or instance), not just the __dict__. - methods_to_set : Iterable[Tuple[str, callable]] - The method names and definitions supplied as a list of (method_name_string, method) pairs. - This format matches the output of inspect.getmembers(). - wrap_func : callable, optional - Function to decorate each method with. Must have the same return type as the method. - """ - for method_name in methods_to_set: - orig_method = getattr(source_cls, method_name) - wrapped_method = ( - wrap_func(orig_method) if wrap_func is not None else orig_method - ) - target_cls_dict[method_name] = wrapped_method - - if wrap_func is map_over_subtree: - # Add a paragraph to the method's docstring explaining how it's been mapped - orig_method_docstring = orig_method.__doc__ - if orig_method_docstring is not None: - if "\n" in orig_method_docstring: - new_method_docstring = orig_method_docstring.replace( - "\n", _MAPPED_DOCSTRING_ADDENDUM, 1 - ) - else: - new_method_docstring = ( - orig_method_docstring + f"\n\n{_MAPPED_DOCSTRING_ADDENDUM}" - ) - setattr(target_cls_dict[method_name], "__doc__", new_method_docstring) - - -class MappedDatasetMethodsMixin: - """ - Mixin to add Dataset methods like .mean(), but wrapped to map over all nodes in the subtree. - """ - - __slots__ = () - _wrap_then_attach_to_cls( - vars(), Dataset, _ALL_DATASET_METHODS_TO_MAP, wrap_func=map_over_subtree - ) - - -class MappedDataWithCoords(DataWithCoords): - # TODO add mapped versions of groupby, weighted, rolling, rolling_exp, coarsen, resample - # TODO re-implement AttrsAccessMixin stuff so that it includes access to child nodes - _wrap_then_attach_to_cls( - vars(), - DataWithCoords, - _DATA_WITH_COORDS_METHODS_TO_MAP, - wrap_func=map_over_subtree, - ) - - -class DataTreeArithmetic(DatasetArithmetic): - """ - Mixin to add Dataset methods like __add__ and .mean(). - """ - - _wrap_then_attach_to_cls( - vars(), - DatasetArithmetic, - _ARITHMETIC_METHODS_TO_MAP, - wrap_func=map_over_subtree, - ) - - class DataTree( TreeNode, DatasetPropertiesMixin, MappedDatasetMethodsMixin, MappedDataWithCoords, - DataTreeArithmetic, + DataTreeArithmeticMixin, ): """ A tree-like hierarchical collection of xarray objects. @@ -858,7 +650,7 @@ def to_netcdf( def to_zarr(self, store, mode: str = "w", encoding=None, **kwargs): """ - Write datatree contents to a netCDF file. + Write datatree contents to a Zarr store. Parameters --------- @@ -875,7 +667,7 @@ def to_zarr(self, store, mode: str = "w", encoding=None, **kwargs): ``{"root/set1": {"my_variable": {"dtype": "int16", "scale_factor": 0.1}, ...}, ...}``. See ``xarray.Dataset.to_zarr`` for available options. kwargs : - Addional keyword arguments to be passed to ``xarray.Dataset.to_zarr`` + Additional keyword arguments to be passed to ``xarray.Dataset.to_zarr`` """ from .io import _datatree_to_zarr diff --git a/datatree/ops.py b/datatree/ops.py new file mode 100644 index 00000000..e411c973 --- /dev/null +++ b/datatree/ops.py @@ -0,0 +1,268 @@ +import textwrap + +from xarray import Dataset + +from .mapping import map_over_subtree + +""" +Module which specifies the subset of xarray.Dataset's API which we wish to copy onto DataTree. + +Structured to mirror the way xarray defines Dataset's various operations internally, but does not actually import from +xarray's internals directly, only the public-facing xarray.Dataset class. +""" + + +_MAPPED_DOCSTRING_ADDENDUM = textwrap.fill( + "This method was copied from 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.", + width=117, +) + +# TODO equals, broadcast_equals etc. +# TODO do dask-related private methods need to be exposed? +_DATASET_DASK_METHODS_TO_MAP = [ + "load", + "compute", + "persist", + "unify_chunks", + "chunk", + "map_blocks", +] +_DATASET_METHODS_TO_MAP = [ + "copy", + "as_numpy", + "__copy__", + "__deepcopy__", + "set_coords", + "reset_coords", + "info", + "isel", + "sel", + "head", + "tail", + "thin", + "broadcast_like", + "reindex_like", + "reindex", + "interp", + "interp_like", + "rename", + "rename_dims", + "rename_vars", + "swap_dims", + "expand_dims", + "set_index", + "reset_index", + "reorder_levels", + "stack", + "unstack", + "update", + "merge", + "drop_vars", + "drop_sel", + "drop_isel", + "drop_dims", + "transpose", + "dropna", + "fillna", + "interpolate_na", + "ffill", + "bfill", + "combine_first", + "reduce", + "map", + "assign", + "diff", + "shift", + "roll", + "sortby", + "quantile", + "rank", + "differentiate", + "integrate", + "cumulative_integrate", + "filter_by_attrs", + "polyfit", + "pad", + "idxmin", + "idxmax", + "argmin", + "argmax", + "query", + "curvefit", +] +_ALL_DATASET_METHODS_TO_MAP = _DATASET_DASK_METHODS_TO_MAP + _DATASET_METHODS_TO_MAP + +_DATA_WITH_COORDS_METHODS_TO_MAP = [ + "squeeze", + "clip", + "assign_coords", + "where", + "close", + "isnull", + "notnull", + "isin", + "astype", +] + +REDUCE_METHODS = ["all", "any"] +NAN_REDUCE_METHODS = [ + "max", + "min", + "mean", + "prod", + "sum", + "std", + "var", + "median", +] +NAN_CUM_METHODS = ["cumsum", "cumprod"] +_TYPED_DATASET_OPS_TO_MAP = [ + "__add__", + "__sub__", + "__mul__", + "__pow__", + "__truediv__", + "__floordiv__", + "__mod__", + "__and__", + "__xor__", + "__or__", + "__lt__", + "__le__", + "__gt__", + "__ge__", + "__eq__", + "__ne__", + "__radd__", + "__rsub__", + "__rmul__", + "__rpow__", + "__rtruediv__", + "__rfloordiv__", + "__rmod__", + "__rand__", + "__rxor__", + "__ror__", + "__iadd__", + "__isub__", + "__imul__", + "__ipow__", + "__itruediv__", + "__ifloordiv__", + "__imod__", + "__iand__", + "__ixor__", + "__ior__", + "__neg__", + "__pos__", + "__abs__", + "__invert__", + "round", + "argsort", + "conj", + "conjugate", +] +# TODO NUM_BINARY_OPS apparently aren't defined on DatasetArithmetic, and don't appear to be injected anywhere... +_ARITHMETIC_METHODS_TO_MAP = ( + REDUCE_METHODS + + NAN_REDUCE_METHODS + + NAN_CUM_METHODS + + _TYPED_DATASET_OPS_TO_MAP + + ["__array_ufunc__"] +) + + +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) + + Result is like having written this in the classes' definition: + ``` + @wrap_func + def method_name(self, *args, **kwargs): + return self.method(*args, **kwargs) + ``` + + Every method attached here needs to have a return value of Dataset or DataArray in order to construct a new tree. + + Parameters + ---------- + target_cls_dict : MappingProxy + The __dict__ attribute of the class which we want the methods to be added to. (The __dict__ attribute can also + be accessed by calling vars() from within that classes' definition.) This will be updated by this function. + source_cls : class + Class object from which we want to copy methods (and optionally wrap them). Should be the actual class object + (or instance), not just the __dict__. + methods_to_set : Iterable[Tuple[str, callable]] + The method names and definitions supplied as a list of (method_name_string, method) pairs. + This format matches the output of inspect.getmembers(). + wrap_func : callable, optional + Function to decorate each method with. Must have the same return type as the method. + """ + for method_name in methods_to_set: + orig_method = getattr(source_cls, method_name) + wrapped_method = ( + wrap_func(orig_method) if wrap_func is not None else orig_method + ) + target_cls_dict[method_name] = wrapped_method + + if wrap_func is map_over_subtree: + # Add a paragraph to the method's docstring explaining how it's been mapped + orig_method_docstring = orig_method.__doc__ + if orig_method_docstring is not None: + if "\n" in orig_method_docstring: + new_method_docstring = orig_method_docstring.replace( + "\n", _MAPPED_DOCSTRING_ADDENDUM, 1 + ) + else: + new_method_docstring = ( + orig_method_docstring + f"\n\n{_MAPPED_DOCSTRING_ADDENDUM}" + ) + setattr(target_cls_dict[method_name], "__doc__", new_method_docstring) + + +class MappedDatasetMethodsMixin: + """ + Mixin to add methods defined specifically on the Dataset class such as .query(), but wrapped to map over all nodes + in the subtree. + """ + + _wrap_then_attach_to_cls( + target_cls_dict=vars(), + source_cls=Dataset, + methods_to_set=_ALL_DATASET_METHODS_TO_MAP, + wrap_func=map_over_subtree, + ) + + +class MappedDataWithCoords: + """ + Mixin to add coordinate-aware Dataset methods such as .where(), but wrapped to map over all nodes in the subtree. + """ + + # TODO add mapped versions of groupby, weighted, rolling, rolling_exp, coarsen, resample + # TODO re-implement AttrsAccessMixin stuff so that it includes access to child nodes + _wrap_then_attach_to_cls( + target_cls_dict=vars(), + source_cls=Dataset, + methods_to_set=_DATA_WITH_COORDS_METHODS_TO_MAP, + wrap_func=map_over_subtree, + ) + + +class DataTreeArithmeticMixin: + """ + Mixin to add Dataset arithmetic operations such as __add__, reduction methods such as .mean(), and enable numpy + ufuncs such as np.sin(), but wrapped to map over all nodes in the subtree. + """ + + _wrap_then_attach_to_cls( + target_cls_dict=vars(), + source_cls=Dataset, + methods_to_set=_ARITHMETIC_METHODS_TO_MAP, + wrap_func=map_over_subtree, + ) diff --git a/datatree/tests/test_dataset_api.py b/datatree/tests/test_dataset_api.py index e930f49f..2c6a4d52 100644 --- a/datatree/tests/test_dataset_api.py +++ b/datatree/tests/test_dataset_api.py @@ -5,6 +5,8 @@ from datatree import DataNode +from .test_datatree import assert_tree_equal, create_test_datatree + class TestDSProperties: def test_properties(self): @@ -88,8 +90,36 @@ def test_cum_method(self): class TestOps: - @pytest.mark.xfail - def test_binary_op(self): + def test_binary_op_on_int(self): + ds1 = xr.Dataset({"a": [5], "b": [3]}) + ds2 = xr.Dataset({"x": [0.1, 0.2], "y": [10, 20]}) + dt = DataNode("root", data=ds1) + DataNode("subnode", data=ds2, parent=dt) + + expected_root = DataNode("root", data=ds1 * 5) + expected_descendant = DataNode("subnode", data=ds2 * 5, parent=expected_root) + result = dt * 5 + + assert_equal(result.ds, expected_root.ds) + assert_equal(result["subnode"].ds, expected_descendant.ds) + + def test_binary_op_on_dataset(self): + ds1 = xr.Dataset({"a": [5], "b": [3]}) + ds2 = xr.Dataset({"x": [0.1, 0.2], "y": [10, 20]}) + dt = DataNode("root", data=ds1) + DataNode("subnode", data=ds2, parent=dt) + other_ds = xr.Dataset({"z": ("z", [0.1, 0.2])}) + + expected_root = DataNode("root", data=ds1 * other_ds) + expected_descendant = DataNode( + "subnode", data=ds2 * other_ds, parent=expected_root + ) + result = dt * other_ds + + assert_equal(result.ds, expected_root.ds) + assert_equal(result["subnode"].ds, expected_descendant.ds) + + def test_binary_op_on_datatree(self): ds1 = xr.Dataset({"a": [5], "b": [3]}) ds2 = xr.Dataset({"x": [0.1, 0.2], "y": [10, 20]}) dt = DataNode("root", data=ds1) @@ -103,7 +133,6 @@ def test_binary_op(self): assert_equal(result["subnode"].ds, expected_descendant.ds) -@pytest.mark.xfail class TestUFuncs: def test_root(self): da = xr.DataArray(name="a", data=[1, 2, 3]) @@ -119,3 +148,9 @@ def test_descendants(self): expected_ds = np.sin(da.to_dataset()) result_ds = np.sin(dt)["results"].ds assert_equal(result_ds, expected_ds) + + def test_tree(self): + dt = create_test_datatree() + expected = create_test_datatree(modify=lambda ds: np.sin(ds)) + result_tree = np.sin(dt) + assert_tree_equal(result_tree, expected)