Skip to content
Draft
Show file tree
Hide file tree
Changes from 6 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions esmvalcore/cmor/_fixes/fix.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
38 changes: 38 additions & 0 deletions esmvalcore/cmor/_fixes/native6/era5.py
Original file line number Diff line number Diff line change
Expand Up @@ -408,6 +408,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."""

Expand Down
42 changes: 32 additions & 10 deletions esmvalcore/cmor/fix.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,13 +9,14 @@

import logging
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.local import _get_start_end_date

if TYPE_CHECKING:
from ..config import Session
Expand Down Expand Up @@ -99,6 +100,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,
Expand Down Expand Up @@ -163,14 +185,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)
Comment on lines 216 to 228

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It may be nicer to define the grouping operation on the Fix object, so this code would look like:

    fixed_cubes = CubeList(cubes)
    for fix in fixes:
        fixed_cubes = CubeList(
            cube
            for group in fix.group_input_for_fix_metadata(fixed_cubes)
            for cube in fix.fix_metadata(group)
        )


Expand Down
1 change: 1 addition & 0 deletions tests/unit/cmor/test_fix.py
Original file line number Diff line number Diff line change
Expand Up @@ -119,6 +119,7 @@ def setUp(self):
self.cube = self._create_mock_cube()
self.fixed_cube = self._create_mock_cube()
self.mock_fix = Mock()
self.mock_fix.GROUP_CUBES_BY_DATE = False
self.mock_fix.fix_metadata.return_value = [self.fixed_cube]
self.expected_get_fixes_call = {
"project": "project",
Expand Down