More lazy fixes and preprocessing functions (#2325)

Co-authored-by: Manuel Schlund <[email protected]>

esmvalcore/cmor/_fixes/cesm/cesm2.py

@@ -78,6 +78,6 @@ def _fix_time(self, cube):
 
         # Fix time coordinate
         time_coord = cube.coord('time')
-        if time_coord.bounds is not None:
-            time_coord.points = time_coord.bounds.mean(axis=-1)
+        if time_coord.has_bounds():
+            time_coord.points = time_coord.core_bounds().mean(axis=-1)
         self.fix_regular_time(cube, coord=time_coord)

esmvalcore/cmor/_fixes/cmip5/hadgem2_cc.py

@@ -1,34 +1,9 @@
-
 """Fix HadGEM2_CC."""
-import numpy as np
-
 from ..fix import Fix
+from .hadgem2_es import AllVars as BaseAllVars
 
 
-class AllVars(Fix):
-    """Fix common errors for all vars."""
-
-    def fix_metadata(self, cubes):
-        """
-        Fix latitude.
-
-        Parameters
-        ----------
-        cube: iris.cube.CubeList
-
-        Returns
-        -------
-        iris.cube.CubeList
-
-        """
-        for cube in cubes:
-            lats = cube.coords('latitude')
-            if lats:
-                lat = cube.coord('latitude')
-                lat.points = np.clip(lat.points, -90., 90.)
-                lat.bounds = np.clip(lat.bounds, -90., 90.)
-
-        return cubes
+AllVars = BaseAllVars
 
 
 class O2(Fix):

esmvalcore/cmor/_fixes/cmip5/hadgem2_es.py

@@ -25,10 +25,10 @@ def fix_metadata(self, cubes):
             lats = cube.coords('latitude')
             if lats:
                 lat = cube.coord('latitude')
-                lat.points = np.clip(lat.points, -90., 90.)
+                lat.points = np.clip(lat.core_points(), -90., 90.)
                 if not lat.has_bounds():
                     lat.guess_bounds()
-                lat.bounds = np.clip(lat.bounds, -90., 90.)
+                lat.bounds = np.clip(lat.core_bounds(), -90., 90.)
 
         return cubes
 

esmvalcore/cmor/_fixes/cmip5/miroc_esm.py

@@ -100,7 +100,7 @@ def fix_metadata(self, cubes):
                                            calendar='gregorian')
                 if cube.coord('time').units != expected_time_units:
                     continue
-                if cube.coord('time').bounds is None:
+                if not cube.coord('time').has_bounds():
                     continue
 
                 # Only apply fix if there is a year < 1 in the first element

esmvalcore/cmor/_fixes/cmip6/access_esm1_5.py

@@ -102,9 +102,9 @@ def fix_metadata(self, cubes):
         """
         cube = self.get_cube_from_list(cubes)
         cube.coord('air_pressure').points = \
-            np.round(cube.coord('air_pressure').points, 0)
+            np.round(cube.coord('air_pressure').core_points(), 0)
         cube.coord('air_pressure').bounds = \
-            np.round(cube.coord('air_pressure').bounds, 0)
+            np.round(cube.coord('air_pressure').core_bounds(), 0)
         return cubes
 
 

esmvalcore/cmor/_fixes/cmip6/cesm2.py

@@ -142,16 +142,12 @@ def fix_metadata(self, cubes):
 
         """
         for cube in cubes:
-            latitude = cube.coord('latitude')
-            if latitude.bounds is None:
-                latitude.guess_bounds()
-            latitude.bounds = latitude.bounds.astype(np.float64)
-            latitude.bounds = np.round(latitude.bounds, 4)
-            longitude = cube.coord('longitude')
-            if longitude.bounds is None:
-                longitude.guess_bounds()
-            longitude.bounds = longitude.bounds.astype(np.float64)
-            longitude.bounds = np.round(longitude.bounds, 4)
+            for coord_name in ['latitude', 'longitude']:
+                coord = cube.coord(coord_name)
+                if not coord.has_bounds():
+                    coord.guess_bounds()
+                coord.bounds = np.round(coord.core_bounds().astype(np.float64),
+                                        4)
 
         return cubes
 

esmvalcore/cmor/_fixes/cmip6/ec_earth3.py

@@ -68,7 +68,7 @@ def fix_metadata(self, cubes):
 
         for cube in cubes:
             latitude = cube.coord('latitude')
-            latitude.points = np.round(latitude.points, 8)
-            latitude.bounds = np.round(latitude.bounds, 8)
+            latitude.points = np.round(latitude.core_points(), 8)
+            latitude.bounds = np.round(latitude.core_bounds(), 8)
 
         return cubes

esmvalcore/cmor/_fixes/cmip6/ec_earth3_veg.py

@@ -72,7 +72,7 @@ def fix_metadata(self, cubes):
 
         for cube in cubes:
             latitude = cube.coord('latitude')
-            latitude.points = np.round(latitude.points, 8)
-            latitude.bounds = np.round(latitude.bounds, 8)
+            latitude.points = np.round(latitude.core_points(), 8)
+            latitude.bounds = np.round(latitude.core_bounds(), 8)
 
         return cubes

esmvalcore/cmor/_fixes/cmip6/giss_e2_1_g.py

@@ -1,6 +1,4 @@
 """Fixes for GISS-E2-1-G model."""
-import dask.array as da
-
 from ..common import ClFixHybridPressureCoord
 from ..fix import Fix
 
@@ -22,7 +20,7 @@ def fix_metadata(self, cubes):
         """
         for cube in cubes:
             if (cube.units == 'degC'
-                    and da.any(cube.core_data().ravel()[:1000] > 100.)):
+                    and cube.core_data().ravel()[:1000].max() > 100.):
                 cube.units = 'K'
                 cube.convert_units(self.vardef.units)
         return cubes

esmvalcore/cmor/_fixes/cmip6/miroc6.py

@@ -28,18 +28,12 @@ def fix_metadata(self, cubes):
         iris.cube.CubeList
         """
         for cube in cubes:
-            latitude = cube.coord('latitude')
-            if latitude.bounds is None:
-                latitude.guess_bounds()
-            latitude.points = latitude.points.astype(np.float32).astype(
-                np.float64)
-            latitude.bounds = latitude.bounds.astype(np.float32).astype(
-                np.float64)
-            longitude = cube.coord('longitude')
-            if longitude.bounds is None:
-                longitude.guess_bounds()
-            longitude.points = longitude.points.astype(np.float32).astype(
-                np.float64)
-            longitude.bounds = longitude.bounds.astype(np.float32).astype(
-                np.float64)
+            for coord_name in ['latitude', 'longitude']:
+                coord = cube.coord(coord_name)
+                coord.points = coord.core_points().astype(np.float32).astype(
+                    np.float64)
+                if not coord.has_bounds():
+                    coord.guess_bounds()
+                coord.bounds = coord.core_bounds().astype(np.float32).astype(
+                    np.float64)
         return cubes

esmvalcore/cmor/_fixes/common.py

@@ -83,7 +83,7 @@ def fix_metadata(self, cubes):
         # This was originally done by iris, but it has to be repeated since
         # a has bounds now
         ap_coord = cube.coord(var_name='ap')
-        if ap_coord.bounds is None:
+        if not ap_coord.has_bounds():
             cube.remove_coord(ap_coord)
             a_coord = cube.coord(var_name='a')
             p0_coord = cube.coord(var_name='p0')

esmvalcore/cmor/_fixes/cordex/cordex_fixes.py

@@ -100,7 +100,7 @@ def fix_metadata(self, cubes):
                 if coord.dtype in ['>f8', '>f4']:
                     coord.points = coord.core_points().astype(
                         np.float64, casting='same_kind')
-                    if coord.bounds is not None:
+                    if coord.has_bounds():
                         coord.bounds = coord.core_bounds().astype(
                             np.float64, casting='same_kind')
         return cubes

esmvalcore/cmor/_fixes/emac/emac.py

@@ -174,7 +174,7 @@ def _fix_alevel(cube, cubes):
         for coord in (ap_coord, b_coord, ps_coord):
             coord.points = coord.core_points().astype(
                 float, casting='same_kind')
-            if coord.bounds is not None:
+            if coord.has_bounds():
                 coord.bounds = coord.core_bounds().astype(
                     float, casting='same_kind')
 

esmvalcore/cmor/_fixes/icon/_base_fixes.py

@@ -450,6 +450,6 @@ def _load_cubes(path):
     @staticmethod
     def _set_range_in_0_360(lon_coord):
         """Convert longitude coordinate to [0, 360]."""
-        lon_coord.points = (lon_coord.points + 360.0) % 360.0
-        if lon_coord.bounds is not None:
-            lon_coord.bounds = (lon_coord.bounds + 360.0) % 360.0
+        lon_coord.points = (lon_coord.core_points() + 360.0) % 360.0
+        if lon_coord.has_bounds():
+            lon_coord.bounds = (lon_coord.core_bounds() + 360.0) % 360.0

esmvalcore/cmor/_fixes/native6/era5.py

@@ -7,9 +7,9 @@
 
 from esmvalcore.iris_helpers import date2num
 
+from ...table import CMOR_TABLES
 from ..fix import Fix
 from ..shared import add_scalar_height_coord
-from ...table import CMOR_TABLES
 
 logger = logging.getLogger(__name__)
 
@@ -382,7 +382,7 @@ def _fix_coordinates(self, cube):
             coord.var_name = coord_def.out_name
             coord.long_name = coord_def.long_name
             coord.points = coord.core_points().astype('float64')
-            if (coord.bounds is None and len(coord.points) > 1
+            if (not coord.has_bounds() and len(coord.core_points()) > 1
                     and coord_def.must_have_bounds == "yes"):
                 coord.guess_bounds()
 

esmvalcore/cmor/_fixes/native6/mswep.py

@@ -113,7 +113,7 @@ def _fix_bounds(self, cube):
 
             coord = cube.coord(axis=coord_def.axis)
 
-            if coord.bounds is None:
+            if not coord.has_bounds():
                 coord.guess_bounds()
 
     def _fix_names(self, cube):

esmvalcore/cmor/_fixes/shared.py

@@ -359,7 +359,7 @@ def fix_bounds(cube, cubes, coord_var_names):
     """
     for coord_var_name in coord_var_names:
         coord = cube.coord(var_name=coord_var_name)
-        if coord.bounds is not None:
+        if coord.has_bounds():
             continue
         bounds_cube = get_bounds_cube(cubes, coord_var_name)
         cube.coord(var_name=coord_var_name).bounds = bounds_cube.core_data()
@@ -398,10 +398,9 @@ def round_coordinates(cubes, decimals=5, coord_names=None):
         else:
             coords = [cube.coord(c) for c in coord_names if cube.coords(c)]
         for coord in coords:
-            coord.points = da.round(da.asarray(coord.core_points()), decimals)
-            if coord.bounds is not None:
-                coord.bounds = da.round(da.asarray(coord.core_bounds()),
-                                        decimals)
+            coord.points = np.round(coord.core_points(), decimals)
+            if coord.has_bounds():
+                coord.bounds = np.round(coord.core_bounds(), decimals)
     return cubes
 
 

esmvalcore/preprocessor/_derive/sm.py

@@ -29,7 +29,7 @@ def calculate(cubes):
         """
         mrsos_cube = cubes.extract_cube(NameConstraint(var_name='mrsos'))
 
-        depth = mrsos_cube.coord('depth').bounds.astype(np.float32)
+        depth = mrsos_cube.coord('depth').core_bounds().astype(np.float32)
         layer_thickness = depth[..., 1] - depth[..., 0]
 
         sm_cube = mrsos_cube / layer_thickness / 998.2

esmvalcore/preprocessor/_units.py

@@ -2,8 +2,11 @@
 
 Allows for unit conversions.
 """
+from __future__ import annotations
+
 import logging
 
+import dask.array as da
 import iris
 import numpy as np
 from cf_units import Unit
@@ -120,18 +123,21 @@ def convert_units(cube, units):
     return cube
 
 
-def accumulate_coordinate(cube, coordinate):
+def accumulate_coordinate(
+    cube: iris.cube.Cube,
+    coordinate: str | iris.coords.DimCoord | iris.coords.AuxCoord
+) -> iris.cube.Cube:
     """Weight data using the bounds from a given coordinate.
 
     The resulting cube will then have units given by
     ``cube_units * coordinate_units``.
 
     Parameters
     ----------
-    cube : iris.cube.Cube
+    cube:
         Data cube for the flux
 
-    coordinate: str
+    coordinate:
         Name of the coordinate that will be used as weights.
 
     Returns
@@ -158,8 +164,9 @@ def accumulate_coordinate(cube, coordinate):
         raise NotImplementedError(
             f'Multidimensional coordinate {coord} not supported.')
 
+    array_module = da if coord.has_lazy_bounds() else np
     factor = iris.coords.AuxCoord(
-        np.diff(coord.bounds)[..., -1],
+        array_module.diff(coord.core_bounds())[..., -1],
         var_name=coord.var_name,
         long_name=coord.long_name,
         units=coord.units,

esmvalcore/preprocessor/_volume.py

@@ -128,7 +128,7 @@ def calculate_volume(cube: Cube) -> da.core.Array:
     z_dim = cube.coord_dims(depth)[0]
 
     # Calculate Z-direction thickness
-    thickness = depth.bounds[..., 1] - depth.bounds[..., 0]
+    thickness = depth.core_bounds()[..., 1] - depth.core_bounds()[..., 0]
 
     # Try to calculate grid cell area
     try:

tests/integration/cmor/_fixes/cmip5/test_hadgem2_es.py

@@ -1,10 +1,16 @@
 """Test HADGEM2-ES fixes."""
 import unittest
 
+import dask.array as da
+import iris.coords
+import iris.cube
+import numpy as np
+
 from esmvalcore.cmor._fixes.cmip5.hadgem2_es import O2, AllVars, Cl
 from esmvalcore.cmor._fixes.common import ClFixHybridHeightCoord
 from esmvalcore.cmor._fixes.fix import GenericFix
 from esmvalcore.cmor.fix import Fix
+from tests import assert_array_equal
 
 
 class TestAllVars(unittest.TestCase):
@@ -16,6 +22,30 @@ def test_get(self):
             Fix.get_fixes('CMIP5', 'HADGEM2-ES', 'Amon', 'tas'),
             [AllVars(None), GenericFix(None)])
 
+    @staticmethod
+    def test_clip_latitude():
+        cube = iris.cube.Cube(
+            da.arange(2, dtype=np.float32),
+            aux_coords_and_dims=[
+                (
+                    iris.coords.AuxCoord(
+                        da.asarray([90., 91.]),
+                        bounds=da.asarray([[89.5, 90.5], [90.5, 91.5]]),
+                        standard_name='latitude',
+                    ),
+                    0,
+                ),
+            ],
+        )
+        fix = AllVars(None)
+        cubes = fix.fix_metadata([cube])
+        assert len(cubes) == 1
+        coord = cubes[0].coord('latitude')
+        assert coord.has_lazy_points()
+        assert coord.has_lazy_bounds()
+        assert_array_equal(coord.points, np.array([90., 90]))
+        assert_array_equal(coord.bounds, np.array([[89.5, 90.], [90., 90.]]))
+
 
 class TestO2(unittest.TestCase):
     """Test o2 fixes."""