diff --git a/datatree/datatree.py b/datatree/datatree.py index 69c2af97..02b1b424 100644 --- a/datatree/datatree.py +++ b/datatree/datatree.py @@ -1,6 +1,7 @@ from __future__ import annotations import functools import textwrap +import inspect from typing import Mapping, Hashable, Union, List, Any, Callable, Iterable, Dict @@ -11,6 +12,7 @@ from xarray.core.variable import Variable from xarray.core.combine import merge from xarray.core import dtypes, utils +from xarray.core._typed_ops import DatasetOpsMixin from .treenode import TreeNode, PathType, _init_single_treenode @@ -31,7 +33,7 @@ | | Variable("far_infrared") |-- DataNode("topography") | |-- DataNode("elevation") -| | |-- Variable("height_above_sea_level") +| | Variable("height_above_sea_level") |-- DataNode("population") """ @@ -75,7 +77,6 @@ def _map_over_subtree(tree, *args, **kwargs): """Internal function which maps func over every node in tree, returning a tree of the results.""" # Recreate and act on root node - # TODO make this of class DataTree out_tree = DataNode(name=tree.name, data=tree.ds) if out_tree.has_data: out_tree.ds = func(out_tree.ds, *args, **kwargs) @@ -91,11 +92,15 @@ def _map_over_subtree(tree, *args, **kwargs): return _map_over_subtree +_DATASET_PROPERTIES_TO_EXPOSE = ['dims', 'variables', 'encoding', 'sizes', 'attrs', 'nbytes', 'indexes', 'xindexes', + 'xindexes', 'coords', 'data_vars', 'chunks', 'real', 'imag'] + + class DatasetPropertiesMixin: """Expose properties of wrapped Dataset""" - # TODO a neater / more succinct way of doing this? - # we wouldn't need it at all if we inherited directly from Dataset... + # TODO a neater way of setting all of these? + # We wouldn't need this at all if we inherited directly from Dataset... @property def dims(self): @@ -132,14 +137,141 @@ def attrs(self): else: raise AttributeError("property is not defined for a node with no data") + + @property + def nbytes(self) -> int: + return sum(node.ds.nbytes for node in self.subtree_nodes) + + @property + def indexes(self): + if self.has_data: + return self.ds.indexes + else: + raise AttributeError("property is not defined for a node with no data") + + @property + def xindexes(self): + if self.has_data: + return self.ds.xindexes + else: + raise AttributeError("property is not defined for a node with no data") + + @property + def coords(self): + if self.has_data: + return self.ds.coords + else: + raise AttributeError("property is not defined for a node with no data") + + @property + def data_vars(self): + if self.has_data: + return self.ds.data_vars + else: + raise AttributeError("property is not defined for a node with no data") + + # TODO should this instead somehow give info about the chunking of every node? + @property + def chunks(self): + if self.has_data: + return self.ds.chunks + else: + raise AttributeError("property is not defined for a node with no data") + + @property + def real(self): + if self.has_data: + return self.ds.real + else: + raise AttributeError("property is not defined for a node with no data") + + @property + def imag(self): + if self.has_data: + return self.ds.imag + else: + raise AttributeError("property is not defined for a node with no data") + + # TODO .loc + dims.__doc__ = Dataset.dims.__doc__ variables.__doc__ = Dataset.variables.__doc__ encoding.__doc__ = Dataset.encoding.__doc__ sizes.__doc__ = Dataset.sizes.__doc__ attrs.__doc__ = Dataset.attrs.__doc__ + indexes.__doc__ = Dataset.indexes.__doc__ + xindexes.__doc__ = Dataset.xindexes.__doc__ + coords.__doc__ = Dataset.coords.__doc__ + data_vars.__doc__ = Dataset.data_vars.__doc__ + 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` decorator for more details.", width=117) + + +def _expose_methods_wrapped_to_map_over_subtree(obj, method_name, method): + """ + Expose given method on node object, but wrapped to map over whole subtree, not just that node object. + + Result is like having written this in obj's class definition: + + ``` + @map_over_subtree + def method_name(self, *args, **kwargs): + return self.method(*args, **kwargs) + ``` + """ + + # Expose Dataset method, but wrapped to map over whole subtree when called + # TODO should we be using functools.partialmethod here instead? + mapped_over_tree = functools.partial(map_over_subtree(method), obj) + setattr(obj, method_name, mapped_over_tree) + + # TODO do we really need this for ops like __add__? + # Add a line to the method's docstring explaining how it's been mapped + method_docstring = method.__doc__ + if method_docstring is not None: + updated_method_docstring = method_docstring.replace('\n', _MAPPED_DOCSTRING_ADDENDUM, 1) + setattr(obj, f'{method_name}.__doc__', updated_method_docstring) + + +# TODO equals, broadcast_equals etc. +# TODO do dask-related private methods need to be exposed? +_DATASET_DASK_METHODS_TO_EXPOSE = ['load', 'compute', 'persist', 'unify_chunks', 'chunk', 'map_blocks'] +_DATASET_METHODS_TO_EXPOSE = ['copy', 'as_numpy', '__copy__', '__deepcopy__', '__contains__', '__len__', + '__bool__', '__iter__', '__array__', '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'] +_DATASET_OPS_TO_EXPOSE = ['_unary_op', '_binary_op', '_inplace_binary_op'] +_ALL_DATASET_METHODS_TO_EXPOSE = _DATASET_DASK_METHODS_TO_EXPOSE + _DATASET_METHODS_TO_EXPOSE + _DATASET_OPS_TO_EXPOSE +# TODO methods which should not or cannot act over the whole tree, such as .to_array -class DataTree(TreeNode, DatasetPropertiesMixin): + +class DatasetMethodsMixin: + """Mixin to add Dataset methods like .mean(), but wrapped to map over all nodes in the subtree.""" + + # TODO is there a way to put this code in the class definition so we don't have to specifically call this method? + def _add_dataset_methods(self): + methods_to_expose = [(method_name, getattr(Dataset, method_name)) + for method_name in _ALL_DATASET_METHODS_TO_EXPOSE] + + for method_name, method in methods_to_expose: + _expose_methods_wrapped_to_map_over_subtree(self, method_name, method) + + +# TODO implement ArrayReduce type methods + + +class DataTree(TreeNode, DatasetPropertiesMixin, DatasetMethodsMixin): """ A tree-like hierarchical collection of xarray objects. @@ -178,14 +310,6 @@ class DataTree(TreeNode, DatasetPropertiesMixin): # TODO do we need a watch out for if methods intended only for root nodes are called on non-root nodes? # TODO add any other properties (maybe dask ones?) - _DS_PROPERTIES = ['variables', 'attrs', 'encoding', 'dims', 'sizes'] - - # TODO add all the other methods to dispatch - _DS_METHODS_TO_MAP_OVER_SUBTREES = ['isel', 'sel', 'min', 'max', 'mean', '__array_ufunc__'] - _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 datatree.map_over_subtree decorator for more details.", - width=117) # TODO currently allows self.ds = None, should we instead always store at least an empty Dataset? @@ -218,24 +342,14 @@ def __init__( new_node = self.get_node(path) new_node[path] = data - self._add_method_api() - - def _add_method_api(self): - # Add methods defined in Dataset's class definition to this classes API, but wrapped to map over descendants too - for method_name in self._DS_METHODS_TO_MAP_OVER_SUBTREES: - # Expose Dataset method, but wrapped to map over whole subtree - ds_method = getattr(Dataset, method_name) - setattr(self, method_name, map_over_subtree(ds_method)) + # TODO this has to be + self._add_all_dataset_api() - # Add a line to the method's docstring explaining how it's been mapped - ds_method_docstring = getattr(Dataset, f'{method_name}').__doc__ - if ds_method_docstring is not None: - updated_method_docstring = ds_method_docstring.replace('\n', self._MAPPED_DOCSTRING_ADDENDUM, 1) - setattr(self, f'{method_name}.__doc__', updated_method_docstring) + def _add_all_dataset_api(self): + # Add methods like .isel(), but wrapped to map over subtrees + self._add_dataset_methods() - # TODO wrap methods for ops too, such as those in DatasetOpsMixin - - # TODO map applied ufuncs over all leaves + # TODO add dataset ops here @property def ds(self) -> Dataset: @@ -257,7 +371,7 @@ def has_data(self): def _init_single_datatree_node( cls, name: Hashable, - data: Dataset = None, + data: Union[Dataset, DataArray] = None, parent: TreeNode = None, children: List[TreeNode] = None, ): @@ -285,6 +399,9 @@ def _init_single_datatree_node( obj = object.__new__(cls) obj = _init_single_treenode(obj, name=name, parent=parent, children=children) obj.ds = data + + obj._add_all_dataset_api() + return obj def __str__(self): @@ -559,13 +676,6 @@ def get_any(self, *tags: Hashable) -> DataTree: if any(tag in c.tags for tag in tags)} return DataTree(data_objects=matching_children) - @property - def chunks(self): - raise NotImplementedError - - def chunk(self): - raise NotImplementedError - def merge(self, datatree: DataTree) -> DataTree: """Merge all the leaves of a second DataTree into this one.""" raise NotImplementedError @@ -586,7 +696,7 @@ def merge_child_datasets( datasets = [self.get(path).ds for path in paths] return merge(datasets, compat=compat, join=join, fill_value=fill_value, combine_attrs=combine_attrs) - def as_dataarray(self) -> DataArray: + def as_array(self) -> DataArray: return self.ds.as_dataarray() @property diff --git a/datatree/tests/test_dataset_api.py b/datatree/tests/test_dataset_api.py index ea3cd920..c6d0f150 100644 --- a/datatree/tests/test_dataset_api.py +++ b/datatree/tests/test_dataset_api.py @@ -1,5 +1,7 @@ import pytest +import numpy as np + import xarray as xr from xarray.testing import assert_equal @@ -78,6 +80,7 @@ def test_properties(self): assert dt.sizes == dt.ds.sizes assert dt.variables == dt.ds.variables + def test_no_data_no_properties(self): dt = DataNode('root', data=None) with pytest.raises(AttributeError): @@ -93,12 +96,39 @@ def test_no_data_no_properties(self): class TestDSMethodInheritance: + def test_root(self): + da = xr.DataArray(name='a', data=[1, 2, 3], dims='x') + dt = DataNode('root', data=da) + expected_ds = da.to_dataset().isel(x=1) + result_ds = dt.isel(x=1).ds + assert_equal(result_ds, expected_ds) + + def test_descendants(self): + da = xr.DataArray(name='a', data=[1, 2, 3], dims='x') + dt = DataNode('root') + DataNode('results', parent=dt, data=da) + expected_ds = da.to_dataset().isel(x=1) + result_ds = dt.isel(x=1)['results'].ds + assert_equal(result_ds, expected_ds) + + +class TestOps: ... -class TestBinaryOps: - ... - - +@pytest.mark.xfail class TestUFuncs: - ... + def test_root(self): + da = xr.DataArray(name='a', data=[1, 2, 3]) + dt = DataNode('root', data=da) + expected_ds = np.sin(da.to_dataset()) + result_ds = np.sin(dt).ds + assert_equal(result_ds, expected_ds) + + def test_descendants(self): + da = xr.DataArray(name='a', data=[1, 2, 3]) + dt = DataNode('root') + DataNode('results', parent=dt, data=da) + expected_ds = np.sin(da.to_dataset()) + result_ds = np.sin(dt)['results'].ds + assert_equal(result_ds, expected_ds)