diff --git a/asv_bench/benchmarks/dataset_io.py b/asv_bench/benchmarks/dataset_io.py index dcc2de0473b..0956be67dad 100644 --- a/asv_bench/benchmarks/dataset_io.py +++ b/asv_bench/benchmarks/dataset_io.py @@ -7,6 +7,8 @@ import pandas as pd import xarray as xr +from xarray.backends.api import open_datatree +from xarray.core.datatree import DataTree from . import _skip_slow, parameterized, randint, randn, requires_dask @@ -16,7 +18,6 @@ except ImportError: pass - os.environ["HDF5_USE_FILE_LOCKING"] = "FALSE" _ENGINES = tuple(xr.backends.list_engines().keys() - {"store"}) @@ -469,6 +470,116 @@ def create_delayed_write(): return ds.to_netcdf("file.nc", engine="netcdf4", compute=False) +class IONestedDataTree: + """ + A few examples that benchmark reading/writing a heavily nested netCDF datatree with + xarray + """ + + timeout = 300.0 + repeat = 1 + number = 5 + + def make_datatree(self, nchildren=10): + # multiple Dataset + self.ds = xr.Dataset() + self.nt = 1000 + self.nx = 90 + self.ny = 45 + self.nchildren = nchildren + + self.block_chunks = { + "time": self.nt / 4, + "lon": self.nx / 3, + "lat": self.ny / 3, + } + + self.time_chunks = {"time": int(self.nt / 36)} + + times = pd.date_range("1970-01-01", periods=self.nt, freq="D") + lons = xr.DataArray( + np.linspace(0, 360, self.nx), + dims=("lon",), + attrs={"units": "degrees east", "long_name": "longitude"}, + ) + lats = xr.DataArray( + np.linspace(-90, 90, self.ny), + dims=("lat",), + attrs={"units": "degrees north", "long_name": "latitude"}, + ) + self.ds["foo"] = xr.DataArray( + randn((self.nt, self.nx, self.ny), frac_nan=0.2), + coords={"lon": lons, "lat": lats, "time": times}, + dims=("time", "lon", "lat"), + name="foo", + attrs={"units": "foo units", "description": "a description"}, + ) + self.ds["bar"] = xr.DataArray( + randn((self.nt, self.nx, self.ny), frac_nan=0.2), + coords={"lon": lons, "lat": lats, "time": times}, + dims=("time", "lon", "lat"), + name="bar", + attrs={"units": "bar units", "description": "a description"}, + ) + self.ds["baz"] = xr.DataArray( + randn((self.nx, self.ny), frac_nan=0.2).astype(np.float32), + coords={"lon": lons, "lat": lats}, + dims=("lon", "lat"), + name="baz", + attrs={"units": "baz units", "description": "a description"}, + ) + + self.ds.attrs = {"history": "created for xarray benchmarking"} + + self.oinds = { + "time": randint(0, self.nt, 120), + "lon": randint(0, self.nx, 20), + "lat": randint(0, self.ny, 10), + } + self.vinds = { + "time": xr.DataArray(randint(0, self.nt, 120), dims="x"), + "lon": xr.DataArray(randint(0, self.nx, 120), dims="x"), + "lat": slice(3, 20), + } + root = {f"group_{group}": self.ds for group in range(self.nchildren)} + nested_tree1 = { + f"group_{group}/subgroup_1": xr.Dataset() for group in range(self.nchildren) + } + nested_tree2 = { + f"group_{group}/subgroup_2": xr.DataArray(np.arange(1, 10)).to_dataset( + name="a" + ) + for group in range(self.nchildren) + } + nested_tree3 = { + f"group_{group}/subgroup_2/sub-subgroup_1": self.ds + for group in range(self.nchildren) + } + dtree = root | nested_tree1 | nested_tree2 | nested_tree3 + self.dtree = DataTree.from_dict(dtree) + + +class IOReadDataTreeNetCDF4(IONestedDataTree): + def setup(self): + # TODO: Lazily skipped in CI as it is very demanding and slow. + # Improve times and remove errors. + _skip_slow() + + requires_dask() + + self.make_datatree() + self.format = "NETCDF4" + self.filepath = "datatree.nc4.nc" + dtree = self.dtree + dtree.to_netcdf(filepath=self.filepath) + + def time_load_datatree_netcdf4(self): + open_datatree(self.filepath, engine="netcdf4").load() + + def time_open_datatree_netcdf4(self): + open_datatree(self.filepath, engine="netcdf4") + + class IOWriteNetCDFDask: timeout = 60 repeat = 1 diff --git a/doc/whats-new.rst b/doc/whats-new.rst index c58f73cb1fa..8c6b3a099c2 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -24,6 +24,8 @@ New Features ~~~~~~~~~~~~ - Allow chunking for arrays with duplicated dimension names (:issue:`8759`, :pull:`9099`). By `Martin Raspaud `_. +- Extract the source url from fsspec objects (:issue:`9142`, :pull:`8923`). + By `Justus Magin `_. Breaking changes ~~~~~~~~~~~~~~~~ @@ -35,11 +37,21 @@ Deprecations Bug fixes ~~~~~~~~~ +- Don't convert custom indexes to ``pandas`` indexes when computing a diff (:pull:`9157`) + By `Justus Magin `_. - Make :py:func:`testing.assert_allclose` work with numpy 2.0 (:issue:`9165`, :pull:`9166`). By `Pontus Lurcock `_. +- Allow diffing objects with array attributes on variables (:issue:`9153`, :pull:`9169`). + By `Justus Magin `_. - Promote floating-point numeric datetimes before decoding (:issue:`9179`, :pull:`9182`). By `Justus Magin `_. - +- Fiy static typing of tolerance arguments by allowing `str` type (:issue:`8892`, :pull:`9194`). + By `Michael Niklas `_. +- Dark themes are now properly detected for ``html[data-theme=dark]``-tags (:pull:`9200`). + By `Dieter Werthmüller `_. +- Reductions no longer fail for ``np.complex_`` dtype arrays when numbagg is + installed. + By `Maximilian Roos `_ Documentation ~~~~~~~~~~~~~ diff --git a/properties/test_pandas_roundtrip.py b/properties/test_pandas_roundtrip.py index 0249aa59d5b..9e0d4640171 100644 --- a/properties/test_pandas_roundtrip.py +++ b/properties/test_pandas_roundtrip.py @@ -9,7 +9,6 @@ import pytest import xarray as xr -from xarray.tests import has_pandas_3 pytest.importorskip("hypothesis") import hypothesis.extra.numpy as npst # isort:skip @@ -25,22 +24,34 @@ numeric_series = numeric_dtypes.flatmap(lambda dt: pdst.series(dtype=dt)) + +@st.composite +def dataframe_strategy(draw): + tz = draw(st.timezones()) + dtype = pd.DatetimeTZDtype(unit="ns", tz=tz) + + datetimes = st.datetimes( + min_value=pd.Timestamp("1677-09-21T00:12:43.145224193"), + max_value=pd.Timestamp("2262-04-11T23:47:16.854775807"), + timezones=st.just(tz), + ) + + df = pdst.data_frames( + [ + pdst.column("datetime_col", elements=datetimes), + pdst.column("other_col", elements=st.integers()), + ], + index=pdst.range_indexes(min_size=1, max_size=10), + ) + return draw(df).astype({"datetime_col": dtype}) + + an_array = npst.arrays( dtype=numeric_dtypes, shape=npst.array_shapes(max_dims=2), # can only convert 1D/2D to pandas ) -datetime_with_tz_strategy = st.datetimes(timezones=st.timezones()) -dataframe_strategy = pdst.data_frames( - [ - pdst.column("datetime_col", elements=datetime_with_tz_strategy), - pdst.column("other_col", elements=st.integers()), - ], - index=pdst.range_indexes(min_size=1, max_size=10), -) - - @st.composite def datasets_1d_vars(draw) -> xr.Dataset: """Generate datasets with only 1D variables @@ -111,11 +122,7 @@ def test_roundtrip_pandas_dataframe(df) -> None: xr.testing.assert_identical(arr, roundtripped.to_xarray()) -@pytest.mark.skipif( - has_pandas_3, - reason="fails to roundtrip on pandas 3 (see https://github.com/pydata/xarray/issues/9098)", -) -@given(df=dataframe_strategy) +@given(df=dataframe_strategy()) def test_roundtrip_pandas_dataframe_datetime(df) -> None: # Need to name the indexes, otherwise Xarray names them 'dim_0', 'dim_1'. df.index.name = "rows" diff --git a/properties/test_properties.py b/properties/test_properties.py new file mode 100644 index 00000000000..fc0a1955539 --- /dev/null +++ b/properties/test_properties.py @@ -0,0 +1,17 @@ +import pytest + +pytest.importorskip("hypothesis") + +from hypothesis import given + +import xarray as xr +import xarray.testing.strategies as xrst + + +@given(attrs=xrst.simple_attrs) +def test_assert_identical(attrs): + v = xr.Variable(dims=(), data=0, attrs=attrs) + xr.testing.assert_identical(v, v.copy(deep=True)) + + ds = xr.Dataset(attrs=attrs) + xr.testing.assert_identical(ds, ds.copy(deep=True)) diff --git a/xarray/backends/api.py b/xarray/backends/api.py index 7054c62126e..521bdf65e6a 100644 --- a/xarray/backends/api.py +++ b/xarray/backends/api.py @@ -382,8 +382,11 @@ def _dataset_from_backend_dataset( ds.set_close(backend_ds._close) # Ensure source filename always stored in dataset object - if "source" not in ds.encoding and isinstance(filename_or_obj, (str, os.PathLike)): - ds.encoding["source"] = _normalize_path(filename_or_obj) + if "source" not in ds.encoding: + path = getattr(filename_or_obj, "path", filename_or_obj) + + if isinstance(path, (str, os.PathLike)): + ds.encoding["source"] = _normalize_path(path) return ds diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index 9796fcbf9e2..85a1a6e214c 100644 --- a/xarray/backends/zarr.py +++ b/xarray/backends/zarr.py @@ -446,7 +446,7 @@ def open_store( stacklevel=stacklevel, zarr_version=zarr_version, ) - group_paths = [str(group / node[1:]) for node in _iter_zarr_groups(zarr_group)] + group_paths = [node for node in _iter_zarr_groups(zarr_group, parent=group)] return { group: cls( zarr_group.get(group), diff --git a/xarray/core/alignment.py b/xarray/core/alignment.py index 13e3400d170..44fc7319170 100644 --- a/xarray/core/alignment.py +++ b/xarray/core/alignment.py @@ -137,7 +137,7 @@ def __init__( exclude_dims: str | Iterable[Hashable] = frozenset(), exclude_vars: Iterable[Hashable] = frozenset(), method: str | None = None, - tolerance: int | float | Iterable[int | float] | None = None, + tolerance: float | Iterable[float] | str | None = None, copy: bool = True, fill_value: Any = dtypes.NA, sparse: bool = False, @@ -965,7 +965,7 @@ def reindex( obj: T_Alignable, indexers: Mapping[Any, Any], method: str | None = None, - tolerance: int | float | Iterable[int | float] | None = None, + tolerance: float | Iterable[float] | str | None = None, copy: bool = True, fill_value: Any = dtypes.NA, sparse: bool = False, @@ -1004,7 +1004,7 @@ def reindex_like( obj: T_Alignable, other: Dataset | DataArray, method: str | None = None, - tolerance: int | float | Iterable[int | float] | None = None, + tolerance: float | Iterable[float] | str | None = None, copy: bool = True, fill_value: Any = dtypes.NA, ) -> T_Alignable: diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py index 1ac89d9ceb9..eec1c182a0e 100644 --- a/xarray/core/dataarray.py +++ b/xarray/core/dataarray.py @@ -1910,7 +1910,7 @@ def reindex_like( other: T_DataArrayOrSet, *, method: ReindexMethodOptions = None, - tolerance: int | float | Iterable[int | float] | None = None, + tolerance: float | Iterable[float] | str | None = None, copy: bool = True, fill_value=dtypes.NA, ) -> Self: @@ -1937,7 +1937,7 @@ def reindex_like( - backfill / bfill: propagate next valid index value backward - nearest: use nearest valid index value - tolerance : optional + tolerance : float | Iterable[float] | str | None, default: None Maximum distance between original and new labels for inexact matches. The values of the index at the matching locations must satisfy the equation ``abs(index[indexer] - target) <= tolerance``. @@ -2097,7 +2097,7 @@ def reindex( indexers: Mapping[Any, Any] | None = None, *, method: ReindexMethodOptions = None, - tolerance: float | Iterable[float] | None = None, + tolerance: float | Iterable[float] | str | None = None, copy: bool = True, fill_value=dtypes.NA, **indexers_kwargs: Any, @@ -2127,7 +2127,7 @@ def reindex( - backfill / bfill: propagate next valid index value backward - nearest: use nearest valid index value - tolerance : float | Iterable[float] | None, default: None + tolerance : float | Iterable[float] | str | None, default: None Maximum distance between original and new labels for inexact matches. The values of the index at the matching locations must satisfy the equation ``abs(index[indexer] - target) <= tolerance``. diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index 209f47ba835..fc02d4f712c 100644 --- a/xarray/core/dataset.py +++ b/xarray/core/dataset.py @@ -3500,7 +3500,7 @@ def reindex_like( self, other: T_Xarray, method: ReindexMethodOptions = None, - tolerance: int | float | Iterable[int | float] | None = None, + tolerance: float | Iterable[float] | str | None = None, copy: bool = True, fill_value: Any = xrdtypes.NA, ) -> Self: @@ -3527,7 +3527,7 @@ def reindex_like( - "backfill" / "bfill": propagate next valid index value backward - "nearest": use nearest valid index value - tolerance : optional + tolerance : float | Iterable[float] | str | None, default: None Maximum distance between original and new labels for inexact matches. The values of the index at the matching locations must satisfy the equation ``abs(index[indexer] - target) <= tolerance``. @@ -3570,7 +3570,7 @@ def reindex( self, indexers: Mapping[Any, Any] | None = None, method: ReindexMethodOptions = None, - tolerance: int | float | Iterable[int | float] | None = None, + tolerance: float | Iterable[float] | str | None = None, copy: bool = True, fill_value: Any = xrdtypes.NA, **indexers_kwargs: Any, @@ -3595,7 +3595,7 @@ def reindex( - "backfill" / "bfill": propagate next valid index value backward - "nearest": use nearest valid index value - tolerance : optional + tolerance : float | Iterable[float] | str | None, default: None Maximum distance between original and new labels for inexact matches. The values of the index at the matching locations must satisfy the equation ``abs(index[indexer] - target) <= tolerance``. diff --git a/xarray/core/datatree.py b/xarray/core/datatree.py index c923ca2eb87..38f8f8cd495 100644 --- a/xarray/core/datatree.py +++ b/xarray/core/datatree.py @@ -1,7 +1,8 @@ from __future__ import annotations -import copy import itertools +import textwrap +from collections import ChainMap from collections.abc import Hashable, Iterable, Iterator, Mapping, MutableMapping from html import escape from typing import ( @@ -16,6 +17,7 @@ ) from xarray.core import utils +from xarray.core.alignment import align from xarray.core.common import TreeAttrAccessMixin from xarray.core.coordinates import DatasetCoordinates from xarray.core.dataarray import DataArray @@ -31,7 +33,7 @@ MappedDataWithCoords, ) from xarray.core.datatree_render import RenderDataTree -from xarray.core.formatting import datatree_repr +from xarray.core.formatting import datatree_repr, dims_and_coords_repr from xarray.core.formatting_html import ( datatree_repr as datatree_repr_html, ) @@ -79,11 +81,24 @@ T_Path = Union[str, NodePath] +def _collect_data_and_coord_variables( + data: Dataset, +) -> tuple[dict[Hashable, Variable], dict[Hashable, Variable]]: + data_variables = {} + coord_variables = {} + for k, v in data.variables.items(): + if k in data._coord_names: + coord_variables[k] = v + else: + data_variables[k] = v + return data_variables, coord_variables + + def _coerce_to_dataset(data: Dataset | DataArray | None) -> Dataset: if isinstance(data, DataArray): ds = data.to_dataset() elif isinstance(data, Dataset): - ds = data + ds = data.copy(deep=False) elif data is None: ds = Dataset() else: @@ -93,14 +108,57 @@ def _coerce_to_dataset(data: Dataset | DataArray | None) -> Dataset: return ds -def _check_for_name_collisions( - children: Iterable[str], variables: Iterable[Hashable] +def _join_path(root: str, name: str) -> str: + return str(NodePath(root) / name) + + +def _inherited_dataset(ds: Dataset, parent: Dataset) -> Dataset: + return Dataset._construct_direct( + variables=parent._variables | ds._variables, + coord_names=parent._coord_names | ds._coord_names, + dims=parent._dims | ds._dims, + attrs=ds._attrs, + indexes=parent._indexes | ds._indexes, + encoding=ds._encoding, + close=ds._close, + ) + + +def _without_header(text: str) -> str: + return "\n".join(text.split("\n")[1:]) + + +def _indented(text: str) -> str: + return textwrap.indent(text, prefix=" ") + + +def _check_alignment( + path: str, + node_ds: Dataset, + parent_ds: Dataset | None, + children: Mapping[str, DataTree], ) -> None: - colliding_names = set(children).intersection(set(variables)) - if colliding_names: - raise KeyError( - f"Some names would collide between variables and children: {list(colliding_names)}" - ) + if parent_ds is not None: + try: + align(node_ds, parent_ds, join="exact") + except ValueError as e: + node_repr = _indented(_without_header(repr(node_ds))) + parent_repr = _indented(dims_and_coords_repr(parent_ds)) + raise ValueError( + f"group {path!r} is not aligned with its parents:\n" + f"Group:\n{node_repr}\nFrom parents:\n{parent_repr}" + ) from e + + if children: + if parent_ds is not None: + base_ds = _inherited_dataset(node_ds, parent_ds) + else: + base_ds = node_ds + + for child_name, child in children.items(): + child_path = str(NodePath(path) / child_name) + child_ds = child.to_dataset(inherited=False) + _check_alignment(child_path, child_ds, base_ds, child.children) class DatasetView(Dataset): @@ -118,7 +176,7 @@ class DatasetView(Dataset): __slots__ = ( "_attrs", - "_cache", + "_cache", # used by _CachedAccessor "_coord_names", "_dims", "_encoding", @@ -136,21 +194,27 @@ def __init__( raise AttributeError("DatasetView objects are not to be initialized directly") @classmethod - def _from_node( + def _constructor( cls, - wrapping_node: DataTree, + variables: dict[Any, Variable], + coord_names: set[Hashable], + dims: dict[Any, int], + attrs: dict | None, + indexes: dict[Any, Index], + encoding: dict | None, + close: Callable[[], None] | None, ) -> DatasetView: - """Constructor, using dataset attributes from wrapping node""" - + """Private constructor, from Dataset attributes.""" + # We override Dataset._construct_direct below, so we need a new + # constructor for creating DatasetView objects. obj: DatasetView = object.__new__(cls) - obj._variables = wrapping_node._variables - obj._coord_names = wrapping_node._coord_names - obj._dims = wrapping_node._dims - obj._indexes = wrapping_node._indexes - obj._attrs = wrapping_node._attrs - obj._close = wrapping_node._close - obj._encoding = wrapping_node._encoding - + obj._variables = variables + obj._coord_names = coord_names + obj._dims = dims + obj._indexes = indexes + obj._attrs = attrs + obj._close = close + obj._encoding = encoding return obj def __setitem__(self, key, val) -> None: @@ -337,27 +401,27 @@ class DataTree( _name: str | None _parent: DataTree | None _children: dict[str, DataTree] + _cache: dict[str, Any] # used by _CachedAccessor + _data_variables: dict[Hashable, Variable] + _node_coord_variables: dict[Hashable, Variable] + _node_dims: dict[Hashable, int] + _node_indexes: dict[Hashable, Index] _attrs: dict[Hashable, Any] | None - _cache: dict[str, Any] - _coord_names: set[Hashable] - _dims: dict[Hashable, int] _encoding: dict[Hashable, Any] | None _close: Callable[[], None] | None - _indexes: dict[Hashable, Index] - _variables: dict[Hashable, Variable] __slots__ = ( "_name", "_parent", "_children", + "_cache", # used by _CachedAccessor + "_data_variables", + "_node_coord_variables", + "_node_dims", + "_node_indexes", "_attrs", - "_cache", - "_coord_names", - "_dims", "_encoding", "_close", - "_indexes", - "_variables", ) def __init__( @@ -370,14 +434,15 @@ def __init__( """ Create a single node of a DataTree. - The node may optionally contain data in the form of data and coordinate variables, stored in the same way as - data is stored in an xarray.Dataset. + The node may optionally contain data in the form of data and coordinate + variables, stored in the same way as data is stored in an + xarray.Dataset. Parameters ---------- data : Dataset, DataArray, or None, optional - Data to store under the .ds attribute of this node. DataArrays will be promoted to Datasets. - Default is None. + Data to store under the .ds attribute of this node. DataArrays will + be promoted to Datasets. Default is None. parent : DataTree, optional Parent node to this node. Default is None. children : Mapping[str, DataTree], optional @@ -393,30 +458,48 @@ def __init__( -------- DataTree.from_dict """ - - # validate input if children is None: children = {} - ds = _coerce_to_dataset(data) - _check_for_name_collisions(children, ds.variables) super().__init__(name=name) + self._set_node_data(_coerce_to_dataset(data)) + self.parent = parent + self.children = children - # set data attributes - self._replace( - inplace=True, - variables=ds._variables, - coord_names=ds._coord_names, - dims=ds._dims, - indexes=ds._indexes, - attrs=ds._attrs, - encoding=ds._encoding, - ) + def _set_node_data(self, ds: Dataset): + data_vars, coord_vars = _collect_data_and_coord_variables(ds) + self._data_variables = data_vars + self._node_coord_variables = coord_vars + self._node_dims = ds._dims + self._node_indexes = ds._indexes + self._encoding = ds._encoding + self._attrs = ds._attrs self._close = ds._close - # set tree attributes (must happen after variables set to avoid initialization errors) - self.children = children - self.parent = parent + def _pre_attach(self: DataTree, parent: DataTree, name: str) -> None: + super()._pre_attach(parent, name) + if name in parent.ds.variables: + raise KeyError( + f"parent {parent.name} already contains a variable named {name}" + ) + path = str(NodePath(parent.path) / name) + node_ds = self.to_dataset(inherited=False) + parent_ds = parent._to_dataset_view(rebuild_dims=False) + _check_alignment(path, node_ds, parent_ds, self.children) + + @property + def _coord_variables(self) -> ChainMap[Hashable, Variable]: + return ChainMap( + self._node_coord_variables, *(p._node_coord_variables for p in self.parents) + ) + + @property + def _dims(self) -> ChainMap[Hashable, int]: + return ChainMap(self._node_dims, *(p._node_dims for p in self.parents)) + + @property + def _indexes(self) -> ChainMap[Hashable, Index]: + return ChainMap(self._node_indexes, *(p._node_indexes for p in self.parents)) @property def parent(self: DataTree) -> DataTree | None: @@ -429,71 +512,87 @@ def parent(self: DataTree, new_parent: DataTree) -> None: raise ValueError("Cannot set an unnamed node as a child of another node") self._set_parent(new_parent, self.name) + def _to_dataset_view(self, rebuild_dims: bool) -> DatasetView: + variables = dict(self._data_variables) + variables |= self._coord_variables + if rebuild_dims: + dims = calculate_dimensions(variables) + else: + # Note: rebuild_dims=False can create technically invalid Dataset + # objects because it may not contain all dimensions on its direct + # member variables, e.g., consider: + # tree = DataTree.from_dict( + # { + # "/": xr.Dataset({"a": (("x",), [1, 2])}), # x has size 2 + # "/b/c": xr.Dataset({"d": (("x",), [3])}), # x has size1 + # } + # ) + # However, they are fine for internal use cases, for align() or + # building a repr(). + dims = dict(self._dims) + return DatasetView._constructor( + variables=variables, + coord_names=set(self._coord_variables), + dims=dims, + attrs=self._attrs, + indexes=dict(self._indexes), + encoding=self._encoding, + close=None, + ) + @property def ds(self) -> DatasetView: """ An immutable Dataset-like view onto the data in this node. - For a mutable Dataset containing the same data as in this node, use `.to_dataset()` instead. + Includes inherited coordinates and indexes from parent nodes. + + For a mutable Dataset containing the same data as in this node, use + `.to_dataset()` instead. See Also -------- DataTree.to_dataset """ - return DatasetView._from_node(self) + return self._to_dataset_view(rebuild_dims=True) @ds.setter def ds(self, data: Dataset | DataArray | None = None) -> None: - # Known mypy issue for setters with different type to property: - # https://github.com/python/mypy/issues/3004 ds = _coerce_to_dataset(data) + self._replace_node(ds) - _check_for_name_collisions(self.children, ds.variables) - - self._replace( - inplace=True, - variables=ds._variables, - coord_names=ds._coord_names, - dims=ds._dims, - indexes=ds._indexes, - attrs=ds._attrs, - encoding=ds._encoding, - ) - self._close = ds._close - - def _pre_attach(self: DataTree, parent: DataTree) -> None: - """ - Method which superclass calls before setting parent, here used to prevent having two - children with duplicate names (or a data variable with the same name as a child). - """ - super()._pre_attach(parent) - if self.name in list(parent.ds.variables): - raise KeyError( - f"parent {parent.name} already contains a data variable named {self.name}" - ) - - def to_dataset(self) -> Dataset: + def to_dataset(self, inherited: bool = True) -> Dataset: """ Return the data in this node as a new xarray.Dataset object. + Parameters + ---------- + inherited : bool, optional + If False, only include coordinates and indexes defined at the level + of this DataTree node, excluding inherited coordinates. + See Also -------- DataTree.ds """ + coord_vars = self._coord_variables if inherited else self._node_coord_variables + variables = dict(self._data_variables) + variables |= coord_vars + dims = calculate_dimensions(variables) if inherited else dict(self._node_dims) return Dataset._construct_direct( - self._variables, - self._coord_names, - self._dims, - self._attrs, - self._indexes, - self._encoding, + variables, + set(coord_vars), + dims, + None if self._attrs is None else dict(self._attrs), + dict(self._indexes if inherited else self._node_indexes), + None if self._encoding is None else dict(self._encoding), self._close, ) @property - def has_data(self): - """Whether or not there are any data variables in this node.""" - return len(self._variables) > 0 + def has_data(self) -> bool: + """Whether or not there are any variables in this node.""" + return bool(self._data_variables or self._node_coord_variables) @property def has_attrs(self) -> bool: @@ -518,7 +617,7 @@ def variables(self) -> Mapping[Hashable, Variable]: Dataset invariants. It contains all variable objects constituting this DataTree node, including both data variables and coordinates. """ - return Frozen(self._variables) + return Frozen(self._data_variables | self._coord_variables) @property def attrs(self) -> dict[Hashable, Any]: @@ -579,7 +678,7 @@ def _attr_sources(self) -> Iterable[Mapping[Hashable, Any]]: def _item_sources(self) -> Iterable[Mapping[Any, Any]]: """Places to look-up items for key-completion""" yield self.data_vars - yield HybridMappingProxy(keys=self._coord_names, mapping=self.coords) + yield HybridMappingProxy(keys=self._coord_variables, mapping=self.coords) # virtual coordinates yield HybridMappingProxy(keys=self.dims, mapping=self) @@ -621,10 +720,10 @@ def __contains__(self, key: object) -> bool: return key in self.variables or key in self.children def __bool__(self) -> bool: - return bool(self.ds.data_vars) or bool(self.children) + return bool(self._data_variables) or bool(self._children) def __iter__(self) -> Iterator[Hashable]: - return itertools.chain(self.ds.data_vars, self.children) + return itertools.chain(self._data_variables, self._children) def __array__(self, dtype=None, copy=None): raise TypeError( @@ -646,122 +745,32 @@ def _repr_html_(self): return f"
{escape(repr(self))}
" return datatree_repr_html(self) - @classmethod - def _construct_direct( - cls, - variables: dict[Any, Variable], - coord_names: set[Hashable], - dims: dict[Any, int] | None = None, - attrs: dict | None = None, - indexes: dict[Any, Index] | None = None, - encoding: dict | None = None, - name: str | None = None, - parent: DataTree | None = None, - children: dict[str, DataTree] | None = None, - close: Callable[[], None] | None = None, - ) -> DataTree: - """Shortcut around __init__ for internal use when we want to skip costly validation.""" + def _replace_node( + self: DataTree, + data: Dataset | Default = _default, + children: dict[str, DataTree] | Default = _default, + ) -> None: - # data attributes - if dims is None: - dims = calculate_dimensions(variables) - if indexes is None: - indexes = {} - if children is None: - children = dict() + ds = self.to_dataset(inherited=False) if data is _default else data - obj: DataTree = object.__new__(cls) - obj._variables = variables - obj._coord_names = coord_names - obj._dims = dims - obj._indexes = indexes - obj._attrs = attrs - obj._close = close - obj._encoding = encoding - - # tree attributes - obj._name = name - obj._children = children - obj._parent = parent + if children is _default: + children = self._children - return obj + for child_name in children: + if child_name in ds.variables: + raise ValueError(f"node already contains a variable named {child_name}") - def _replace( - self: DataTree, - variables: dict[Hashable, Variable] | None = None, - coord_names: set[Hashable] | None = None, - dims: dict[Any, int] | None = None, - attrs: dict[Hashable, Any] | None | Default = _default, - indexes: dict[Hashable, Index] | None = None, - encoding: dict | None | Default = _default, - name: str | None | Default = _default, - parent: DataTree | None | Default = _default, - children: dict[str, DataTree] | None = None, - inplace: bool = False, - ) -> DataTree: - """ - Fastpath constructor for internal use. + parent_ds = ( + self.parent._to_dataset_view(rebuild_dims=False) + if self.parent is not None + else None + ) + _check_alignment(self.path, ds, parent_ds, children) - Returns an object with optionally replaced attributes. + if data is not _default: + self._set_node_data(ds) - Explicitly passed arguments are *not* copied when placed on the new - datatree. It is up to the caller to ensure that they have the right type - and are not used elsewhere. - """ - # TODO Adding new children inplace using this method will cause bugs. - # You will end up with an inconsistency between the name of the child node and the key the child is stored under. - # Use ._set() instead for now - if inplace: - if variables is not None: - self._variables = variables - if coord_names is not None: - self._coord_names = coord_names - if dims is not None: - self._dims = dims - if attrs is not _default: - self._attrs = attrs - if indexes is not None: - self._indexes = indexes - if encoding is not _default: - self._encoding = encoding - if name is not _default: - self._name = name - if parent is not _default: - self._parent = parent - if children is not None: - self._children = children - obj = self - else: - if variables is None: - variables = self._variables.copy() - if coord_names is None: - coord_names = self._coord_names.copy() - if dims is None: - dims = self._dims.copy() - if attrs is _default: - attrs = copy.copy(self._attrs) - if indexes is None: - indexes = self._indexes.copy() - if encoding is _default: - encoding = copy.copy(self._encoding) - if name is _default: - name = self._name # no need to copy str objects or None - if parent is _default: - parent = copy.copy(self._parent) - if children is _default: - children = copy.copy(self._children) - obj = self._construct_direct( - variables, - coord_names, - dims, - attrs, - indexes, - encoding, - name, - parent, - children, - ) - return obj + self._children = children def copy( self: DataTree, @@ -813,9 +822,8 @@ def _copy_node( deep: bool = False, ) -> DataTree: """Copy just one node of a tree""" - new_node: DataTree = DataTree() - new_node.name = self.name - new_node.ds = self.to_dataset().copy(deep=deep) # type: ignore[assignment] + data = self.ds.copy(deep=deep) + new_node: DataTree = DataTree(data, name=self.name) return new_node def __copy__(self: DataTree) -> DataTree: @@ -963,11 +971,12 @@ def update( raise TypeError(f"Type {type(v)} cannot be assigned to a DataTree") vars_merge_result = dataset_update_method(self.to_dataset(), new_variables) + data = Dataset._construct_direct(**vars_merge_result._asdict()) + # TODO are there any subtleties with preserving order of children like this? merged_children = {**self.children, **new_children} - self._replace( - inplace=True, children=merged_children, **vars_merge_result._asdict() - ) + + self._replace_node(data, children=merged_children) def assign( self, items: Mapping[Any, Any] | None = None, **items_kwargs: Any @@ -1042,10 +1051,12 @@ def drop_nodes( if extra: raise KeyError(f"Cannot drop all nodes - nodes {extra} not present") + result = self.copy() children_to_keep = { - name: child for name, child in self.children.items() if name not in names + name: child for name, child in result.children.items() if name not in names } - return self._replace(children=children_to_keep) + result._replace_node(children=children_to_keep) + return result @classmethod def from_dict( @@ -1137,7 +1148,9 @@ def indexes(self) -> Indexes[pd.Index]: @property def xindexes(self) -> Indexes[Index]: """Mapping of xarray Index objects used for label based indexing.""" - return Indexes(self._indexes, {k: self._variables[k] for k in self._indexes}) + return Indexes( + self._indexes, {k: self._coord_variables[k] for k in self._indexes} + ) @property def coords(self) -> DatasetCoordinates: diff --git a/xarray/core/datatree_io.py b/xarray/core/datatree_io.py index 1473e624d9e..36665a0d153 100644 --- a/xarray/core/datatree_io.py +++ b/xarray/core/datatree_io.py @@ -85,7 +85,7 @@ def _datatree_to_netcdf( unlimited_dims = {} for node in dt.subtree: - ds = node.ds + ds = node.to_dataset(inherited=False) group_path = node.path if ds is None: _create_empty_netcdf_group(filepath, group_path, mode, engine) @@ -151,7 +151,7 @@ def _datatree_to_zarr( ) for node in dt.subtree: - ds = node.ds + ds = node.to_dataset(inherited=False) group_path = node.path if ds is None: _create_empty_zarr_group(store, group_path, mode) diff --git a/xarray/core/formatting.py b/xarray/core/formatting.py index c15df34b5b1..6571b288fae 100644 --- a/xarray/core/formatting.py +++ b/xarray/core/formatting.py @@ -748,6 +748,27 @@ def dataset_repr(ds): return "\n".join(summary) +def dims_and_coords_repr(ds) -> str: + """Partial Dataset repr for use inside DataTree inheritance errors.""" + summary = [] + + col_width = _calculate_col_width(ds.coords) + max_rows = OPTIONS["display_max_rows"] + + dims_start = pretty_print("Dimensions:", col_width) + dims_values = dim_summary_limited(ds, col_width=col_width + 1, max_rows=max_rows) + summary.append(f"{dims_start}({dims_values})") + + if ds.coords: + summary.append(coords_repr(ds.coords, col_width=col_width, max_rows=max_rows)) + + unindexed_dims_str = unindexed_dims_repr(ds.dims, ds.coords, max_rows=max_rows) + if unindexed_dims_str: + summary.append(unindexed_dims_str) + + return "\n".join(summary) + + def diff_dim_summary(a, b): if a.sizes != b.sizes: return f"Differing dimensions:\n ({dim_summary(a)}) != ({dim_summary(b)})" @@ -765,6 +786,12 @@ def _diff_mapping_repr( a_indexes=None, b_indexes=None, ): + def compare_attr(a, b): + if is_duck_array(a) or is_duck_array(b): + return array_equiv(a, b) + else: + return a == b + def extra_items_repr(extra_keys, mapping, ab_side, kwargs): extra_repr = [ summarizer(k, mapping[k], col_width, **kwargs[k]) for k in extra_keys @@ -801,11 +828,7 @@ def extra_items_repr(extra_keys, mapping, ab_side, kwargs): is_variable = True except AttributeError: # compare attribute value - if is_duck_array(a_mapping[k]) or is_duck_array(b_mapping[k]): - compatible = array_equiv(a_mapping[k], b_mapping[k]) - else: - compatible = a_mapping[k] == b_mapping[k] - + compatible = compare_attr(a_mapping[k], b_mapping[k]) is_variable = False if not compatible: @@ -821,7 +844,11 @@ def extra_items_repr(extra_keys, mapping, ab_side, kwargs): attrs_to_print = set(a_attrs) ^ set(b_attrs) attrs_to_print.update( - {k for k in set(a_attrs) & set(b_attrs) if a_attrs[k] != b_attrs[k]} + { + k + for k in set(a_attrs) & set(b_attrs) + if not compare_attr(a_attrs[k], b_attrs[k]) + } ) for m in (a_mapping, b_mapping): attr_s = "\n".join( @@ -871,8 +898,8 @@ def diff_coords_repr(a, b, compat, col_width=None): "Coordinates", summarize_variable, col_width=col_width, - a_indexes=a.indexes, - b_indexes=b.indexes, + a_indexes=a.xindexes, + b_indexes=b.xindexes, ) @@ -1024,7 +1051,7 @@ def diff_datatree_repr(a: DataTree, b: DataTree, compat): def _single_node_repr(node: DataTree) -> str: """Information about this node, not including its relationships to other nodes.""" if node.has_data or node.has_attrs: - ds_info = "\n" + repr(node.ds) + ds_info = "\n" + repr(node._to_dataset_view(rebuild_dims=False)) else: ds_info = "" return f"Group: {node.path}{ds_info}" diff --git a/xarray/core/formatting_html.py b/xarray/core/formatting_html.py index 9bf5befbe3f..24b290031eb 100644 --- a/xarray/core/formatting_html.py +++ b/xarray/core/formatting_html.py @@ -386,7 +386,7 @@ def summarize_datatree_children(children: Mapping[str, DataTree]) -> str: def datatree_node_repr(group_title: str, dt: DataTree) -> str: header_components = [f"
{escape(group_title)}
"] - ds = dt.ds + ds = dt._to_dataset_view(rebuild_dims=False) sections = [ children_section(dt.children), diff --git a/xarray/core/nputils.py b/xarray/core/nputils.py index 6970d37402f..1d30fe9db9e 100644 --- a/xarray/core/nputils.py +++ b/xarray/core/nputils.py @@ -191,7 +191,7 @@ def f(values, axis=None, **kwargs): or kwargs.get("ddof", 0) == 1 ) # TODO: bool? - and values.dtype.kind in "uifc" + and values.dtype.kind in "uif" # and values.dtype.isnative and (dtype is None or np.dtype(dtype) == values.dtype) # numbagg.nanquantile only available after 0.8.0 and with linear method diff --git a/xarray/core/resample.py b/xarray/core/resample.py index ceab0a891c9..ec86f2a283f 100644 --- a/xarray/core/resample.py +++ b/xarray/core/resample.py @@ -66,12 +66,12 @@ def _drop_coords(self) -> T_Xarray: obj = obj.drop_vars([k]) return obj - def pad(self, tolerance: float | Iterable[float] | None = None) -> T_Xarray: + def pad(self, tolerance: float | Iterable[float] | str | None = None) -> T_Xarray: """Forward fill new values at up-sampled frequency. Parameters ---------- - tolerance : float | Iterable[float] | None, default: None + tolerance : float | Iterable[float] | str | None, default: None Maximum distance between original and new labels to limit the up-sampling method. Up-sampled data with indices that satisfy the equation @@ -91,12 +91,14 @@ def pad(self, tolerance: float | Iterable[float] | None = None) -> T_Xarray: ffill = pad - def backfill(self, tolerance: float | Iterable[float] | None = None) -> T_Xarray: + def backfill( + self, tolerance: float | Iterable[float] | str | None = None + ) -> T_Xarray: """Backward fill new values at up-sampled frequency. Parameters ---------- - tolerance : float | Iterable[float] | None, default: None + tolerance : float | Iterable[float] | str | None, default: None Maximum distance between original and new labels to limit the up-sampling method. Up-sampled data with indices that satisfy the equation @@ -116,13 +118,15 @@ def backfill(self, tolerance: float | Iterable[float] | None = None) -> T_Xarray bfill = backfill - def nearest(self, tolerance: float | Iterable[float] | None = None) -> T_Xarray: + def nearest( + self, tolerance: float | Iterable[float] | str | None = None + ) -> T_Xarray: """Take new values from nearest original coordinate to up-sampled frequency coordinates. Parameters ---------- - tolerance : float | Iterable[float] | None, default: None + tolerance : float | Iterable[float] | str | None, default: None Maximum distance between original and new labels to limit the up-sampling method. Up-sampled data with indices that satisfy the equation diff --git a/xarray/core/treenode.py b/xarray/core/treenode.py index 6f51e1ffa38..77e7ed23a51 100644 --- a/xarray/core/treenode.py +++ b/xarray/core/treenode.py @@ -138,14 +138,14 @@ def _attach(self, parent: Tree | None, child_name: str | None = None) -> None: "To directly set parent, child needs a name, but child is unnamed" ) - self._pre_attach(parent) + self._pre_attach(parent, child_name) parentchildren = parent._children assert not any( child is self for child in parentchildren ), "Tree is corrupt." parentchildren[child_name] = self self._parent = parent - self._post_attach(parent) + self._post_attach(parent, child_name) else: self._parent = None @@ -415,11 +415,11 @@ def _post_detach(self: Tree, parent: Tree) -> None: """Method call after detaching from `parent`.""" pass - def _pre_attach(self: Tree, parent: Tree) -> None: + def _pre_attach(self: Tree, parent: Tree, name: str) -> None: """Method call before attaching to `parent`.""" pass - def _post_attach(self: Tree, parent: Tree) -> None: + def _post_attach(self: Tree, parent: Tree, name: str) -> None: """Method call after attaching to `parent`.""" pass @@ -567,6 +567,9 @@ def same_tree(self, other: Tree) -> bool: return self.root is other.root +AnyNamedNode = TypeVar("AnyNamedNode", bound="NamedNode") + + class NamedNode(TreeNode, Generic[Tree]): """ A TreeNode which knows its own name. @@ -606,10 +609,9 @@ def __repr__(self, level=0): def __str__(self) -> str: return f"NamedNode('{self.name}')" if self.name else "NamedNode()" - def _post_attach(self: NamedNode, parent: NamedNode) -> None: + def _post_attach(self: AnyNamedNode, parent: AnyNamedNode, name: str) -> None: """Ensures child has name attribute corresponding to key under which it has been stored.""" - key = next(k for k, v in parent.children.items() if v is self) - self.name = key + self.name = name @property def path(self) -> str: diff --git a/xarray/static/css/style.css b/xarray/static/css/style.css index e0a51312b10..dbe61e311c1 100644 --- a/xarray/static/css/style.css +++ b/xarray/static/css/style.css @@ -14,6 +14,7 @@ } html[theme=dark], +html[data-theme=dark], body[data-theme=dark], body.vscode-dark { --xr-font-color0: rgba(255, 255, 255, 1); diff --git a/xarray/testing/strategies.py b/xarray/testing/strategies.py index 449d0c793cc..085b70e518b 100644 --- a/xarray/testing/strategies.py +++ b/xarray/testing/strategies.py @@ -192,10 +192,14 @@ def dimension_sizes( max_side=2, max_dims=2, ), - dtype=npst.scalar_dtypes(), + dtype=npst.scalar_dtypes() + | npst.byte_string_dtypes() + | npst.unicode_string_dtypes(), ) _attr_values = st.none() | st.booleans() | _readable_strings | _small_arrays +simple_attrs = st.dictionaries(_attr_keys, _attr_values) + def attrs() -> st.SearchStrategy[Mapping[Hashable, Any]]: """ diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index 177700a5404..15485dc178a 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -5151,6 +5151,21 @@ def test_source_encoding_always_present_with_pathlib() -> None: assert ds.encoding["source"] == tmp +@requires_h5netcdf +@requires_fsspec +def test_source_encoding_always_present_with_fsspec() -> None: + import fsspec + + rnddata = np.random.randn(10) + original = Dataset({"foo": ("x", rnddata)}) + with create_tmp_file() as tmp: + original.to_netcdf(tmp) + + fs = fsspec.filesystem("file") + with fs.open(tmp) as f, open_dataset(f) as ds: + assert ds.encoding["source"] == tmp + + def _assert_no_dates_out_of_range_warning(record): undesired_message = "dates out of range" for warning in record: diff --git a/xarray/tests/test_backends_datatree.py b/xarray/tests/test_backends_datatree.py index 4e819eec0b5..b4c4f481359 100644 --- a/xarray/tests/test_backends_datatree.py +++ b/xarray/tests/test_backends_datatree.py @@ -1,10 +1,12 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, cast import pytest +import xarray as xr from xarray.backends.api import open_datatree +from xarray.core.datatree import DataTree from xarray.testing import assert_equal from xarray.tests import ( requires_h5netcdf, @@ -13,11 +15,11 @@ ) if TYPE_CHECKING: - from xarray.backends.api import T_NetcdfEngine + from xarray.core.datatree_io import T_DataTreeNetcdfEngine class DatatreeIOBase: - engine: T_NetcdfEngine | None = None + engine: T_DataTreeNetcdfEngine | None = None def test_to_netcdf(self, tmpdir, simple_datatree): filepath = tmpdir / "test.nc" @@ -27,6 +29,21 @@ def test_to_netcdf(self, tmpdir, simple_datatree): roundtrip_dt = open_datatree(filepath, engine=self.engine) assert_equal(original_dt, roundtrip_dt) + def test_to_netcdf_inherited_coords(self, tmpdir): + filepath = tmpdir / "test.nc" + original_dt = DataTree.from_dict( + { + "/": xr.Dataset({"a": (("x",), [1, 2])}, coords={"x": [3, 4]}), + "/sub": xr.Dataset({"b": (("x",), [5, 6])}), + } + ) + original_dt.to_netcdf(filepath, engine=self.engine) + + roundtrip_dt = open_datatree(filepath, engine=self.engine) + assert_equal(original_dt, roundtrip_dt) + subtree = cast(DataTree, roundtrip_dt["/sub"]) + assert "x" not in subtree.to_dataset(inherited=False).coords + def test_netcdf_encoding(self, tmpdir, simple_datatree): filepath = tmpdir / "test.nc" original_dt = simple_datatree @@ -48,12 +65,12 @@ def test_netcdf_encoding(self, tmpdir, simple_datatree): @requires_netCDF4 class TestNetCDF4DatatreeIO(DatatreeIOBase): - engine: T_NetcdfEngine | None = "netcdf4" + engine: T_DataTreeNetcdfEngine | None = "netcdf4" @requires_h5netcdf class TestH5NetCDFDatatreeIO(DatatreeIOBase): - engine: T_NetcdfEngine | None = "h5netcdf" + engine: T_DataTreeNetcdfEngine | None = "h5netcdf" @requires_zarr @@ -119,3 +136,18 @@ def test_to_zarr_default_write_mode(self, tmpdir, simple_datatree): # with default settings, to_zarr should not overwrite an existing dir with pytest.raises(zarr.errors.ContainsGroupError): simple_datatree.to_zarr(tmpdir) + + def test_to_zarr_inherited_coords(self, tmpdir): + original_dt = DataTree.from_dict( + { + "/": xr.Dataset({"a": (("x",), [1, 2])}, coords={"x": [3, 4]}), + "/sub": xr.Dataset({"b": (("x",), [5, 6])}), + } + ) + filepath = tmpdir / "test.zarr" + original_dt.to_zarr(filepath) + + roundtrip_dt = open_datatree(filepath, engine="zarr") + assert_equal(original_dt, roundtrip_dt) + subtree = cast(DataTree, roundtrip_dt["/sub"]) + assert "x" not in subtree.to_dataset(inherited=False).coords diff --git a/xarray/tests/test_computation.py b/xarray/tests/test_computation.py index 080447cede0..4b9b95b27bb 100644 --- a/xarray/tests/test_computation.py +++ b/xarray/tests/test_computation.py @@ -2598,3 +2598,11 @@ def test_cross(a, b, ae, be, dim: str, axis: int, use_dask: bool) -> None: actual = xr.cross(a, b, dim=dim) xr.testing.assert_duckarray_allclose(expected, actual) + + +@pytest.mark.parametrize("compute_backend", ["numbagg"], indirect=True) +def test_complex_number_reduce(compute_backend): + da = xr.DataArray(np.ones((2,), dtype=np.complex_), dims=["x"]) + # Check that xarray doesn't call into numbagg, which doesn't compile for complex + # numbers at the moment (but will when numba supports dynamic compilation) + da.min() diff --git a/xarray/tests/test_datatree.py b/xarray/tests/test_datatree.py index b0dc2accd3e..f2b58fa2489 100644 --- a/xarray/tests/test_datatree.py +++ b/xarray/tests/test_datatree.py @@ -1,3 +1,5 @@ +import re +import typing from copy import copy, deepcopy from textwrap import dedent @@ -149,22 +151,37 @@ def test_is_hollow(self): assert not eve.is_hollow +class TestToDataset: + def test_to_dataset(self): + base = xr.Dataset(coords={"a": 1}) + sub = xr.Dataset(coords={"b": 2}) + tree = DataTree.from_dict({"/": base, "/sub": sub}) + subtree = typing.cast(DataTree, tree["sub"]) + + assert_identical(tree.to_dataset(inherited=False), base) + assert_identical(subtree.to_dataset(inherited=False), sub) + + sub_and_base = xr.Dataset(coords={"a": 1, "b": 2}) + assert_identical(tree.to_dataset(inherited=True), base) + assert_identical(subtree.to_dataset(inherited=True), sub_and_base) + + class TestVariablesChildrenNameCollisions: def test_parent_already_has_variable_with_childs_name(self): dt: DataTree = DataTree(data=xr.Dataset({"a": [0], "b": 1})) - with pytest.raises(KeyError, match="already contains a data variable named a"): + with pytest.raises(KeyError, match="already contains a variable named a"): DataTree(name="a", data=None, parent=dt) def test_assign_when_already_child_with_variables_name(self): dt: DataTree = DataTree(data=None) DataTree(name="a", data=None, parent=dt) - with pytest.raises(KeyError, match="names would collide"): + with pytest.raises(ValueError, match="node already contains a variable"): dt.ds = xr.Dataset({"a": 0}) # type: ignore[assignment] dt.ds = xr.Dataset() # type: ignore[assignment] new_ds = dt.to_dataset().assign(a=xr.DataArray(0)) - with pytest.raises(KeyError, match="names would collide"): + with pytest.raises(ValueError, match="node already contains a variable"): dt.ds = new_ds # type: ignore[assignment] @@ -561,7 +578,9 @@ def test_methods(self): def test_arithmetic(self, create_test_datatree): dt = create_test_datatree() - expected = create_test_datatree(modify=lambda ds: 10.0 * ds)["set1"] + expected = create_test_datatree(modify=lambda ds: 10.0 * ds)[ + "set1" + ].to_dataset() result = 10.0 * dt["set1"].ds assert result.identical(expected) @@ -648,13 +667,18 @@ def test_repr(self): │ Data variables: │ e (x) float64 16B 1.0 2.0 └── Group: /b - │ Dimensions: (y: 1) + │ Dimensions: (x: 2, y: 1) + │ Coordinates: + │ * x (x) float64 16B 2.0 3.0 │ Dimensions without coordinates: y │ Data variables: │ f (y) float64 8B 3.0 ├── Group: /b/c └── Group: /b/d - Dimensions: () + Dimensions: (x: 2, y: 1) + Coordinates: + * x (x) float64 16B 2.0 3.0 + Dimensions without coordinates: y Data variables: g float64 8B 4.0 """ @@ -666,13 +690,18 @@ def test_repr(self): """ Group: /b - │ Dimensions: (y: 1) + │ Dimensions: (x: 2, y: 1) + │ Coordinates: + │ * x (x) float64 16B 2.0 3.0 │ Dimensions without coordinates: y │ Data variables: │ f (y) float64 8B 3.0 ├── Group: /b/c └── Group: /b/d - Dimensions: () + Dimensions: (x: 2, y: 1) + Coordinates: + * x (x) float64 16B 2.0 3.0 + Dimensions without coordinates: y Data variables: g float64 8B 4.0 """ @@ -680,6 +709,192 @@ def test_repr(self): assert result == expected +def _exact_match(message: str) -> str: + return re.escape(dedent(message).strip()) + return "^" + re.escape(dedent(message.rstrip())) + "$" + + +class TestInheritance: + def test_inherited_dims(self): + dt = DataTree.from_dict( + { + "/": xr.Dataset({"d": (("x",), [1, 2])}), + "/b": xr.Dataset({"e": (("y",), [3])}), + "/c": xr.Dataset({"f": (("y",), [3, 4, 5])}), + } + ) + assert dt.sizes == {"x": 2} + # nodes should include inherited dimensions + assert dt.b.sizes == {"x": 2, "y": 1} + assert dt.c.sizes == {"x": 2, "y": 3} + # dataset objects created from nodes should not + assert dt.b.ds.sizes == {"y": 1} + assert dt.b.to_dataset(inherited=True).sizes == {"y": 1} + assert dt.b.to_dataset(inherited=False).sizes == {"y": 1} + + def test_inherited_coords_index(self): + dt = DataTree.from_dict( + { + "/": xr.Dataset({"d": (("x",), [1, 2])}, coords={"x": [2, 3]}), + "/b": xr.Dataset({"e": (("y",), [3])}), + } + ) + assert "x" in dt["/b"].indexes + assert "x" in dt["/b"].coords + xr.testing.assert_identical(dt["/x"], dt["/b/x"]) + + def test_inherited_coords_override(self): + dt = DataTree.from_dict( + { + "/": xr.Dataset(coords={"x": 1, "y": 2}), + "/b": xr.Dataset(coords={"x": 4, "z": 3}), + } + ) + assert dt.coords.keys() == {"x", "y"} + root_coords = {"x": 1, "y": 2} + sub_coords = {"x": 4, "y": 2, "z": 3} + xr.testing.assert_equal(dt["/x"], xr.DataArray(1, coords=root_coords)) + xr.testing.assert_equal(dt["/y"], xr.DataArray(2, coords=root_coords)) + assert dt["/b"].coords.keys() == {"x", "y", "z"} + xr.testing.assert_equal(dt["/b/x"], xr.DataArray(4, coords=sub_coords)) + xr.testing.assert_equal(dt["/b/y"], xr.DataArray(2, coords=sub_coords)) + xr.testing.assert_equal(dt["/b/z"], xr.DataArray(3, coords=sub_coords)) + + def test_inconsistent_dims(self): + expected_msg = _exact_match( + """ + group '/b' is not aligned with its parents: + Group: + Dimensions: (x: 1) + Dimensions without coordinates: x + Data variables: + c (x) float64 8B 3.0 + From parents: + Dimensions: (x: 2) + Dimensions without coordinates: x + """ + ) + + with pytest.raises(ValueError, match=expected_msg): + DataTree.from_dict( + { + "/": xr.Dataset({"a": (("x",), [1.0, 2.0])}), + "/b": xr.Dataset({"c": (("x",), [3.0])}), + } + ) + + dt: DataTree = DataTree() + dt["/a"] = xr.DataArray([1.0, 2.0], dims=["x"]) + with pytest.raises(ValueError, match=expected_msg): + dt["/b/c"] = xr.DataArray([3.0], dims=["x"]) + + b: DataTree = DataTree(data=xr.Dataset({"c": (("x",), [3.0])})) + with pytest.raises(ValueError, match=expected_msg): + DataTree( + data=xr.Dataset({"a": (("x",), [1.0, 2.0])}), + children={"b": b}, + ) + + def test_inconsistent_child_indexes(self): + expected_msg = _exact_match( + """ + group '/b' is not aligned with its parents: + Group: + Dimensions: (x: 1) + Coordinates: + * x (x) float64 8B 2.0 + Data variables: + *empty* + From parents: + Dimensions: (x: 1) + Coordinates: + * x (x) float64 8B 1.0 + """ + ) + + with pytest.raises(ValueError, match=expected_msg): + DataTree.from_dict( + { + "/": xr.Dataset(coords={"x": [1.0]}), + "/b": xr.Dataset(coords={"x": [2.0]}), + } + ) + + dt: DataTree = DataTree() + dt.ds = xr.Dataset(coords={"x": [1.0]}) # type: ignore + dt["/b"] = DataTree() + with pytest.raises(ValueError, match=expected_msg): + dt["/b"].ds = xr.Dataset(coords={"x": [2.0]}) + + b: DataTree = DataTree(xr.Dataset(coords={"x": [2.0]})) + with pytest.raises(ValueError, match=expected_msg): + DataTree(data=xr.Dataset(coords={"x": [1.0]}), children={"b": b}) + + def test_inconsistent_grandchild_indexes(self): + expected_msg = _exact_match( + """ + group '/b/c' is not aligned with its parents: + Group: + Dimensions: (x: 1) + Coordinates: + * x (x) float64 8B 2.0 + Data variables: + *empty* + From parents: + Dimensions: (x: 1) + Coordinates: + * x (x) float64 8B 1.0 + """ + ) + + with pytest.raises(ValueError, match=expected_msg): + DataTree.from_dict( + { + "/": xr.Dataset(coords={"x": [1.0]}), + "/b/c": xr.Dataset(coords={"x": [2.0]}), + } + ) + + dt: DataTree = DataTree() + dt.ds = xr.Dataset(coords={"x": [1.0]}) # type: ignore + dt["/b/c"] = DataTree() + with pytest.raises(ValueError, match=expected_msg): + dt["/b/c"].ds = xr.Dataset(coords={"x": [2.0]}) + + c: DataTree = DataTree(xr.Dataset(coords={"x": [2.0]})) + b: DataTree = DataTree(children={"c": c}) + with pytest.raises(ValueError, match=expected_msg): + DataTree(data=xr.Dataset(coords={"x": [1.0]}), children={"b": b}) + + def test_inconsistent_grandchild_dims(self): + expected_msg = _exact_match( + """ + group '/b/c' is not aligned with its parents: + Group: + Dimensions: (x: 1) + Dimensions without coordinates: x + Data variables: + d (x) float64 8B 3.0 + From parents: + Dimensions: (x: 2) + Dimensions without coordinates: x + """ + ) + + with pytest.raises(ValueError, match=expected_msg): + DataTree.from_dict( + { + "/": xr.Dataset({"a": (("x",), [1.0, 2.0])}), + "/b/c": xr.Dataset({"d": (("x",), [3.0])}), + } + ) + + dt: DataTree = DataTree() + dt["/a"] = xr.DataArray([1.0, 2.0], dims=["x"]) + with pytest.raises(ValueError, match=expected_msg): + dt["/b/c/d"] = xr.DataArray([3.0], dims=["x"]) + + class TestRestructuring: def test_drop_nodes(self): sue = DataTree.from_dict({"Mary": None, "Kate": None, "Ashley": None}) diff --git a/xarray/tests/test_formatting.py b/xarray/tests/test_formatting.py index d7a46eeaefc..9d0eb81bace 100644 --- a/xarray/tests/test_formatting.py +++ b/xarray/tests/test_formatting.py @@ -10,9 +10,20 @@ import xarray as xr from xarray.core import formatting from xarray.core.datatree import DataTree # TODO: Remove when can do xr.DataTree +from xarray.core.indexes import Index from xarray.tests import requires_cftime, requires_dask, requires_netCDF4 +class CustomIndex(Index): + names: tuple[str, ...] + + def __init__(self, names: tuple[str, ...]): + self.names = names + + def __repr__(self): + return f"CustomIndex(coords={self.names})" + + class TestFormatting: def test_get_indexer_at_least_n_items(self) -> None: cases = [ @@ -219,17 +230,6 @@ def test_attribute_repr(self) -> None: assert "\t" not in tabs def test_index_repr(self) -> None: - from xarray.core.indexes import Index - - class CustomIndex(Index): - names: tuple[str, ...] - - def __init__(self, names: tuple[str, ...]): - self.names = names - - def __repr__(self): - return f"CustomIndex(coords={self.names})" - coord_names = ("x", "y") index = CustomIndex(coord_names) names = ("x",) @@ -258,15 +258,6 @@ def _repr_inline_(self, max_width: int): ), ) def test_index_repr_grouping(self, names) -> None: - from xarray.core.indexes import Index - - class CustomIndex(Index): - def __init__(self, names): - self.names = names - - def __repr__(self): - return f"CustomIndex(coords={self.names})" - index = CustomIndex(names) normal = formatting.summarize_index(names, index, col_width=20) @@ -337,6 +328,51 @@ def test_diff_array_repr(self) -> None: # depending on platform, dtype may not be shown in numpy array repr assert actual == expected.replace(", dtype=int64", "") + da_a = xr.DataArray( + np.array([[1, 2, 3], [4, 5, 6]], dtype="int8"), + dims=("x", "y"), + coords=xr.Coordinates( + { + "x": np.array([True, False], dtype="bool"), + "y": np.array([1, 2, 3], dtype="int16"), + }, + indexes={"y": CustomIndex(("y",))}, + ), + ) + + da_b = xr.DataArray( + np.array([1, 2], dtype="int8"), + dims="x", + coords=xr.Coordinates( + { + "x": np.array([True, False], dtype="bool"), + "label": ("x", np.array([1, 2], dtype="int16")), + }, + indexes={"label": CustomIndex(("label",))}, + ), + ) + + expected = dedent( + """\ + Left and right DataArray objects are not equal + Differing dimensions: + (x: 2, y: 3) != (x: 2) + Differing values: + L + array([[1, 2, 3], + [4, 5, 6]], dtype=int8) + R + array([1, 2], dtype=int8) + Coordinates only on the left object: + * y (y) int16 6B 1 2 3 + Coordinates only on the right object: + * label (x) int16 4B 1 2 + """.rstrip() + ) + + actual = formatting.diff_array_repr(da_a, da_b, "equals") + assert actual == expected + va = xr.Variable( "x", np.array([1, 2, 3], dtype="int64"), {"title": "test Variable"} ) @@ -399,6 +435,36 @@ def test_diff_attrs_repr_with_array(self) -> None: actual = formatting.diff_attrs_repr(attrs_a, attrs_c, "equals") assert expected == actual + def test__diff_mapping_repr_array_attrs_on_variables(self) -> None: + a = { + "a": xr.DataArray( + dims="x", + data=np.array([1], dtype="int16"), + attrs={"b": np.array([1, 2], dtype="int8")}, + ) + } + b = { + "a": xr.DataArray( + dims="x", + data=np.array([1], dtype="int16"), + attrs={"b": np.array([2, 3], dtype="int8")}, + ) + } + actual = formatting.diff_data_vars_repr(a, b, compat="identical", col_width=8) + expected = dedent( + """\ + Differing data variables: + L a (x) int16 2B 1 + Differing variable attributes: + b: [1 2] + R a (x) int16 2B 1 + Differing variable attributes: + b: [2 3] + """.rstrip() + ) + + assert actual == expected + def test_diff_dataset_repr(self) -> None: ds_a = xr.Dataset( data_vars={ diff --git a/xarray/tests/test_groupby.py b/xarray/tests/test_groupby.py index 7c0bd03bab7..40ba1a23f73 100644 --- a/xarray/tests/test_groupby.py +++ b/xarray/tests/test_groupby.py @@ -2053,17 +2053,17 @@ def test_upsample_tolerance(self) -> None: array = DataArray(np.arange(2), [("time", times)]) # Forward fill - actual = array.resample(time="6h").ffill(tolerance="12h") # type: ignore[arg-type] # TODO: tolerance also allows strings, same issue in .reindex. + actual = array.resample(time="6h").ffill(tolerance="12h") expected = DataArray([0.0, 0.0, 0.0, np.nan, 1.0], [("time", times_upsampled)]) assert_identical(expected, actual) # Backward fill - actual = array.resample(time="6h").bfill(tolerance="12h") # type: ignore[arg-type] # TODO: tolerance also allows strings, same issue in .reindex. + actual = array.resample(time="6h").bfill(tolerance="12h") expected = DataArray([0.0, np.nan, 1.0, 1.0, 1.0], [("time", times_upsampled)]) assert_identical(expected, actual) # Nearest - actual = array.resample(time="6h").nearest(tolerance="6h") # type: ignore[arg-type] # TODO: tolerance also allows strings, same issue in .reindex. + actual = array.resample(time="6h").nearest(tolerance="6h") expected = DataArray([0, 0, np.nan, 1, 1], [("time", times_upsampled)]) assert_identical(expected, actual) diff --git a/xarray/tests/test_missing.py b/xarray/tests/test_missing.py index da9513a7c71..bd75f633b82 100644 --- a/xarray/tests/test_missing.py +++ b/xarray/tests/test_missing.py @@ -421,46 +421,37 @@ def test_ffill(): assert_equal(actual, expected) -def test_ffill_use_bottleneck_numbagg(): +@pytest.mark.parametrize("compute_backend", [None], indirect=True) +@pytest.mark.parametrize("method", ["ffill", "bfill"]) +def test_b_ffill_use_bottleneck_numbagg(method, compute_backend): + """ + bfill & ffill fail if both bottleneck and numba are disabled + """ da = xr.DataArray(np.array([4, 5, np.nan], dtype=np.float64), dims="x") - with xr.set_options(use_bottleneck=False, use_numbagg=False): - with pytest.raises(RuntimeError): - da.ffill("x") + with pytest.raises(RuntimeError): + getattr(da, method)("x") @requires_dask -def test_ffill_use_bottleneck_dask(): +@pytest.mark.parametrize("compute_backend", [None], indirect=True) +@pytest.mark.parametrize("method", ["ffill", "bfill"]) +def test_b_ffill_use_bottleneck_dask(method, compute_backend): + """ + ffill fails if both bottleneck and numba are disabled, on dask arrays + """ da = xr.DataArray(np.array([4, 5, np.nan], dtype=np.float64), dims="x") - da = da.chunk({"x": 1}) - with xr.set_options(use_bottleneck=False, use_numbagg=False): - with pytest.raises(RuntimeError): - da.ffill("x") + with pytest.raises(RuntimeError): + getattr(da, method)("x") @requires_numbagg @requires_dask -def test_ffill_use_numbagg_dask(): - with xr.set_options(use_bottleneck=False): - da = xr.DataArray(np.array([4, 5, np.nan], dtype=np.float64), dims="x") - da = da.chunk(x=-1) - # Succeeds with a single chunk: - _ = da.ffill("x").compute() - - -def test_bfill_use_bottleneck(): - da = xr.DataArray(np.array([4, 5, np.nan], dtype=np.float64), dims="x") - with xr.set_options(use_bottleneck=False, use_numbagg=False): - with pytest.raises(RuntimeError): - da.bfill("x") - - -@requires_dask -def test_bfill_use_bottleneck_dask(): +@pytest.mark.parametrize("compute_backend", ["numbagg"], indirect=True) +def test_ffill_use_numbagg_dask(compute_backend): da = xr.DataArray(np.array([4, 5, np.nan], dtype=np.float64), dims="x") - da = da.chunk({"x": 1}) - with xr.set_options(use_bottleneck=False, use_numbagg=False): - with pytest.raises(RuntimeError): - da.bfill("x") + da = da.chunk(x=-1) + # Succeeds with a single chunk: + _ = da.ffill("x").compute() @requires_bottleneck diff --git a/xarray/tests/test_utils.py b/xarray/tests/test_utils.py index 50061c774a8..ecec3ca507b 100644 --- a/xarray/tests/test_utils.py +++ b/xarray/tests/test_utils.py @@ -7,7 +7,11 @@ import pytest from xarray.core import duck_array_ops, utils -from xarray.core.utils import either_dict_or_kwargs, infix_dims, iterate_nested +from xarray.core.utils import ( + either_dict_or_kwargs, + infix_dims, + iterate_nested, +) from xarray.tests import assert_array_equal, requires_dask