Merge remote-tracking branch 'pydata/main' into DAS-2066-migrate-form…

…atting-html

.github/workflows/ci-additional.yaml

@@ -127,7 +127,7 @@ jobs:
           python -m mypy --install-types --non-interactive --cobertura-xml-report mypy_report xarray/
 
       - name: Upload mypy coverage to Codecov
-        uses: codecov/codecov-action@v4.2.0
+        uses: codecov/codecov-action@v4.3.0
         with:
           file: mypy_report/cobertura.xml
           flags: mypy
@@ -181,7 +181,7 @@ jobs:
           python -m mypy --install-types --non-interactive --cobertura-xml-report mypy_report xarray/
 
       - name: Upload mypy coverage to Codecov
-        uses: codecov/codecov-action@v4.2.0
+        uses: codecov/codecov-action@v4.3.0
         with:
           file: mypy_report/cobertura.xml
           flags: mypy39
@@ -242,7 +242,7 @@ jobs:
           python -m pyright xarray/
 
       - name: Upload pyright coverage to Codecov
-        uses: codecov/codecov-action@v4.2.0
+        uses: codecov/codecov-action@v4.3.0
         with:
           file: pyright_report/cobertura.xml
           flags: pyright
@@ -301,7 +301,7 @@ jobs:
           python -m pyright xarray/
 
       - name: Upload pyright coverage to Codecov
-        uses: codecov/codecov-action@v4.2.0
+        uses: codecov/codecov-action@v4.3.0
         with:
           file: pyright_report/cobertura.xml
           flags: pyright39

.github/workflows/ci.yaml

@@ -151,7 +151,7 @@ jobs:
           path: pytest.xml
 
       - name: Upload code coverage to Codecov
-        uses: codecov/codecov-action@v4.2.0
+        uses: codecov/codecov-action@v4.3.0
         with:
           file: ./coverage.xml
           flags: unittests

.github/workflows/upstream-dev-ci.yaml

@@ -143,7 +143,7 @@ jobs:
         run: |
           python -m mypy --install-types --non-interactive --cobertura-xml-report mypy_report
       - name: Upload mypy coverage to Codecov
-        uses: codecov/codecov-action@v4.2.0
+        uses: codecov/codecov-action@v4.3.0
         with:
           file: mypy_report/cobertura.xml
           flags: mypy

ci/requirements/doc.yml

@@ -38,9 +38,9 @@ dependencies:
   - sphinx-design
   - sphinx-inline-tabs
   - sphinx>=5.0
+  - sphinxext-opengraph
+  - sphinxext-rediraffe
   - zarr>=2.10
   - pip:
-      - sphinxext-rediraffe
-      - sphinxext-opengraph
       # relative to this file. Needs to be editable to be accepted.
       - -e ../..

doc/whats-new.rst

@@ -22,6 +22,8 @@ v2024.04.0 (unreleased)
 
 New Features
 ~~~~~~~~~~~~
+- New "random" method for converting to and from 360_day calendars (:pull:`8603`).
+  By `Pascal Bourgault <https://github.com/aulemahal>`_.
 
 
 Breaking changes
@@ -37,6 +39,9 @@ Internal Changes
 - Migrates ``formatting_html`` functionality for `DataTree` into ``xarray/core`` (:pull: `8930`)
   By `Eni Awowale <https://github.com/eni-awowale>`_, `Julia Signell <https://github.com/jsignell>`_
   and `Tom Nicholas <https://github.com/TomNicholas>`_.
+- Migrates ``datatree_mapping`` functionality into ``xarray/core`` (:pull:`8948`)
+  By `Matt Savoie <https://github.com/flamingbear>`_ `Owen Littlejohns
+  <https://github.com/owenlittlejohns>` and `Tom Nicholas <https://github.com/TomNicholas>`_.
 
 
 .. _whats-new.2024.03.0:

pyproject.toml

@@ -33,6 +33,7 @@ accel = ["scipy", "bottleneck", "numbagg", "flox", "opt_einsum"]
 complete = ["xarray[accel,io,parallel,viz,dev]"]
 dev = [
   "hypothesis",
+  "mypy",
   "pre-commit",
   "pytest",
   "pytest-cov",
@@ -86,8 +87,8 @@ exclude_lines = ["pragma: no cover", "if TYPE_CHECKING"]
 [tool.mypy]
 enable_error_code = "redundant-self"
 exclude = [
-	'xarray/util/generate_.*\.py',
-	'xarray/datatree_/.*\.py',
+  'xarray/util/generate_.*\.py',
+  'xarray/datatree_/.*\.py',
 ]
 files = "xarray"
 show_error_codes = true
@@ -98,8 +99,8 @@ warn_unused_ignores = true
 
 # Ignore mypy errors for modules imported from datatree_.
 [[tool.mypy.overrides]]
-module = "xarray.datatree_.*"
 ignore_errors = true
+module = "xarray.datatree_.*"
 
 # Much of the numerical computing stack doesn't have type annotations yet.
 [[tool.mypy.overrides]]
@@ -255,6 +256,9 @@ target-version = "py39"
 # E402: module level import not at top of file
 # E501: line too long - let black worry about that
 # E731: do not assign a lambda expression, use a def
+extend-safe-fixes = [
+  "TID252", # absolute imports
+]
 ignore = [
   "E402",
   "E501",
@@ -268,9 +272,6 @@ select = [
   "I", # isort
   "UP", # Pyupgrade
 ]
-extend-safe-fixes = [
-  "TID252", # absolute imports
-]
 
 [tool.ruff.lint.per-file-ignores]
 # don't enforce absolute imports

xarray/backends/api.py

@@ -1416,8 +1416,6 @@ def save_mfdataset(
         these locations will be overwritten.
     format : {"NETCDF4", "NETCDF4_CLASSIC", "NETCDF3_64BIT", \
               "NETCDF3_CLASSIC"}, optional
-    **kwargs : additional arguments are passed along to ``to_netcdf``
-
         File format for the resulting netCDF file:
 
         * NETCDF4: Data is stored in an HDF5 file, using netCDF4 API
@@ -1449,10 +1447,11 @@ def save_mfdataset(
     compute : bool
         If true compute immediately, otherwise return a
         ``dask.delayed.Delayed`` object that can be computed later.
+    **kwargs : dict, optional
+        Additional arguments are passed along to ``to_netcdf``.
 
     Examples
     --------
-
     Save a dataset into one netCDF per year of data:
 
     >>> ds = xr.Dataset(

xarray/coding/calendar_ops.py

@@ -64,7 +64,7 @@ def convert_calendar(
       The target calendar name.
     dim : str
       Name of the time coordinate in the input DataArray or Dataset.
-    align_on : {None, 'date', 'year'}
+    align_on : {None, 'date', 'year', 'random'}
       Must be specified when either the source or target is a `"360_day"`
       calendar; ignored otherwise. See Notes.
     missing : any, optional
@@ -143,6 +143,16 @@ def convert_calendar(
       will be dropped as there are no equivalent dates in a standard calendar.
 
       This option is best used with data on a frequency coarser than daily.
+
+    "random"
+      Similar to "year", each day of year of the source is mapped to another day of year
+      of the target. However, instead of having always the same missing days according
+      the source and target years, here 5 days are chosen randomly, one for each fifth
+      of the year. However, February 29th is always missing when converting to a leap year,
+      or its value is dropped when converting from a leap year. This is similar to the method
+      used in the LOCA dataset (see Pierce, Cayan, and Thrasher (2014). doi:10.1175/JHM-D-14-0082.1).
+
+      This option is best used on daily data.
     """
     from xarray.core.dataarray import DataArray
 
@@ -174,14 +184,20 @@ def convert_calendar(
 
     out = obj.copy()
 
-    if align_on == "year":
+    if align_on in ["year", "random"]:
         # Special case for conversion involving 360_day calendar
-        # Instead of translating dates directly, this tries to keep the position within a year similar.
-
-        new_doy = time.groupby(f"{dim}.year").map(
-            _interpolate_day_of_year, target_calendar=calendar, use_cftime=use_cftime
-        )
-
+        if align_on == "year":
+            # Instead of translating dates directly, this tries to keep the position within a year similar.
+            new_doy = time.groupby(f"{dim}.year").map(
+                _interpolate_day_of_year,
+                target_calendar=calendar,
+                use_cftime=use_cftime,
+            )
+        elif align_on == "random":
+            # The 5 days to remove are randomly chosen, one for each of the five 72-days periods of the year.
+            new_doy = time.groupby(f"{dim}.year").map(
+                _random_day_of_year, target_calendar=calendar, use_cftime=use_cftime
+            )
         # Convert the source datetimes, but override the day of year with our new day of years.
         out[dim] = DataArray(
             [
@@ -229,6 +245,27 @@ def _interpolate_day_of_year(time, target_calendar, use_cftime):
     ).astype(int)
 
 
+def _random_day_of_year(time, target_calendar, use_cftime):
+    """Return a day of year in the new calendar.
+
+    Removes Feb 29th and five other days chosen randomly within five sections of 72 days.
+    """
+    year = int(time.dt.year[0])
+    source_calendar = time.dt.calendar
+    new_doy = np.arange(360) + 1
+    rm_idx = np.random.default_rng().integers(0, 72, 5) + 72 * np.arange(5)
+    if source_calendar == "360_day":
+        for idx in rm_idx:
+            new_doy[idx + 1 :] = new_doy[idx + 1 :] + 1
+        if _days_in_year(year, target_calendar, use_cftime) == 366:
+            new_doy[new_doy >= 60] = new_doy[new_doy >= 60] + 1
+    elif target_calendar == "360_day":
+        new_doy = np.insert(new_doy, rm_idx - np.arange(5), -1)
+        if _days_in_year(year, source_calendar, use_cftime) == 366:
+            new_doy = np.insert(new_doy, 60, -1)
+    return new_doy[time.dt.dayofyear - 1]
+
+
 def _convert_to_new_calendar_with_new_day_of_year(
     date, day_of_year, calendar, use_cftime
 ):

xarray/core/datatree.py

@@ -18,6 +18,11 @@
 from xarray.core.coordinates import DatasetCoordinates
 from xarray.core.dataarray import DataArray
 from xarray.core.dataset import Dataset, DataVariables
+from xarray.core.datatree_mapping import (
+    TreeIsomorphismError,
+    check_isomorphic,
+    map_over_subtree,
+)
 from xarray.core.formatting_html import (
     datatree_repr as datatree_repr_html,
 )
@@ -36,11 +41,6 @@
 from xarray.core.variable import Variable
 from xarray.datatree_.datatree.common import TreeAttrAccessMixin
 from xarray.datatree_.datatree.formatting import datatree_repr
-from xarray.datatree_.datatree.mapping import (
-    TreeIsomorphismError,
-    check_isomorphic,
-    map_over_subtree,
-)
 from xarray.datatree_.datatree.ops import (
     DataTreeArithmeticMixin,
     MappedDatasetMethodsMixin,

xarray/datatree_/datatree/mapping.py → xarray/core/datatree_mapping.py

@@ -4,10 +4,9 @@
 import sys
 from itertools import repeat
 from textwrap import dedent
-from typing import TYPE_CHECKING, Callable, Tuple
+from typing import TYPE_CHECKING, Callable
 
 from xarray import DataArray, Dataset
-
 from xarray.core.iterators import LevelOrderIter
 from xarray.core.treenode import NodePath, TreeNode
 
@@ -84,14 +83,13 @@ def diff_treestructure(a: DataTree, b: DataTree, require_names_equal: bool) -> s
     for node_a, node_b in zip(LevelOrderIter(a), LevelOrderIter(b)):
         path_a, path_b = node_a.path, node_b.path
 
-        if require_names_equal:
-            if node_a.name != node_b.name:
-                diff = dedent(
-                    f"""\
+        if require_names_equal and node_a.name != node_b.name:
+            diff = dedent(
+                f"""\
                 Node '{path_a}' in the left object has name '{node_a.name}'
                 Node '{path_b}' in the right object has name '{node_b.name}'"""
-                )
-                return diff
+            )
+            return diff
 
         if len(node_a.children) != len(node_b.children):
             diff = dedent(
@@ -125,7 +123,7 @@ def map_over_subtree(func: Callable) -> Callable:
     func : callable
         Function to apply to datasets with signature:
 
-        `func(*args, **kwargs) -> Union[Dataset, Iterable[Dataset]]`.
+        `func(*args, **kwargs) -> Union[DataTree, Iterable[DataTree]]`.
 
         (i.e. func must accept at least one Dataset and return at least one Dataset.)
         Function will not be applied to any nodes without datasets.
@@ -154,7 +152,7 @@ def map_over_subtree(func: Callable) -> Callable:
     # TODO inspect function to work out immediately if the wrong number of arguments were passed for it?
 
     @functools.wraps(func)
-    def _map_over_subtree(*args, **kwargs) -> DataTree | Tuple[DataTree, ...]:
+    def _map_over_subtree(*args, **kwargs) -> DataTree | tuple[DataTree, ...]:
         """Internal function which maps func over every node in tree, returning a tree of the results."""
         from xarray.core.datatree import DataTree
 
@@ -259,19 +257,18 @@ def _map_over_subtree(*args, **kwargs) -> DataTree | Tuple[DataTree, ...]:
     return _map_over_subtree
 
 
-def _handle_errors_with_path_context(path):
+def _handle_errors_with_path_context(path: str):
     """Wraps given function so that if it fails it also raises path to node on which it failed."""
 
     def decorator(func):
         def wrapper(*args, **kwargs):
             try:
                 return func(*args, **kwargs)
             except Exception as e:
-                if sys.version_info >= (3, 11):
-                    # Add the context information to the error message
-                    e.add_note(
-                        f"Raised whilst mapping function over node with path {path}"
-                    )
+                # Add the context information to the error message
+                add_note(
+                    e, f"Raised whilst mapping function over node with path {path}"
+                )
                 raise
 
         return wrapper
@@ -287,7 +284,9 @@ def add_note(err: BaseException, msg: str) -> None:
         err.add_note(msg)
 
 
-def _check_single_set_return_values(path_to_node, obj):
+def _check_single_set_return_values(
+    path_to_node: str, obj: Dataset | DataArray | tuple[Dataset | DataArray]
+):
     """Check types returned from single evaluation of func, and return number of return values received from func."""
     if isinstance(obj, (Dataset, DataArray)):
         return 1

xarray/core/iterators.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-from collections import abc
 from collections.abc import Iterator
 from typing import Callable
 
@@ -9,7 +8,7 @@
 """These iterators are copied from anytree.iterators, with minor modifications."""
 
 
-class LevelOrderIter(abc.Iterator):
+class LevelOrderIter(Iterator):
     """Iterate over tree applying level-order strategy starting at `node`.
     This is the iterator used by `DataTree` to traverse nodes.
 

xarray/datatree_/datatree/__init__.py

@@ -1,11 +1,8 @@
 # import public API
-from .mapping import TreeIsomorphismError, map_over_subtree
 from xarray.core.treenode import InvalidTreeError, NotFoundInTreeError
 
 
 __all__ = (
-    "TreeIsomorphismError",
     "InvalidTreeError",
     "NotFoundInTreeError",
-    "map_over_subtree",
 )

xarray/datatree_/datatree/formatting.py

@@ -2,7 +2,7 @@
 
 from xarray.core.formatting import _compat_to_str, diff_dataset_repr
 
-from xarray.datatree_.datatree.mapping import diff_treestructure
+from xarray.core.datatree_mapping import diff_treestructure
 from xarray.datatree_.datatree.render import RenderTree
 
 if TYPE_CHECKING:

xarray/datatree_/datatree/ops.py

@@ -2,7 +2,7 @@
 
 from xarray import Dataset
 
-from .mapping import map_over_subtree
+from xarray.core.datatree_mapping import map_over_subtree
 
 """
 Module which specifies the subset of xarray.Dataset's API which we wish to copy onto DataTree.