Skip to content

Refactor modules from core into xarray.computation #10132

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 8 commits into from
Mar 15, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
64 changes: 32 additions & 32 deletions doc/api-hidden.rst
Original file line number Diff line number Diff line change
Expand Up @@ -28,22 +28,22 @@
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

core.weighted.DatasetWeighted.obj
core.weighted.DatasetWeighted.weights
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

computation.weighted.DatasetWeighted.obj
computation.weighted.DatasetWeighted.weights

Dataset.load_store
Dataset.dump_to_store
Expand All @@ -70,22 +70,22 @@
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

core.weighted.DataArrayWeighted.obj
core.weighted.DataArrayWeighted.weights
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

computation.weighted.DataArrayWeighted.obj
computation.weighted.DataArrayWeighted.weights

core.coordinates.DataTreeCoordinates.get
core.coordinates.DataTreeCoordinates.items
Expand Down
6 changes: 3 additions & 3 deletions doc/api.rst
Original file line number Diff line number Diff line change
Expand Up @@ -1334,7 +1334,7 @@ Grouper Objects
Rolling objects
===============

.. currentmodule:: xarray.core.rolling
.. currentmodule:: xarray.computation.rolling

Dataset
-------
Expand Down Expand Up @@ -1427,7 +1427,7 @@ DataArray
Exponential rolling objects
===========================

.. currentmodule:: xarray.core.rolling_exp
.. currentmodule:: xarray.computation.rolling_exp

.. autosummary::
:toctree: generated/
Expand All @@ -1439,7 +1439,7 @@ Exponential rolling objects
Weighted objects
================

.. currentmodule:: xarray.core.weighted
.. currentmodule:: xarray.computation.weighted

Dataset
-------
Expand Down
4 changes: 2 additions & 2 deletions doc/user-guide/reshaping.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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<compute.coarsen>`),
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

Expand All @@ -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

Expand Down
24 changes: 12 additions & 12 deletions doc/whats-new.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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 <https://github.com/kmuehlbauer>`_.
- 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 <https://github.com/p4perf4ce>`_ and `Michael Niklas <https://github.com/headtr1ck>`_.
- Calling plot with kwargs ``col``, ``row`` or ``hue`` no longer squeezes dimensions passed via these arguments
Expand Down Expand Up @@ -2488,8 +2488,8 @@ New Features

- The ``zarr`` backend is now able to read NCZarr.
By `Mattia Almansi <https://github.com/malmans2>`_.
- 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 <https://github.com/cjauvin>`_ and `David Huard <https://github.com/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
Expand Down Expand Up @@ -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 <https://github.com/cjauvin>`_.
- Added a :py:func:`get_options` method to xarray's root namespace (:issue:`5698`, :pull:`5716`)
By `Pushkar Kopparla <https://github.com/pkopparla>`_.
Expand Down Expand Up @@ -3507,7 +3507,7 @@ New Features
By `Justus Magin <https://github.com/keewis>`_.
- Allow installing from git archives (:pull:`4897`).
By `Justus Magin <https://github.com/keewis>`_.
- :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 <https://github.com/spencerkclark>`_.
Expand Down Expand Up @@ -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 <https://github.com/dcherian>`_.
- 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 <https://github.com/dcherian>`_
- Add ``meta`` kwarg to :py:func:`~xarray.apply_ufunc`;
this is passed on to :py:func:`dask.array.blockwise`. (:pull:`3660`)
Expand Down Expand Up @@ -4705,7 +4705,7 @@ Bug fixes
- Fix error in concatenating unlabeled dimensions (:pull:`3362`).
By `Deepak Cherian <https://github.com/dcherian>`_.
- 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 <https://github.com/dcherian>`_.

Documentation
Expand Down Expand Up @@ -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 <https://github.com/max-sixty>`_.
- 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 <https://github.com/fujiisoup>`_.
- Handle variables with different values for ``missing_value`` and
Expand Down Expand Up @@ -6016,8 +6016,8 @@ Enhancements
NumPy. By `Stephan Hoyer <https://github.com/shoyer>`_.

- 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.
Expand Down Expand Up @@ -6791,7 +6791,7 @@ Enhancements
By `Stephan Hoyer <https://github.com/shoyer>`_ and
`Phillip J. Wolfram <https://github.com/pwolfram>`_.

- 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
Expand Down
11 changes: 5 additions & 6 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"]
Expand Down Expand Up @@ -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.*",
]
Expand Down
10 changes: 5 additions & 5 deletions xarray/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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
Expand Down
2 changes: 1 addition & 1 deletion xarray/coding/calendar_ops.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Empty file added xarray/computation/__init__.py
Empty file.
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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

Expand All @@ -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
Expand All @@ -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", ())
Expand Down
Original file line number Diff line number Diff line change
@@ -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
Expand Down
File renamed without changes.
12 changes: 7 additions & 5 deletions xarray/core/ops.py → xarray/computation/ops.py
Original file line number Diff line number Diff line change
@@ -1,13 +1,14 @@
"""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.
"""

from __future__ import annotations

import operator
from typing import Literal

import numpy as np

Expand Down Expand Up @@ -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,
Expand All @@ -157,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
Expand All @@ -172,10 +174,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,
Expand Down
4 changes: 2 additions & 2 deletions xarray/core/rolling.py → xarray/computation/rolling.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading
Loading