From 0884afeefb0ad100aa295721100681419f0e3f88 Mon Sep 17 00:00:00 2001 From: Manuel Schlund Date: Fri, 11 Oct 2024 10:44:31 +0200 Subject: [PATCH 1/3] Allow grouping cubes by data for fix_metadata --- esmvalcore/cmor/_fixes/fix.py | 10 +++++++++ esmvalcore/cmor/fix.py | 42 ++++++++++++++++++++++++++--------- 2 files changed, 42 insertions(+), 10 deletions(-) diff --git a/esmvalcore/cmor/_fixes/fix.py b/esmvalcore/cmor/_fixes/fix.py index 5aa41f6486..d93dc0de35 100644 --- a/esmvalcore/cmor/_fixes/fix.py +++ b/esmvalcore/cmor/_fixes/fix.py @@ -40,6 +40,16 @@ class Fix: """Base class for dataset fixes.""" + GROUP_CUBES_BY_DATE = False + """Flag for grouping cubes for fix_metadata. + + Fixes are applied to each group element individually. + + If ``False`` (default), group cubes by file. If ``True``, group cubes by + date. + + """ + def __init__( self, vardef: VariableInfo, diff --git a/esmvalcore/cmor/fix.py b/esmvalcore/cmor/fix.py index 2e3209897d..d509980fc9 100644 --- a/esmvalcore/cmor/fix.py +++ b/esmvalcore/cmor/fix.py @@ -10,15 +10,16 @@ import logging import warnings from collections import defaultdict -from collections.abc import Sequence +from collections.abc import Iterable, Sequence from pathlib import Path -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Any, Optional from iris.cube import Cube, CubeList from esmvalcore.cmor._fixes.fix import Fix from esmvalcore.cmor.check import CheckLevels, _get_cmor_checker from esmvalcore.exceptions import ESMValCoreDeprecationWarning +from esmvalcore.local import _get_start_end_date if TYPE_CHECKING: from ..config import Session @@ -102,6 +103,27 @@ def fix_file( return file +def _group_cubes(fixes: Iterable[Fix], cubes: CubeList) -> dict[Any, CubeList]: + """Group cubes for fix_metadata; each group is processed individually.""" + grouped_cubes: dict[Any, CubeList] = defaultdict(CubeList) + + # Group by date + if any(fix.GROUP_CUBES_BY_DATE for fix in fixes): + for cube in cubes: + if "source_file" in cube.attributes: + dates = _get_start_end_date(cube.attributes["source_file"]) + else: + dates = None + grouped_cubes[dates].append(cube) + + # Group by file name + else: + for cube in cubes: + grouped_cubes[cube.attributes.get("source_file", "")].append(cube) + + return grouped_cubes + + def fix_metadata( cubes: Sequence[Cube], short_name: str, @@ -192,14 +214,14 @@ def fix_metadata( ) fixed_cubes = CubeList() - # Group cubes by input file and apply all fixes to each group element - # (i.e., each file) individually - by_file = defaultdict(list) - for cube in cubes: - by_file[cube.attributes.get("source_file", "")].append(cube) - - for cube_list in by_file.values(): - cube_list = CubeList(cube_list) + # Group cubes and apply all fixes to each group element individually. There + # are two options for grouping: + # (1) By input file name (default). + # (2) By time range (can be enabled by setting the attribute + # GROUP_CUBES_BY_DATE=True for the fix class; see + # _fixes.native6.era5.Rsut for an example). + grouped_cubes = _group_cubes(fixes, cubes) + for cube_list in grouped_cubes.values(): for fix in fixes: cube_list = fix.fix_metadata(cube_list) From 168c9f2a107ad7786f07f06a49774e5ab9a31a16 Mon Sep 17 00:00:00 2001 From: Manuel Schlund Date: Fri, 11 Oct 2024 10:44:40 +0200 Subject: [PATCH 2/3] Added ERA5 fix for rsut --- esmvalcore/cmor/_fixes/native6/era5.py | 38 ++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/esmvalcore/cmor/_fixes/native6/era5.py b/esmvalcore/cmor/_fixes/native6/era5.py index aa4beae984..3d25adf191 100644 --- a/esmvalcore/cmor/_fixes/native6/era5.py +++ b/esmvalcore/cmor/_fixes/native6/era5.py @@ -304,6 +304,44 @@ def fix_metadata(self, cubes): return cubes +class Rsut(Fix): + """Fixes for rsut.""" + + # Enable grouping cubes by date for fix_metadata since multiple variables + # from multiple files are needed + GROUP_CUBES_BY_DATE = True + + def fix_metadata(self, cubes): + """Fix metadata. + + Derive rsut as + + rsut = rsdt - rsnt + + with + + rsut = TOA Outgoing Shortwave Radiation + rsdt = TOA Incoming Shortwave Radiation + rsnt = TOA Net Incoming Shortwave Radiation + + """ + rsdt_cube = cubes.extract_cube( + iris.NameConstraint(long_name="TOA incident solar radiation") + ) + rsnt_cube = cubes.extract_cube( + iris.NameConstraint( + long_name="Mean top net short-wave radiation flux" + ) + ) + rsdt_cube = Rsdt(None).fix_metadata([rsdt_cube])[0] + rsdt_cube.convert_units(self.vardef.units) + + rsdt_cube.data = rsdt_cube.core_data() - rsnt_cube.core_data() + rsdt_cube.attributes["positive"] = "up" + + return iris.cube.CubeList([rsdt_cube]) + + class Rss(Fix): """Fixes for Rss.""" From 73450fa33b902053418c1bff37953a9780289fa0 Mon Sep 17 00:00:00 2001 From: Manuel Schlund Date: Fri, 11 Oct 2024 11:41:21 +0200 Subject: [PATCH 3/3] Fix test --- tests/unit/cmor/test_fix.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unit/cmor/test_fix.py b/tests/unit/cmor/test_fix.py index 01038d2786..c51073a2fb 100644 --- a/tests/unit/cmor/test_fix.py +++ b/tests/unit/cmor/test_fix.py @@ -121,6 +121,7 @@ def setUp(self): self.fixed_cube = self._create_mock_cube() self.mock_fix = Mock() self.mock_fix.fix_metadata.return_value = [self.intermediate_cube] + self.mock_fix.GROUP_CUBES_BY_DATE = False self.checker = Mock() self.check_metadata = self.checker.return_value.check_metadata self.expected_get_fixes_call = {