Switch to new iris >= 3.10.0 API (#2500)

Co-authored-by: Valeriu Predoi <[email protected]>

environment.yml

@@ -18,7 +18,7 @@ dependencies:
   - fire
   - geopy
   - humanfriendly
-  - iris >=3.9.0
+  - iris >=3.10.0
   - iris-esmf-regrid >=0.10.0  # github.com/SciTools-incubator/iris-esmf-regrid/pull/342
   - iris-grib
   - isodate

esmvalcore/cmor/_fixes/icon/_base_fixes.py

@@ -16,7 +16,7 @@
 import requests
 from iris import NameConstraint
 from iris.cube import Cube, CubeList
-from iris.experimental.ugrid import Connectivity, Mesh
+from iris.mesh import Connectivity, MeshXY
 
 from esmvalcore.cmor._fixes.native_datasets import NativeDatasetFix
 from esmvalcore.local import _get_data_sources
@@ -38,24 +38,24 @@ def __init__(self, *args, **kwargs):
         self._horizontal_grids = {}
         self._meshes = {}
 
-    def _create_mesh(self, cube):
+    def _create_mesh(self, cube: Cube) -> MeshXY:
         """Create mesh from horizontal grid file.
 
         Note
         ----
-        This functions creates a new :class:`iris.experimental.ugrid.Mesh` from
-        the ``clat`` (already present in the cube), ``clon`` (already present
-        in the cube), ``vertex_index``, ``vertex_of_cell``, ``vlat``, and
-        ``vlon`` variables of the horizontal grid file.
-
-        We do not use :func:`iris.experimental.ugrid.Mesh.from_coords` with the
-        existing latitude and longitude coordinates here because this would
-        produce lots of duplicated entries for the node coordinates. The reason
-        for this is that the node coordinates are constructed from the bounds;
-        since each node is contained 6 times in the bounds array (each node is
-        shared by 6 neighboring cells) the number of nodes is 6 times higher
-        with :func:`iris.experimental.ugrid.Mesh.from_coords` compared to using
-        the information already present in the horizontal grid file.
+        This functions creates a new :class:`iris.mesh.MeshXY` from the
+        ``clat`` (already present in the cube), ``clon`` (already present in
+        the cube), ``vertex_index``, ``vertex_of_cell``, ``vlat``, and ``vlon``
+        variables of the horizontal grid file.
+
+        We do not use :func:`iris.mesh.MeshXY.from_coords` with the existing
+        latitude and longitude coordinates here because this would produce lots
+        of duplicated entries for the node coordinates. The reason for this is
+        that the node coordinates are constructed from the bounds; since each
+        node is contained 6 times in the bounds array (each node is shared by 6
+        neighboring cells) the number of nodes is 6 times higher with
+        :func:`iris.mesh.MeshXY.from_coords` compared to using the information
+        already present in the horizontal grid file.
 
         """
         horizontal_grid = self.get_horizontal_grid(cube)
@@ -91,7 +91,7 @@ def _create_mesh(self, cube):
         if not np.allclose(
                 face_lat.bounds,
                 node_lat.points[conn_node_inds],
-                **close_kwargs,
+                **close_kwargs,  # type: ignore
         ):
             logger.warning(
                 "Latitude bounds of the face coordinate ('clat_vertices' in "
@@ -108,7 +108,7 @@ def _create_mesh(self, cube):
         # differ by 360°, which is also okay.
         face_lon_bounds_to_check = face_lon.bounds % 360
         node_lon_conn_to_check = node_lon.points[conn_node_inds] % 360
-        idx_notclose = ~np.isclose(
+        idx_notclose = ~np.isclose(  # type: ignore
             face_lon_bounds_to_check,
             node_lon_conn_to_check,
             **close_kwargs,
@@ -131,7 +131,7 @@ def _create_mesh(self, cube):
             start_index=start_index,
             location_axis=0,
         )
-        mesh = Mesh(
+        mesh = MeshXY(
             topology_dimension=2,
             node_coords_and_axes=[(node_lat, 'y'), (node_lon, 'x')],
             connectivities=[connectivity],
@@ -429,8 +429,8 @@ def get_mesh(self, cube):
 
         Returns
         -------
-        iris.experimental.ugrid.Mesh
-            Mesh.
+        iris.mesh.MeshXY
+            Mesh of the cube.
 
         Raises
         ------

esmvalcore/preprocessor/_area.py

@@ -6,7 +6,6 @@
 from __future__ import annotations
 
 import logging
-import warnings
 from pathlib import Path
 from typing import TYPE_CHECKING, Iterable, Literal, Optional
 
@@ -16,17 +15,15 @@
 import shapely
 import shapely.ops
 from dask import array as da
-from iris.coords import AuxCoord, CellMeasure
+from iris.coords import AuxCoord
 from iris.cube import Cube, CubeList
-from iris.exceptions import CoordinateMultiDimError, CoordinateNotFoundError
+from iris.exceptions import CoordinateNotFoundError
 
-from esmvalcore.iris_helpers import has_regular_grid
-from esmvalcore.preprocessor._regrid import broadcast_to_shape
 from esmvalcore.preprocessor._shared import (
     get_iris_aggregator,
     get_normalized_cube,
-    guess_bounds,
     preserve_float_dtype,
+    try_adding_calculated_cell_area,
     update_weights_kwargs,
 )
 from esmvalcore.preprocessor._supplementary_vars import (
@@ -291,112 +288,6 @@ def meridional_statistics(
     return result
 
 
-def compute_area_weights(cube):
-    """Compute area weights."""
-    with warnings.catch_warnings(record=True) as caught_warnings:
-        warnings.filterwarnings(
-            'always',
-            message="Using DEFAULT_SPHERICAL_EARTH_RADIUS.",
-            category=UserWarning,
-            module='iris.analysis.cartography',
-        )
-        # TODO: replace the following line with
-        # weights = iris.analysis.cartography.area_weights(
-        #     cube, compute=not cube.has_lazy_data()
-        # )
-        # once https://github.com/SciTools/iris/pull/5658 is available
-        weights = _get_area_weights(cube)
-
-        for warning in caught_warnings:
-            logger.debug(
-                "%s while computing area weights of the following cube:\n%s",
-                warning.message, cube)
-    return weights
-
-
-def _get_area_weights(cube: Cube) -> np.ndarray | da.Array:
-    """Get area weights.
-
-    For non-lazy data, simply use the according iris function. For lazy data,
-    calculate area weights for a single lat-lon slice and broadcast it to the
-    correct shape.
-
-    Note
-    ----
-    This is a temporary workaround to get lazy area weights. Can be removed
-    once https://github.com/SciTools/iris/pull/5658 is available.
-
-    """
-    if not cube.has_lazy_data():
-        return iris.analysis.cartography.area_weights(cube)
-
-    lat_lon_dims = sorted(
-        tuple(set(cube.coord_dims('latitude') + cube.coord_dims('longitude')))
-    )
-    lat_lon_slice = next(cube.slices(['latitude', 'longitude'], ordered=False))
-    weights_2d = iris.analysis.cartography.area_weights(lat_lon_slice)
-    weights = broadcast_to_shape(
-        da.array(weights_2d),
-        cube.shape,
-        lat_lon_dims,
-        chunks=cube.lazy_data().chunks,
-    )
-    return weights
-
-
-def _try_adding_calculated_cell_area(cube: Cube) -> None:
-    """Try to add calculated cell measure 'cell_area' to cube (in-place)."""
-    if cube.cell_measures('cell_area'):
-        return
-
-    logger.debug(
-        "Found no cell measure 'cell_area' in cube %s. Check availability of "
-        "supplementary variables",
-        cube.summary(shorten=True),
-    )
-    logger.debug("Attempting to calculate grid cell area")
-
-    rotated_pole_grid = all([
-        cube.coord('latitude').core_points().ndim == 2,
-        cube.coord('longitude').core_points().ndim == 2,
-        cube.coords('grid_latitude'),
-        cube.coords('grid_longitude'),
-    ])
-
-    # For regular grids, calculate grid cell areas with iris function
-    if has_regular_grid(cube):
-        cube = guess_bounds(cube, ['latitude', 'longitude'])
-        logger.debug("Calculating grid cell areas for regular grid")
-        cell_areas = compute_area_weights(cube)
-
-    # For rotated pole grids, use grid_latitude and grid_longitude to calculate
-    # grid cell areas
-    elif rotated_pole_grid:
-        cube = guess_bounds(cube, ['grid_latitude', 'grid_longitude'])
-        cube_tmp = cube.copy()
-        cube_tmp.remove_coord('latitude')
-        cube_tmp.coord('grid_latitude').rename('latitude')
-        cube_tmp.remove_coord('longitude')
-        cube_tmp.coord('grid_longitude').rename('longitude')
-        logger.debug("Calculating grid cell areas for rotated pole grid")
-        cell_areas = compute_area_weights(cube_tmp)
-
-    # For all other cases, grid cell areas cannot be calculated
-    else:
-        logger.error(
-            "Supplementary variables are needed to calculate grid cell "
-            "areas for irregular or unstructured grid of cube %s",
-            cube.summary(shorten=True),
-        )
-        raise CoordinateMultiDimError(cube.coord('latitude'))
-
-    # Add new cell measure
-    cell_measure = CellMeasure(
-        cell_areas, standard_name='cell_area', units='m2', measure='area',
-    )
-    cube.add_cell_measure(cell_measure, np.arange(cube.ndim))
-
-
 @register_supplementaries(
     variables=['areacella', 'areacello'],
     required='prefer_at_least_one',
@@ -454,7 +345,7 @@ def area_statistics(
     # Get aggregator and correct kwargs (incl. weights)
     (agg, agg_kwargs) = get_iris_aggregator(operator, **operator_kwargs)
     agg_kwargs = update_weights_kwargs(
-        agg, agg_kwargs, 'cell_area', cube, _try_adding_calculated_cell_area
+        agg, agg_kwargs, 'cell_area', cube, try_adding_calculated_cell_area
     )
 
     result = cube.collapsed(['latitude', 'longitude'], agg, **agg_kwargs)

esmvalcore/preprocessor/_regrid.py

@@ -21,6 +21,7 @@
 from geopy.geocoders import Nominatim
 from iris.analysis import AreaWeighted, Linear, Nearest
 from iris.cube import Cube
+from iris.util import broadcast_to_shape
 
 from esmvalcore.cmor._fixes.shared import (
     add_altitude_from_plev,
@@ -30,7 +31,6 @@
 from esmvalcore.exceptions import ESMValCoreDeprecationWarning
 from esmvalcore.iris_helpers import has_irregular_grid, has_unstructured_grid
 from esmvalcore.preprocessor._shared import (
-    broadcast_to_shape,
     get_array_module,
     preserve_float_dtype,
 )

esmvalcore/preprocessor/_shared.py

@@ -10,7 +10,7 @@
 import warnings
 from collections import defaultdict
 from collections.abc import Callable, Iterable
-from functools import partial, wraps
+from functools import wraps
 from typing import Any, Literal, Optional
 
 import dask.array as da
@@ -19,6 +19,7 @@
 from iris.coords import CellMeasure, Coord, DimCoord
 from iris.cube import Cube
 from iris.exceptions import CoordinateMultiDimError, CoordinateNotFoundError
+from iris.util import broadcast_to_shape
 
 from esmvalcore.exceptions import ESMValCoreDeprecationWarning
 from esmvalcore.iris_helpers import has_regular_grid
@@ -318,52 +319,6 @@ def get_array_module(*args):
     return np
 
 
-def broadcast_to_shape(array, shape, dim_map, chunks=None):
-    """Copy of `iris.util.broadcast_to_shape` that allows specifying chunks."""
-    if isinstance(array, da.Array):
-        if chunks is not None:
-            chunks = list(chunks)
-            for src_idx, tgt_idx in enumerate(dim_map):
-                # Only use the specified chunks along new dimensions or on
-                # dimensions that have size 1 in the source array.
-                if array.shape[src_idx] != 1:
-                    chunks[tgt_idx] = array.chunks[src_idx]
-        broadcast = partial(da.broadcast_to, shape=shape, chunks=chunks)
-    else:
-        broadcast = partial(np.broadcast_to, shape=shape)
-
-    n_orig_dims = len(array.shape)
-    n_new_dims = len(shape) - n_orig_dims
-    array = array.reshape(array.shape + (1,) * n_new_dims)
-
-    # Get dims in required order.
-    array = np.moveaxis(array, range(n_orig_dims), dim_map)
-    new_array = broadcast(array)
-
-    if np.ma.isMA(array):
-        # broadcast_to strips masks so we need to handle them explicitly.
-        mask = np.ma.getmask(array)
-        if mask is np.ma.nomask:
-            new_mask = np.ma.nomask
-        else:
-            new_mask = broadcast(mask)
-        new_array = np.ma.array(new_array, mask=new_mask)
-
-    elif _is_lazy_masked_data(array):
-        # broadcast_to strips masks so we need to handle them explicitly.
-        mask = da.ma.getmaskarray(array)
-        new_mask = broadcast(mask)
-        new_array = da.ma.masked_array(new_array, new_mask)
-
-    return new_array
-
-
-def _is_lazy_masked_data(array):
-    """Similar to `iris._lazy_data.is_lazy_masked_data`."""
-    return isinstance(array, da.Array) and isinstance(
-        da.utils.meta_from_array(array), np.ma.MaskedArray)
-
-
 def get_weights(
     cube: Cube,
     coords: Iterable[Coord] | Iterable[str],
@@ -505,50 +460,18 @@ def _compute_area_weights(cube):
             category=UserWarning,
             module='iris.analysis.cartography',
         )
-        # TODO: replace the following line with
-        # weights = iris.analysis.cartography.area_weights(
-        #     cube, compute=not cube.has_lazy_data()
-        # )
-        # once https://github.com/SciTools/iris/pull/5658 is available
-        weights = _get_area_weights(cube)
-
+        if cube.has_lazy_data():
+            kwargs = {'compute': False, 'chunks': cube.lazy_data().chunks}
+        else:
+            kwargs = {'compute': True}
+        weights = iris.analysis.cartography.area_weights(cube, **kwargs)
         for warning in caught_warnings:
             logger.debug(
                 "%s while computing area weights of the following cube:\n%s",
                 warning.message, cube)
     return weights
 
 
-def _get_area_weights(cube: Cube) -> np.ndarray | da.Array:
-    """Get area weights.
-
-    For non-lazy data, simply use the according iris function. For lazy data,
-    calculate area weights for a single lat-lon slice and broadcast it to the
-    correct shape.
-
-    Note
-    ----
-    This is a temporary workaround to get lazy area weights. Can be removed
-    once https://github.com/SciTools/iris/pull/5658 is available.
-
-    """
-    if not cube.has_lazy_data():
-        return iris.analysis.cartography.area_weights(cube)
-
-    lat_lon_dims = sorted(
-        tuple(set(cube.coord_dims('latitude') + cube.coord_dims('longitude')))
-    )
-    lat_lon_slice = next(cube.slices(['latitude', 'longitude'], ordered=False))
-    weights_2d = iris.analysis.cartography.area_weights(lat_lon_slice)
-    weights = broadcast_to_shape(
-        da.array(weights_2d),
-        cube.shape,
-        lat_lon_dims,
-        chunks=cube.lazy_data().chunks,
-    )
-    return weights
-
-
 def get_all_coords(
     cube: Cube,
     coords: Iterable[Coord] | Iterable[str] | None,

esmvalcore/preprocessor/_volume.py

@@ -14,9 +14,9 @@
 import numpy as np
 from iris.coords import AuxCoord, CellMeasure
 from iris.cube import Cube
+from iris.util import broadcast_to_shape
 
 from ._shared import (
-    broadcast_to_shape,
     get_iris_aggregator,
     get_normalized_cube,
     preserve_float_dtype,