From 3f610e714a8ee62ff9d7cba3b6668263c5b19aa3 Mon Sep 17 00:00:00 2001 From: Bouwe Andela Date: Thu, 15 Feb 2024 19:16:17 +0100 Subject: [PATCH] More lazy fixes and preprocessing functions (#2325) Co-authored-by: Manuel Schlund <32543114+schlunma@users.noreply.github.com> --- esmvalcore/cmor/_fixes/cesm/cesm2.py | 4 +-- esmvalcore/cmor/_fixes/cmip5/hadgem2_cc.py | 29 ++--------------- esmvalcore/cmor/_fixes/cmip5/hadgem2_es.py | 4 +-- esmvalcore/cmor/_fixes/cmip5/miroc_esm.py | 2 +- esmvalcore/cmor/_fixes/cmip6/access_esm1_5.py | 4 +-- esmvalcore/cmor/_fixes/cmip6/cesm2.py | 16 ++++------ esmvalcore/cmor/_fixes/cmip6/ec_earth3.py | 4 +-- esmvalcore/cmor/_fixes/cmip6/ec_earth3_veg.py | 4 +-- esmvalcore/cmor/_fixes/cmip6/giss_e2_1_g.py | 4 +-- esmvalcore/cmor/_fixes/cmip6/miroc6.py | 22 +++++-------- esmvalcore/cmor/_fixes/common.py | 2 +- esmvalcore/cmor/_fixes/cordex/cordex_fixes.py | 2 +- esmvalcore/cmor/_fixes/emac/emac.py | 2 +- esmvalcore/cmor/_fixes/icon/_base_fixes.py | 6 ++-- esmvalcore/cmor/_fixes/native6/era5.py | 4 +-- esmvalcore/cmor/_fixes/native6/mswep.py | 2 +- esmvalcore/cmor/_fixes/shared.py | 9 +++--- esmvalcore/preprocessor/_derive/sm.py | 2 +- esmvalcore/preprocessor/_units.py | 15 ++++++--- esmvalcore/preprocessor/_volume.py | 2 +- .../cmor/_fixes/cmip5/test_hadgem2_es.py | 30 +++++++++++++++++ tests/unit/preprocessor/_derive/test_sm.py | 32 +++++++++++++++++++ 22 files changed, 116 insertions(+), 85 deletions(-) create mode 100644 tests/unit/preprocessor/_derive/test_sm.py diff --git a/esmvalcore/cmor/_fixes/cesm/cesm2.py b/esmvalcore/cmor/_fixes/cesm/cesm2.py index a76e5a1a48..5b92121555 100644 --- a/esmvalcore/cmor/_fixes/cesm/cesm2.py +++ b/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) diff --git a/esmvalcore/cmor/_fixes/cmip5/hadgem2_cc.py b/esmvalcore/cmor/_fixes/cmip5/hadgem2_cc.py index 33e4347a65..3531f78f06 100644 --- a/esmvalcore/cmor/_fixes/cmip5/hadgem2_cc.py +++ b/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): diff --git a/esmvalcore/cmor/_fixes/cmip5/hadgem2_es.py b/esmvalcore/cmor/_fixes/cmip5/hadgem2_es.py index b2cf33c8d2..f7360dae2b 100644 --- a/esmvalcore/cmor/_fixes/cmip5/hadgem2_es.py +++ b/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 diff --git a/esmvalcore/cmor/_fixes/cmip5/miroc_esm.py b/esmvalcore/cmor/_fixes/cmip5/miroc_esm.py index 0721e62409..fc94bc45cc 100644 --- a/esmvalcore/cmor/_fixes/cmip5/miroc_esm.py +++ b/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 diff --git a/esmvalcore/cmor/_fixes/cmip6/access_esm1_5.py b/esmvalcore/cmor/_fixes/cmip6/access_esm1_5.py index da884b5e3a..7b0497e493 100644 --- a/esmvalcore/cmor/_fixes/cmip6/access_esm1_5.py +++ b/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 diff --git a/esmvalcore/cmor/_fixes/cmip6/cesm2.py b/esmvalcore/cmor/_fixes/cmip6/cesm2.py index 134f68b29c..6ded187fbb 100644 --- a/esmvalcore/cmor/_fixes/cmip6/cesm2.py +++ b/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 diff --git a/esmvalcore/cmor/_fixes/cmip6/ec_earth3.py b/esmvalcore/cmor/_fixes/cmip6/ec_earth3.py index 27ce37961d..49bb6a66a2 100644 --- a/esmvalcore/cmor/_fixes/cmip6/ec_earth3.py +++ b/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 diff --git a/esmvalcore/cmor/_fixes/cmip6/ec_earth3_veg.py b/esmvalcore/cmor/_fixes/cmip6/ec_earth3_veg.py index edcf5624d2..73bf689725 100644 --- a/esmvalcore/cmor/_fixes/cmip6/ec_earth3_veg.py +++ b/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 diff --git a/esmvalcore/cmor/_fixes/cmip6/giss_e2_1_g.py b/esmvalcore/cmor/_fixes/cmip6/giss_e2_1_g.py index 5d7033fbdf..f4fb27d35c 100644 --- a/esmvalcore/cmor/_fixes/cmip6/giss_e2_1_g.py +++ b/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 diff --git a/esmvalcore/cmor/_fixes/cmip6/miroc6.py b/esmvalcore/cmor/_fixes/cmip6/miroc6.py index b440d81c64..cf1d40ca86 100644 --- a/esmvalcore/cmor/_fixes/cmip6/miroc6.py +++ b/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 diff --git a/esmvalcore/cmor/_fixes/common.py b/esmvalcore/cmor/_fixes/common.py index 6557e65853..d68bbac538 100644 --- a/esmvalcore/cmor/_fixes/common.py +++ b/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') diff --git a/esmvalcore/cmor/_fixes/cordex/cordex_fixes.py b/esmvalcore/cmor/_fixes/cordex/cordex_fixes.py index 9ccbfd4ed3..3a7a26d1e9 100644 --- a/esmvalcore/cmor/_fixes/cordex/cordex_fixes.py +++ b/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 diff --git a/esmvalcore/cmor/_fixes/emac/emac.py b/esmvalcore/cmor/_fixes/emac/emac.py index 27a7d3c74d..38ae862916 100644 --- a/esmvalcore/cmor/_fixes/emac/emac.py +++ b/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') diff --git a/esmvalcore/cmor/_fixes/icon/_base_fixes.py b/esmvalcore/cmor/_fixes/icon/_base_fixes.py index 4fc67f490e..858cb41124 100644 --- a/esmvalcore/cmor/_fixes/icon/_base_fixes.py +++ b/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 diff --git a/esmvalcore/cmor/_fixes/native6/era5.py b/esmvalcore/cmor/_fixes/native6/era5.py index d76fd21701..6c67494aaa 100644 --- a/esmvalcore/cmor/_fixes/native6/era5.py +++ b/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() diff --git a/esmvalcore/cmor/_fixes/native6/mswep.py b/esmvalcore/cmor/_fixes/native6/mswep.py index 9441b3a12d..7ed50fcdcb 100644 --- a/esmvalcore/cmor/_fixes/native6/mswep.py +++ b/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): diff --git a/esmvalcore/cmor/_fixes/shared.py b/esmvalcore/cmor/_fixes/shared.py index e9ae61499f..b50deb422b 100644 --- a/esmvalcore/cmor/_fixes/shared.py +++ b/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 diff --git a/esmvalcore/preprocessor/_derive/sm.py b/esmvalcore/preprocessor/_derive/sm.py index 2e22df7c7c..e2da2857a6 100644 --- a/esmvalcore/preprocessor/_derive/sm.py +++ b/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 diff --git a/esmvalcore/preprocessor/_units.py b/esmvalcore/preprocessor/_units.py index 069cef1a6d..8c96f78ae1 100644 --- a/esmvalcore/preprocessor/_units.py +++ b/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,7 +123,10 @@ 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 @@ -128,10 +134,10 @@ def accumulate_coordinate(cube, coordinate): 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, diff --git a/esmvalcore/preprocessor/_volume.py b/esmvalcore/preprocessor/_volume.py index 8eb03e11d2..5f64609193 100644 --- a/esmvalcore/preprocessor/_volume.py +++ b/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: diff --git a/tests/integration/cmor/_fixes/cmip5/test_hadgem2_es.py b/tests/integration/cmor/_fixes/cmip5/test_hadgem2_es.py index c0ab2bb8fd..17d928e574 100644 --- a/tests/integration/cmor/_fixes/cmip5/test_hadgem2_es.py +++ b/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.""" diff --git a/tests/unit/preprocessor/_derive/test_sm.py b/tests/unit/preprocessor/_derive/test_sm.py new file mode 100644 index 0000000000..7417b8f190 --- /dev/null +++ b/tests/unit/preprocessor/_derive/test_sm.py @@ -0,0 +1,32 @@ +import dask.array as da +import iris.coords +import iris.cube +import numpy as np + +from esmvalcore.preprocessor._derive.sm import DerivedVariable +from tests import assert_array_equal + + +def test_sm(): + + points = da.arange(0, 4, 2).astype(np.float32) + bounds = da.asarray([[-1., 1.], [1., 3]]) + + depth = iris.coords.AuxCoord( + points, + bounds=bounds, + standard_name='depth', + ) + cube = iris.cube.Cube( + da.asarray([0, 998.2]), + var_name='mrsos', + aux_coords_and_dims=[ + (depth, 0), + ], + ) + + result = DerivedVariable.calculate(iris.cube.CubeList([cube])) + assert result.has_lazy_data() + assert result.coord('depth').has_lazy_points() + assert result.coord('depth').has_lazy_bounds() + assert_array_equal(result.data, np.array([0, 0.5]))