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."""