From aa6a3ea02d023fa0eaadbfacf4e99a117af67b28 Mon Sep 17 00:00:00 2001 From: Maximilian Roos Date: Fri, 14 Mar 2025 21:11:27 -0700 Subject: [PATCH 1/8] Refactor modules from `core` into `xarray.computation` --- pyproject.toml | 3 +-- xarray/__init__.py | 10 +++++----- xarray/coding/calendar_ops.py | 2 +- xarray/computation/__init__.py | 0 xarray/{core => computation}/arithmetic.py | 7 ++++--- xarray/{core => computation}/computation.py | 3 +++ xarray/{core => computation}/nanops.py | 0 xarray/{core => computation}/ops.py | 9 +++++---- xarray/{core => computation}/rolling.py | 4 ++-- xarray/{core => computation}/rolling_exp.py | 2 +- xarray/{core => computation}/weighted.py | 2 +- xarray/core/_typed_ops.py | 3 ++- xarray/core/accessor_str.py | 2 +- xarray/core/common.py | 21 +++++++++++--------- xarray/core/dataarray.py | 19 +++++++++--------- xarray/core/dataset.py | 22 ++++++++++----------- xarray/core/duck_array_ops.py | 4 ++-- xarray/core/groupby.py | 8 ++++++-- xarray/core/missing.py | 2 +- xarray/core/variable.py | 15 +++++++------- xarray/groupers.py | 2 +- xarray/tests/test_computation.py | 4 ++-- xarray/util/generate_ops.py | 3 ++- 23 files changed, 81 insertions(+), 66 deletions(-) create mode 100644 xarray/computation/__init__.py rename xarray/{core => computation}/arithmetic.py (94%) rename xarray/{core => computation}/computation.py (99%) rename xarray/{core => computation}/nanops.py (100%) rename xarray/{core => computation}/ops.py (96%) rename xarray/{core => computation}/rolling.py (99%) rename xarray/{core => computation}/rolling_exp.py (99%) rename xarray/{core => computation}/weighted.py (99%) diff --git a/pyproject.toml b/pyproject.toml index 85cb8f1bc0c..0e9817d32b8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -170,8 +170,7 @@ module = [ "xarray.core.accessor_dt", "xarray.core.accessor_str", "xarray.core.alignment", - "xarray.core.computation", - "xarray.core.rolling_exp", + "xarray.computation.*", "xarray.indexes.*", "xarray.tests.*", ] diff --git a/xarray/__init__.py b/xarray/__init__.py index 8af936ed27a..bad37cecbf7 100644 --- a/xarray/__init__.py +++ b/xarray/__init__.py @@ -15,11 +15,7 @@ from xarray.coding.cftime_offsets import cftime_range, date_range, date_range_like from xarray.coding.cftimeindex import CFTimeIndex from xarray.coding.frequencies import infer_freq -from xarray.conventions import SerializationWarning, decode_cf -from xarray.core.alignment import align, broadcast -from xarray.core.combine import combine_by_coords, combine_nested -from xarray.core.common import ALL_DIMS, full_like, ones_like, zeros_like -from xarray.core.computation import ( +from xarray.computation.computation import ( apply_ufunc, corr, cov, @@ -29,6 +25,10 @@ unify_chunks, where, ) +from xarray.conventions import SerializationWarning, decode_cf +from xarray.core.alignment import align, broadcast +from xarray.core.combine import combine_by_coords, combine_nested +from xarray.core.common import ALL_DIMS, full_like, ones_like, zeros_like from xarray.core.concat import concat from xarray.core.coordinates import Coordinates from xarray.core.dataarray import DataArray diff --git a/xarray/coding/calendar_ops.py b/xarray/coding/calendar_ops.py index 22a19a63871..60afd47a148 100644 --- a/xarray/coding/calendar_ops.py +++ b/xarray/coding/calendar_ops.py @@ -9,12 +9,12 @@ _should_cftime_be_used, convert_times, ) +from xarray.computation.computation import apply_ufunc from xarray.core.common import ( _contains_datetime_like_objects, full_like, is_np_datetime_like, ) -from xarray.core.computation import apply_ufunc try: import cftime diff --git a/xarray/computation/__init__.py b/xarray/computation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/xarray/core/arithmetic.py b/xarray/computation/arithmetic.py similarity index 94% rename from xarray/core/arithmetic.py rename to xarray/computation/arithmetic.py index 734d7b328de..7a3f342d6c0 100644 --- a/xarray/core/arithmetic.py +++ b/xarray/computation/arithmetic.py @@ -6,6 +6,8 @@ import numpy as np +from xarray.computation.ops import IncludeNumpySameMethods, IncludeReduceMethods + # _typed_ops.py is a generated file from xarray.core._typed_ops import ( DataArrayGroupByOpsMixin, @@ -15,7 +17,6 @@ VariableOpsMixin, ) from xarray.core.common import ImplementsArrayReduce, ImplementsDatasetReduce -from xarray.core.ops import IncludeNumpySameMethods, IncludeReduceMethods from xarray.core.options import OPTIONS, _get_keep_attrs from xarray.namedarray.utils import is_duck_array @@ -29,7 +30,7 @@ class SupportsArithmetic: __slots__ = () # TODO: implement special methods for arithmetic here rather than injecting - # them in xarray/core/ops.py. Ideally, do so by inheriting from + # them in xarray/computation/ops.py. Ideally, do so by inheriting from # numpy.lib.mixins.NDArrayOperatorsMixin. # TODO: allow extending this with some sort of registration system @@ -41,7 +42,7 @@ class SupportsArithmetic: ) def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): - from xarray.core.computation import apply_ufunc + from xarray.computation.computation import apply_ufunc # See the docstring example for numpy.lib.mixins.NDArrayOperatorsMixin. out = kwargs.get("out", ()) diff --git a/xarray/core/computation.py b/xarray/computation/computation.py similarity index 99% rename from xarray/core/computation.py rename to xarray/computation/computation.py index 39da85982a8..f782ebb61f4 100644 --- a/xarray/core/computation.py +++ b/xarray/computation/computation.py @@ -1,5 +1,8 @@ """ Functions for applying functions that act on arrays to xarray's labeled data. + +NOTE: This module is currently large and contains various computational functionality. +The long-term plan is to break it down into more focused submodules. """ from __future__ import annotations diff --git a/xarray/core/nanops.py b/xarray/computation/nanops.py similarity index 100% rename from xarray/core/nanops.py rename to xarray/computation/nanops.py diff --git a/xarray/core/ops.py b/xarray/computation/ops.py similarity index 96% rename from xarray/core/ops.py rename to xarray/computation/ops.py index c67b46692b8..4477a3f181d 100644 --- a/xarray/core/ops.py +++ b/xarray/computation/ops.py @@ -1,6 +1,6 @@ """Define core operations for xarray objects. -TODO(shoyer): rewrite this module, making use of xarray.core.computation, +TODO(shoyer): rewrite this module, making use of xarray.computation.computation, NumPy's __array_ufunc__ and mixin classes instead of the unintuitive "inject" functions. """ @@ -8,6 +8,7 @@ from __future__ import annotations import operator +from typing import Literal import numpy as np @@ -143,7 +144,7 @@ def fillna(data, other, join="left", dataset_join="left"): - "left": take only variables from the first object - "right": take only variables from the last object """ - from xarray.core.computation import apply_ufunc + from xarray.computation.computation import apply_ufunc return apply_ufunc( duck_array_ops.fillna, @@ -172,10 +173,10 @@ def where_method(self, cond, other=dtypes.NA): ------- Same type as caller. """ - from xarray.core.computation import apply_ufunc + from xarray.computation.computation import apply_ufunc # alignment for three arguments is complicated, so don't support it yet - join = "inner" if other is dtypes.NA else "exact" + join: Literal["inner", "exact"] = "inner" if other is dtypes.NA else "exact" return apply_ufunc( duck_array_ops.where_method, self, diff --git a/xarray/core/rolling.py b/xarray/computation/rolling.py similarity index 99% rename from xarray/core/rolling.py rename to xarray/computation/rolling.py index 3ed3020ba89..cc54bc6c14c 100644 --- a/xarray/core/rolling.py +++ b/xarray/computation/rolling.py @@ -10,8 +10,8 @@ import numpy as np from xarray.compat import dask_array_ops +from xarray.computation.arithmetic import CoarsenArithmetic from xarray.core import dtypes, duck_array_ops, utils -from xarray.core.arithmetic import CoarsenArithmetic from xarray.core.options import OPTIONS, _get_keep_attrs from xarray.core.types import CoarsenBoundaryOptions, SideOptions, T_Xarray from xarray.core.utils import ( @@ -645,7 +645,7 @@ def _bottleneck_reduce(self, func, keep_attrs, **kwargs): # bottleneck doesn't allow min_count to be 0, although it should # work the same as if min_count = 1 # Note bottleneck only works with 1d-rolling. - if self.min_periods is not None and self.min_periods == 0: + if self.min_periods == 0: min_count = 1 else: min_count = self.min_periods diff --git a/xarray/core/rolling_exp.py b/xarray/computation/rolling_exp.py similarity index 99% rename from xarray/core/rolling_exp.py rename to xarray/computation/rolling_exp.py index 3f32bb2de26..45f5db73167 100644 --- a/xarray/core/rolling_exp.py +++ b/xarray/computation/rolling_exp.py @@ -6,7 +6,7 @@ import numpy as np from xarray.compat.pdcompat import count_not_none -from xarray.core.computation import apply_ufunc +from xarray.computation.computation import apply_ufunc from xarray.core.options import _get_keep_attrs from xarray.core.types import T_DataWithCoords from xarray.core.utils import module_available diff --git a/xarray/core/weighted.py b/xarray/computation/weighted.py similarity index 99% rename from xarray/core/weighted.py rename to xarray/computation/weighted.py index cd24091b18e..3715541a443 100644 --- a/xarray/core/weighted.py +++ b/xarray/computation/weighted.py @@ -6,9 +6,9 @@ import numpy as np from numpy.typing import ArrayLike +from xarray.computation.computation import apply_ufunc, dot from xarray.core import duck_array_ops, utils from xarray.core.alignment import align, broadcast -from xarray.core.computation import apply_ufunc, dot from xarray.core.types import Dims, T_DataArray, T_Xarray from xarray.namedarray.utils import is_duck_dask_array diff --git a/xarray/core/_typed_ops.py b/xarray/core/_typed_ops.py index b77bb80e61b..b936a657ae6 100644 --- a/xarray/core/_typed_ops.py +++ b/xarray/core/_typed_ops.py @@ -8,7 +8,8 @@ from collections.abc import Callable from typing import TYPE_CHECKING, Any, overload -from xarray.core import nputils, ops +from xarray.core import nputils +from xarray.computation import ops from xarray.core.types import ( DaCompatible, DsCompatible, diff --git a/xarray/core/accessor_str.py b/xarray/core/accessor_str.py index e44ef75a88b..171436a5646 100644 --- a/xarray/core/accessor_str.py +++ b/xarray/core/accessor_str.py @@ -51,8 +51,8 @@ import numpy as np +from xarray.computation.computation import apply_ufunc from xarray.core import duck_array_ops -from xarray.core.computation import apply_ufunc from xarray.core.types import T_DataArray if TYPE_CHECKING: diff --git a/xarray/core/common.py b/xarray/core/common.py index ceaae42356a..93036a612ea 100644 --- a/xarray/core/common.py +++ b/xarray/core/common.py @@ -11,7 +11,7 @@ import numpy as np import pandas as pd -from xarray.core import dtypes, duck_array_ops, formatting, formatting_html, ops +from xarray.core import dtypes, duck_array_ops, formatting, formatting_html from xarray.core.indexing import BasicIndexer, ExplicitlyIndexed from xarray.core.options import OPTIONS, _get_keep_attrs from xarray.core.types import ResampleCompatible @@ -36,11 +36,11 @@ if TYPE_CHECKING: from numpy.typing import DTypeLike + from xarray.computation.rolling_exp import RollingExp from xarray.core.dataarray import DataArray from xarray.core.dataset import Dataset from xarray.core.indexes import Index from xarray.core.resample import Resample - from xarray.core.rolling_exp import RollingExp from xarray.core.types import ( DatetimeLike, DTypeLikeSave, @@ -491,7 +491,7 @@ def clip( -------- numpy.clip : equivalent function """ - from xarray.core.computation import apply_ufunc + from xarray.computation.computation import apply_ufunc if keep_attrs is None: # When this was a unary func, the default was True, so retaining the @@ -900,7 +900,6 @@ def rolling_exp( -------- core.rolling_exp.RollingExp """ - from xarray.core import rolling_exp if "keep_attrs" in window_kwargs: warnings.warn( @@ -912,7 +911,9 @@ def rolling_exp( window = either_dict_or_kwargs(window, window_kwargs, "rolling_exp") - return rolling_exp.RollingExp(self, window, window_type) + from xarray.computation.rolling_exp import RollingExp + + return RollingExp(self, window, window_type) def _resample( self, @@ -1253,6 +1254,8 @@ def _dataset_indexer(dim: Hashable) -> DataArray: self = self.isel(**indexers) cond = cond.isel(**indexers) + from xarray.computation import ops + return ops.where_method(self, cond, other) def set_close(self, close: Callable[[], None] | None) -> None: @@ -1308,7 +1311,7 @@ def isnull(self, keep_attrs: bool | None = None) -> Self: array([False, True, False]) Dimensions without coordinates: x """ - from xarray.core.computation import apply_ufunc + from xarray.computation.computation import apply_ufunc if keep_attrs is None: keep_attrs = _get_keep_attrs(default=False) @@ -1351,7 +1354,7 @@ def notnull(self, keep_attrs: bool | None = None) -> Self: array([ True, False, True]) Dimensions without coordinates: x """ - from xarray.core.computation import apply_ufunc + from xarray.computation.computation import apply_ufunc if keep_attrs is None: keep_attrs = _get_keep_attrs(default=False) @@ -1390,7 +1393,7 @@ def isin(self, test_elements: Any) -> Self: -------- numpy.isin """ - from xarray.core.computation import apply_ufunc + from xarray.computation.computation import apply_ufunc from xarray.core.dataarray import DataArray from xarray.core.dataset import Dataset from xarray.core.variable import Variable @@ -1473,7 +1476,7 @@ def astype( dask.array.Array.astype sparse.COO.astype """ - from xarray.core.computation import apply_ufunc + from xarray.computation.computation import apply_ufunc kwargs = dict(order=order, casting=casting, subok=subok, copy=copy) kwargs = {k: v for k, v in kwargs.items() if v is not None} diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py index 4324a4587b3..87e95bdc464 100644 --- a/xarray/core/dataarray.py +++ b/xarray/core/dataarray.py @@ -29,7 +29,10 @@ from xarray.coding.calendar_ops import convert_calendar, interp_calendar from xarray.coding.cftimeindex import CFTimeIndex -from xarray.core import alignment, computation, dtypes, indexing, ops, utils +from xarray.computation import computation, ops +from xarray.computation.arithmetic import DataArrayArithmetic +from xarray.computation.computation import unify_chunks +from xarray.core import alignment, dtypes, indexing, utils from xarray.core._aggregations import DataArrayAggregations from xarray.core.accessor_dt import CombinedDatetimelikeAccessor from xarray.core.accessor_str import StringAccessor @@ -38,9 +41,7 @@ _get_broadcast_dims_map_common_coords, align, ) -from xarray.core.arithmetic import DataArrayArithmetic from xarray.core.common import AbstractArray, DataWithCoords, get_chunksizes -from xarray.core.computation import unify_chunks from xarray.core.coordinates import ( Coordinates, DataArrayCoordinates, @@ -97,9 +98,10 @@ from xarray.backends import ZarrStore from xarray.backends.api import T_NetcdfEngine, T_NetcdfTypes + from xarray.computation.rolling import DataArrayCoarsen, DataArrayRolling + from xarray.computation.weighted import DataArrayWeighted from xarray.core.groupby import DataArrayGroupBy from xarray.core.resample import DataArrayResample - from xarray.core.rolling import DataArrayCoarsen, DataArrayRolling from xarray.core.types import ( CoarsenBoundaryOptions, DatetimeLike, @@ -123,7 +125,6 @@ T_ChunksFreq, T_Xarray, ) - from xarray.core.weighted import DataArrayWeighted from xarray.groupers import Grouper, Resampler from xarray.namedarray.parallelcompat import ChunkManagerEntrypoint @@ -7048,7 +7049,7 @@ def weighted(self, weights: DataArray) -> DataArrayWeighted: Tutorial on Weighted Reduction using :py:func:`~xarray.DataArray.weighted` """ - from xarray.core.weighted import DataArrayWeighted + from xarray.computation.weighted import DataArrayWeighted return DataArrayWeighted(self, weights) @@ -7122,7 +7123,7 @@ def rolling( Dataset.rolling core.rolling.DataArrayRolling """ - from xarray.core.rolling import DataArrayRolling + from xarray.computation.rolling import DataArrayRolling dim = either_dict_or_kwargs(dim, window_kwargs, "rolling") return DataArrayRolling(self, dim, min_periods=min_periods, center=center) @@ -7182,7 +7183,7 @@ def cumulative( Dataset.cumulative core.rolling.DataArrayRolling """ - from xarray.core.rolling import DataArrayRolling + from xarray.computation.rolling import DataArrayRolling # Could we abstract this "normalize and check 'dim'" logic? It's currently shared # with the same method in Dataset. @@ -7336,7 +7337,7 @@ def coarsen( Tutorial on windowed computation using :py:func:`~xarray.DataArray.coarsen` """ - from xarray.core.rolling import DataArrayCoarsen + from xarray.computation.rolling import DataArrayCoarsen dim = either_dict_or_kwargs(dim, window_kwargs, "coarsen") return DataArrayCoarsen( diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index 8f55f9b71a0..68f0caa678d 100644 --- a/xarray/core/dataset.py +++ b/xarray/core/dataset.py @@ -42,12 +42,14 @@ from xarray.coding.calendar_ops import convert_calendar, interp_calendar from xarray.coding.cftimeindex import CFTimeIndex, _parse_array_of_cftime_strings from xarray.compat.array_api_compat import to_like_array +from xarray.computation import ops +from xarray.computation.arithmetic import DatasetArithmetic +from xarray.computation.computation import _ensure_numeric, unify_chunks from xarray.core import ( alignment, duck_array_ops, formatting, formatting_html, - ops, utils, ) from xarray.core import dtypes as xrdtypes @@ -57,13 +59,11 @@ _get_broadcast_dims_map_common_coords, align, ) -from xarray.core.arithmetic import DatasetArithmetic from xarray.core.common import ( DataWithCoords, _contains_datetime_like_objects, get_chunksizes, ) -from xarray.core.computation import _ensure_numeric, unify_chunks from xarray.core.coordinates import ( Coordinates, DatasetCoordinates, @@ -140,11 +140,12 @@ from xarray.backends import AbstractDataStore, ZarrStore from xarray.backends.api import T_NetcdfEngine, T_NetcdfTypes + from xarray.computation.rolling import DatasetCoarsen, DatasetRolling + from xarray.computation.weighted import DatasetWeighted from xarray.core.dataarray import DataArray from xarray.core.groupby import DatasetGroupBy from xarray.core.merge import CoercibleMapping, CoercibleValue from xarray.core.resample import DatasetResample - from xarray.core.rolling import DatasetCoarsen, DatasetRolling from xarray.core.types import ( CFCalendar, CoarsenBoundaryOptions, @@ -173,7 +174,6 @@ T_DatasetPadConstantValues, T_Xarray, ) - from xarray.core.weighted import DatasetWeighted from xarray.groupers import Grouper, Resampler from xarray.namedarray.parallelcompat import ChunkManagerEntrypoint @@ -201,7 +201,7 @@ def _initialize_curvefit_params(params, p0, bounds, func_args): """Set initial guess and bounds for curvefit. Priority: 1) passed args 2) func signature 3) scipy defaults """ - from xarray.core.computation import where + from xarray.computation.computation import where def _initialize_feasible(lb, ub): # Mimics functionality of scipy.optimize.minpack._initialize_feasible @@ -9783,8 +9783,8 @@ def curvefit( """ from scipy.optimize import curve_fit + from xarray.computation.computation import apply_ufunc from xarray.core.alignment import broadcast - from xarray.core.computation import apply_ufunc from xarray.core.dataarray import _THIS_ARRAY, DataArray if p0 is None: @@ -10381,7 +10381,7 @@ def weighted(self, weights: DataArray) -> DatasetWeighted: Tutorial on Weighted Reduction using :py:func:`~xarray.Dataset.weighted` """ - from xarray.core.weighted import DatasetWeighted + from xarray.computation.weighted import DatasetWeighted return DatasetWeighted(self, weights) @@ -10421,7 +10421,7 @@ def rolling( DataArray.rolling DataArray.rolling_exp """ - from xarray.core.rolling import DatasetRolling + from xarray.computation.rolling import DatasetRolling dim = either_dict_or_kwargs(dim, window_kwargs, "rolling") return DatasetRolling(self, dim, min_periods=min_periods, center=center) @@ -10453,7 +10453,7 @@ def cumulative( Dataset.rolling Dataset.rolling_exp """ - from xarray.core.rolling import DatasetRolling + from xarray.computation.rolling import DatasetRolling if isinstance(dim, str): if dim not in self.dims: @@ -10514,7 +10514,7 @@ def coarsen( Tutorial on windowed computation using :py:func:`~xarray.Dataset.coarsen` """ - from xarray.core.rolling import DatasetCoarsen + from xarray.computation.rolling import DatasetCoarsen dim = either_dict_or_kwargs(dim, window_kwargs, "coarsen") return DatasetCoarsen( diff --git a/xarray/core/duck_array_ops.py b/xarray/core/duck_array_ops.py index b1d842f30f3..96330a64b68 100644 --- a/xarray/core/duck_array_ops.py +++ b/xarray/core/duck_array_ops.py @@ -467,8 +467,6 @@ def _ignore_warnings_if(condition): def _create_nan_agg_method(name, coerce_strings=False, invariant_0d=False): - from xarray.core import nanops - def f(values, axis=None, skipna=None, **kwargs): if kwargs.pop("out", None) is not None: raise TypeError(f"`out` is not valid for {name}") @@ -495,6 +493,8 @@ def f(values, axis=None, skipna=None, **kwargs): or dtypes.is_object(values.dtype) ) ): + from xarray.computation import nanops + nanname = "nan" + name func = getattr(nanops, nanname) else: diff --git a/xarray/core/groupby.py b/xarray/core/groupby.py index 3bdc30eb1f0..fa0c20d98fe 100644 --- a/xarray/core/groupby.py +++ b/xarray/core/groupby.py @@ -12,13 +12,17 @@ import pandas as pd from packaging.version import Version -from xarray.core import dtypes, duck_array_ops, nputils, ops +from xarray.computation import ops +from xarray.computation.arithmetic import ( + DataArrayGroupbyArithmetic, + DatasetGroupbyArithmetic, +) +from xarray.core import dtypes, duck_array_ops, nputils from xarray.core._aggregations import ( DataArrayGroupByAggregations, DatasetGroupByAggregations, ) from xarray.core.alignment import align, broadcast -from xarray.core.arithmetic import DataArrayGroupbyArithmetic, DatasetGroupbyArithmetic from xarray.core.common import ImplementsArrayReduce, ImplementsDatasetReduce from xarray.core.concat import concat from xarray.core.coordinates import Coordinates, _coordinates_from_variable diff --git a/xarray/core/missing.py b/xarray/core/missing.py index b4ca36b31df..67f7af1270a 100644 --- a/xarray/core/missing.py +++ b/xarray/core/missing.py @@ -12,9 +12,9 @@ import numpy as np import pandas as pd +from xarray.computation.computation import apply_ufunc from xarray.core import utils from xarray.core.common import _contains_datetime_like_objects, ones_like -from xarray.core.computation import apply_ufunc from xarray.core.duck_array_ops import ( datetime_to_numeric, push, diff --git a/xarray/core/variable.py b/xarray/core/variable.py index 5ff6a098016..bc9e6d17a07 100644 --- a/xarray/core/variable.py +++ b/xarray/core/variable.py @@ -17,8 +17,9 @@ import xarray as xr # only for Dataset and DataArray from xarray.compat.array_api_compat import to_like_array -from xarray.core import common, dtypes, duck_array_ops, indexing, nputils, ops, utils -from xarray.core.arithmetic import VariableArithmetic +from xarray.computation import ops +from xarray.computation.arithmetic import VariableArithmetic +from xarray.core import common, dtypes, duck_array_ops, indexing, nputils, utils from xarray.core.common import AbstractArray from xarray.core.extension_array import PandasExtensionArray from xarray.core.indexing import ( @@ -483,7 +484,7 @@ def astype( dask.array.Array.astype sparse.COO.astype """ - from xarray.core.computation import apply_ufunc + from xarray.computation.computation import apply_ufunc kwargs = dict(order=order, casting=casting, subok=subok, copy=copy) kwargs = {k: v for k, v in kwargs.items() if v is not None} @@ -1597,7 +1598,7 @@ def clip(self, min=None, max=None): -------- numpy.clip : equivalent function """ - from xarray.core.computation import apply_ufunc + from xarray.computation.computation import apply_ufunc xp = duck_array_ops.get_array_namespace(self.data) return apply_ufunc(xp.clip, self, min, max, dask="allowed") @@ -1876,7 +1877,7 @@ def quantile( The American Statistician, 50(4), pp. 361-365, 1996 """ - from xarray.core.computation import apply_ufunc + from xarray.computation.computation import apply_ufunc if interpolation is not None: warnings.warn( @@ -2235,7 +2236,7 @@ def isnull(self, keep_attrs: bool | None = None): Size: 3B array([False, True, False]) """ - from xarray.core.computation import apply_ufunc + from xarray.computation.computation import apply_ufunc if keep_attrs is None: keep_attrs = _get_keep_attrs(default=False) @@ -2269,7 +2270,7 @@ def notnull(self, keep_attrs: bool | None = None): Size: 3B array([ True, False, True]) """ - from xarray.core.computation import apply_ufunc + from xarray.computation.computation import apply_ufunc if keep_attrs is None: keep_attrs = _get_keep_attrs(default=False) diff --git a/xarray/groupers.py b/xarray/groupers.py index 32e5e712196..234c9f1398a 100644 --- a/xarray/groupers.py +++ b/xarray/groupers.py @@ -17,7 +17,7 @@ from numpy.typing import ArrayLike from xarray.coding.cftime_offsets import BaseCFTimeOffset, _new_to_legacy_freq -from xarray.core.computation import apply_ufunc +from xarray.computation.computation import apply_ufunc from xarray.core.coordinates import Coordinates, _coordinates_from_variable from xarray.core.dataarray import DataArray from xarray.core.duck_array_ops import array_all, isnull diff --git a/xarray/tests/test_computation.py b/xarray/tests/test_computation.py index 8ca9d2bbbb5..fab118b70e7 100644 --- a/xarray/tests/test_computation.py +++ b/xarray/tests/test_computation.py @@ -10,8 +10,7 @@ from numpy.testing import assert_allclose, assert_array_equal import xarray as xr -from xarray.core.alignment import broadcast -from xarray.core.computation import ( +from xarray.computation.computation import ( _UFuncSignature, apply_ufunc, broadcast_compat_data, @@ -22,6 +21,7 @@ result_name, unified_dim_sizes, ) +from xarray.core.alignment import broadcast from xarray.tests import ( has_dask, raise_if_dask_computes, diff --git a/xarray/util/generate_ops.py b/xarray/util/generate_ops.py index 31e01bfd158..3300bbf594a 100644 --- a/xarray/util/generate_ops.py +++ b/xarray/util/generate_ops.py @@ -260,7 +260,8 @@ def unops() -> list[OpsType]: from collections.abc import Callable from typing import TYPE_CHECKING, Any, overload -from xarray.core import nputils, ops +from xarray.core import nputils +from xarray.computation import ops from xarray.core.types import ( DaCompatible, DsCompatible, From edc36831916a635acbd7236218e4d56ba43d4d43 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 15 Mar 2025 05:22:45 +0000 Subject: [PATCH 2/8] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- xarray/core/_typed_ops.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/core/_typed_ops.py b/xarray/core/_typed_ops.py index b936a657ae6..502ac66a9f6 100644 --- a/xarray/core/_typed_ops.py +++ b/xarray/core/_typed_ops.py @@ -8,8 +8,8 @@ from collections.abc import Callable from typing import TYPE_CHECKING, Any, overload -from xarray.core import nputils from xarray.computation import ops +from xarray.core import nputils from xarray.core.types import ( DaCompatible, DsCompatible, From 54747a92ef6efa99156e1e81012672317dec35ff Mon Sep 17 00:00:00 2001 From: Maximilian Roos Date: Sat, 15 Mar 2025 00:06:06 -0700 Subject: [PATCH 3/8] --- xarray/computation/ops.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/xarray/computation/ops.py b/xarray/computation/ops.py index 4477a3f181d..51a498bf02a 100644 --- a/xarray/computation/ops.py +++ b/xarray/computation/ops.py @@ -176,7 +176,8 @@ def where_method(self, cond, other=dtypes.NA): from xarray.computation.computation import apply_ufunc # alignment for three arguments is complicated, so don't support it yet - join: Literal["inner", "exact"] = "inner" if other is dtypes.NA else "exact" + # Unsure why we get a mypy error here + join: Literal["inner", "exact"] = "inner" if other is dtypes.NA else "exact" # type: ignore[has-type] return apply_ufunc( duck_array_ops.where_method, self, From b493c9cdca97f7c2c9e113af8f5536fb1b105c38 Mon Sep 17 00:00:00 2001 From: Maximilian Roos Date: Sat, 15 Mar 2025 00:25:06 -0700 Subject: [PATCH 4/8] --- xarray/computation/ops.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/xarray/computation/ops.py b/xarray/computation/ops.py index 51a498bf02a..e098088f017 100644 --- a/xarray/computation/ops.py +++ b/xarray/computation/ops.py @@ -158,7 +158,8 @@ def fillna(data, other, join="left", dataset_join="left"): ) -def where_method(self, cond, other=dtypes.NA): +# Unsure why we get a mypy error here +def where_method(self, cond, other=dtypes.NA): # type: ignore[has-type] """Return elements from `self` or `other` depending on `cond`. Parameters @@ -176,8 +177,7 @@ def where_method(self, cond, other=dtypes.NA): from xarray.computation.computation import apply_ufunc # alignment for three arguments is complicated, so don't support it yet - # Unsure why we get a mypy error here - join: Literal["inner", "exact"] = "inner" if other is dtypes.NA else "exact" # type: ignore[has-type] + join: Literal["inner", "exact"] = "inner" if other is dtypes.NA else "exact" return apply_ufunc( duck_array_ops.where_method, self, From 2d87d44c6d346aa815fb02feb68674c19fe115ea Mon Sep 17 00:00:00 2001 From: Maximilian Roos Date: Sat, 15 Mar 2025 00:35:00 -0700 Subject: [PATCH 5/8] --- doc/api.rst | 6 +++--- doc/user-guide/reshaping.rst | 4 ++-- pyproject.toml | 8 ++++---- 3 files changed, 9 insertions(+), 9 deletions(-) diff --git a/doc/api.rst b/doc/api.rst index d4b0fcb2962..d7c2370d348 100644 --- a/doc/api.rst +++ b/doc/api.rst @@ -1334,7 +1334,7 @@ Grouper Objects Rolling objects =============== -.. currentmodule:: xarray.core.rolling +.. currentmodule:: xarray.computation.rolling Dataset ------- @@ -1427,7 +1427,7 @@ DataArray Exponential rolling objects =========================== -.. currentmodule:: xarray.core.rolling_exp +.. currentmodule:: xarray.computation.rolling_exp .. autosummary:: :toctree: generated/ @@ -1439,7 +1439,7 @@ Exponential rolling objects Weighted objects ================ -.. currentmodule:: xarray.core.weighted +.. currentmodule:: xarray.computation.weighted Dataset ------- diff --git a/doc/user-guide/reshaping.rst b/doc/user-guide/reshaping.rst index e3ff836dbe6..aa96190f820 100644 --- a/doc/user-guide/reshaping.rst +++ b/doc/user-guide/reshaping.rst @@ -305,7 +305,7 @@ Reshaping via coarsen Whilst :py:class:`~xarray.DataArray.coarsen` is normally used for reducing your data's resolution by applying a reduction function (see the :ref:`page on computation`), -it can also be used to reorganise your data without applying a computation via :py:meth:`~xarray.core.rolling.DataArrayCoarsen.construct`. +it can also be used to reorganise your data without applying a computation via :py:meth:`~xarray.computation.rolling.DataArrayCoarsen.construct`. Taking our example tutorial air temperature dataset over the Northern US @@ -324,7 +324,7 @@ Taking our example tutorial air temperature dataset over the Northern US @savefig pre_coarsening.png air.isel(time=0).plot(x="lon", y="lat") -we can split this up into sub-regions of size ``(9, 18)`` points using :py:meth:`~xarray.core.rolling.DataArrayCoarsen.construct`: +we can split this up into sub-regions of size ``(9, 18)`` points using :py:meth:`~xarray.computation.rolling.DataArrayCoarsen.construct`: .. ipython:: python diff --git a/pyproject.toml b/pyproject.toml index 0e9817d32b8..fed74619ec9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -102,10 +102,10 @@ fallback_version = "9999" [tool.coverage.run] omit = [ "*/xarray/tests/*", - "*/xarray/core/dask_array_compat.py", - "*/xarray/core/npcompat.py", - "*/xarray/core/pdcompat.py", - "*/xarray/core/pycompat.py", + "*/xarray/compat/dask_array_compat.py", + "*/xarray/compat/npcompat.py", + "*/xarray/compat/pdcompat.py", + "*/xarray/namedarray/pycompat.py", "*/xarray/core/types.py", ] source = ["xarray"] From 651801a0ffcb0a87b8e50e9e99843888eecf1e3c Mon Sep 17 00:00:00 2001 From: Maximilian Roos Date: Sat, 15 Mar 2025 00:56:38 -0700 Subject: [PATCH 6/8] --- doc/whats-new.rst | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 994fc70339c..94912c852dd 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -5936,7 +5936,7 @@ Enhancements supplied list, returning a bool array. See :ref:`selecting values with isin` for full details. Similar to the ``np.isin`` function. By `Maximilian Roos `_. -- Some speed improvement to construct :py:class:`~xarray.core.rolling.DataArrayRolling` +- Some speed improvement to construct :py:class:`~xarray.computation.rolling.DataArrayRolling` object (:issue:`1993`) By `Keisuke Fujii `_. - Handle variables with different values for ``missing_value`` and @@ -6016,8 +6016,8 @@ Enhancements NumPy. By `Stephan Hoyer `_. - Improve :py:func:`~xarray.DataArray.rolling` logic. - :py:func:`~xarray.core.rolling.DataArrayRolling` object now supports - :py:func:`~xarray.core.rolling.DataArrayRolling.construct` method that returns a view + :py:func:`~xarray.computation.rolling.DataArrayRolling` object now supports + :py:func:`~xarray.computation.rolling.DataArrayRolling.construct` method that returns a view of the DataArray / Dataset object with the rolling-window dimension added to the last axis. This enables more flexible operation, such as strided rolling, windowed rolling, ND-rolling, short-time FFT and convolution. From 99fd6ee0937940b979ecf26538637e09aab81568 Mon Sep 17 00:00:00 2001 From: Maximilian Roos Date: Sat, 15 Mar 2025 01:02:43 -0700 Subject: [PATCH 7/8] --- doc/api-hidden.rst | 52 +++++++++++++++++++++++----------------------- doc/whats-new.rst | 12 +++++------ 2 files changed, 32 insertions(+), 32 deletions(-) diff --git a/doc/api-hidden.rst b/doc/api-hidden.rst index 2281d939160..9b60fa5f8d7 100644 --- a/doc/api-hidden.rst +++ b/doc/api-hidden.rst @@ -28,19 +28,19 @@ core.coordinates.DatasetCoordinates.equals core.coordinates.DatasetCoordinates.identical - core.rolling.DatasetCoarsen.boundary - core.rolling.DatasetCoarsen.coord_func - core.rolling.DatasetCoarsen.obj - core.rolling.DatasetCoarsen.side - core.rolling.DatasetCoarsen.trim_excess - core.rolling.DatasetCoarsen.windows - - core.rolling.DatasetRolling.center - core.rolling.DatasetRolling.dim - core.rolling.DatasetRolling.min_periods - core.rolling.DatasetRolling.obj - core.rolling.DatasetRolling.rollings - core.rolling.DatasetRolling.window + computation.rolling.DatasetCoarsen.boundary + computation.rolling.DatasetCoarsen.coord_func + computation.rolling.DatasetCoarsen.obj + computation.rolling.DatasetCoarsen.side + computation.rolling.DatasetCoarsen.trim_excess + computation.rolling.DatasetCoarsen.windows + + computation.rolling.DatasetRolling.center + computation.rolling.DatasetRolling.dim + computation.rolling.DatasetRolling.min_periods + computation.rolling.DatasetRolling.obj + computation.rolling.DatasetRolling.rollings + computation.rolling.DatasetRolling.window core.weighted.DatasetWeighted.obj core.weighted.DatasetWeighted.weights @@ -70,19 +70,19 @@ core.coordinates.DataArrayCoordinates.equals core.coordinates.DataArrayCoordinates.identical - core.rolling.DataArrayCoarsen.boundary - core.rolling.DataArrayCoarsen.coord_func - core.rolling.DataArrayCoarsen.obj - core.rolling.DataArrayCoarsen.side - core.rolling.DataArrayCoarsen.trim_excess - core.rolling.DataArrayCoarsen.windows - - core.rolling.DataArrayRolling.center - core.rolling.DataArrayRolling.dim - core.rolling.DataArrayRolling.min_periods - core.rolling.DataArrayRolling.obj - core.rolling.DataArrayRolling.window - core.rolling.DataArrayRolling.window_labels + computation.rolling.DataArrayCoarsen.boundary + computation.rolling.DataArrayCoarsen.coord_func + computation.rolling.DataArrayCoarsen.obj + computation.rolling.DataArrayCoarsen.side + computation.rolling.DataArrayCoarsen.trim_excess + computation.rolling.DataArrayCoarsen.windows + + computation.rolling.DataArrayRolling.center + computation.rolling.DataArrayRolling.dim + computation.rolling.DataArrayRolling.min_periods + computation.rolling.DataArrayRolling.obj + computation.rolling.DataArrayRolling.window + computation.rolling.DataArrayRolling.window_labels core.weighted.DataArrayWeighted.obj core.weighted.DataArrayWeighted.weights diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 94912c852dd..780036f9b3b 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -1473,7 +1473,7 @@ Bug fixes special case ``NaT`` handling in :py:meth:`~core.accessor_dt.DatetimeAccessor.isocalendar` (:issue:`7928`, :pull:`8084`). By `Kai Mühlbauer `_. -- Fix :py:meth:`~core.rolling.DatasetRolling.construct` with stride on Datasets without indexes. +- Fix :py:meth:`~computation.rolling.DatasetRolling.construct` with stride on Datasets without indexes. (:issue:`7021`, :pull:`7578`). By `Amrest Chinkamol `_ and `Michael Niklas `_. - Calling plot with kwargs ``col``, ``row`` or ``hue`` no longer squeezes dimensions passed via these arguments @@ -3507,7 +3507,7 @@ New Features By `Justus Magin `_. - Allow installing from git archives (:pull:`4897`). By `Justus Magin `_. -- :py:class:`~core.rolling.DataArrayCoarsen` and :py:class:`~core.rolling.DatasetCoarsen` +- :py:class:`~computation.rolling.DataArrayCoarsen` and :py:class:`~computation.rolling.DatasetCoarsen` now implement a ``reduce`` method, enabling coarsening operations with custom reduction functions (:issue:`3741`, :pull:`4939`). By `Spencer Clark `_. @@ -4352,8 +4352,8 @@ New Features - :py:meth:`Dataset.quantile`, :py:meth:`DataArray.quantile` and ``GroupBy.quantile`` now work with dask Variables. By `Deepak Cherian `_. -- Added the ``count`` reduction method to both :py:class:`~core.rolling.DatasetCoarsen` - and :py:class:`~core.rolling.DataArrayCoarsen` objects. (:pull:`3500`) +- Added the ``count`` reduction method to both :py:class:`~computation.rolling.DatasetCoarsen` + and :py:class:`~computation.rolling.DataArrayCoarsen` objects. (:pull:`3500`) By `Deepak Cherian `_ - Add ``meta`` kwarg to :py:func:`~xarray.apply_ufunc`; this is passed on to :py:func:`dask.array.blockwise`. (:pull:`3660`) @@ -4705,7 +4705,7 @@ Bug fixes - Fix error in concatenating unlabeled dimensions (:pull:`3362`). By `Deepak Cherian `_. - Warn if the ``dim`` kwarg is passed to rolling operations. This is redundant since a dimension is - specified when the :py:class:`~core.rolling.DatasetRolling` or :py:class:`~core.rolling.DataArrayRolling` object is created. + specified when the :py:class:`~computation.rolling.DatasetRolling` or :py:class:`~computation.rolling.DataArrayRolling` object is created. (:pull:`3362`). By `Deepak Cherian `_. Documentation @@ -6791,7 +6791,7 @@ Enhancements By `Stephan Hoyer `_ and `Phillip J. Wolfram `_. -- New aggregation on rolling objects :py:meth:`~core.rolling.DataArrayRolling.count` +- New aggregation on rolling objects :py:meth:`~computation.rolling.DataArrayRolling.count` which providing a rolling count of valid values (:issue:`1138`). Bug fixes From 8bee1ac1701f401b4c5c486abfbd07dd0dd45579 Mon Sep 17 00:00:00 2001 From: Maximilian Roos Date: Sat, 15 Mar 2025 01:13:28 -0700 Subject: [PATCH 8/8] --- doc/api-hidden.rst | 8 ++++---- doc/whats-new.rst | 6 +++--- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/doc/api-hidden.rst b/doc/api-hidden.rst index 9b60fa5f8d7..ac8290b3d1b 100644 --- a/doc/api-hidden.rst +++ b/doc/api-hidden.rst @@ -42,8 +42,8 @@ computation.rolling.DatasetRolling.rollings computation.rolling.DatasetRolling.window - core.weighted.DatasetWeighted.obj - core.weighted.DatasetWeighted.weights + computation.weighted.DatasetWeighted.obj + computation.weighted.DatasetWeighted.weights Dataset.load_store Dataset.dump_to_store @@ -84,8 +84,8 @@ computation.rolling.DataArrayRolling.window computation.rolling.DataArrayRolling.window_labels - core.weighted.DataArrayWeighted.obj - core.weighted.DataArrayWeighted.weights + computation.weighted.DataArrayWeighted.obj + computation.weighted.DataArrayWeighted.weights core.coordinates.DataTreeCoordinates.get core.coordinates.DataTreeCoordinates.items diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 780036f9b3b..701d4583512 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -2488,8 +2488,8 @@ New Features - The ``zarr`` backend is now able to read NCZarr. By `Mattia Almansi `_. -- Add a weighted ``quantile`` method to :py:class:`~core.weighted.DatasetWeighted` and - :py:class:`~core.weighted.DataArrayWeighted` (:pull:`6059`). +- Add a weighted ``quantile`` method to :py:class:`.computation.weighted.DatasetWeighted` and + :py:class:`~computation.weighted.DataArrayWeighted` (:pull:`6059`). By `Christian Jauvin `_ and `David Huard `_. - Add a ``create_index=True`` parameter to :py:meth:`Dataset.stack` and :py:meth:`DataArray.stack` so that the creation of multi-indexes is optional @@ -2871,7 +2871,7 @@ Thomas Nicholas, Tomas Chor, Tom Augspurger, Victor Negîrneac, Zachary Blackwoo New Features ~~~~~~~~~~~~ -- Add ``std``, ``var``, ``sum_of_squares`` to :py:class:`~core.weighted.DatasetWeighted` and :py:class:`~core.weighted.DataArrayWeighted`. +- Add ``std``, ``var``, ``sum_of_squares`` to :py:class:`~computation.weighted.DatasetWeighted` and :py:class:`~computation.weighted.DataArrayWeighted`. By `Christian Jauvin `_. - Added a :py:func:`get_options` method to xarray's root namespace (:issue:`5698`, :pull:`5716`) By `Pushkar Kopparla `_.