Repository: ESMValGroup/ESMValCore

Issue:
Introduce `mip` level fixes
As proposed by @tomaslovato in #1069, it would be nice to have fixes at `mip` level for each dataset. For example, this fix

```py
class Amon(Fix):
    """Fixes for all `Amon` variables."""
```

would be applied to all `Amon` variables of the given dataset. This would probably avoid lots of code duplication.

The implementation should be fairly easy as pointed out by @tomaslovato [here](https://github.com/ESMValGroup/ESMValCore/issues/1069#issuecomment-822432524).

I really like this idea! What say you @valeriupredoi @zklaus @jvegasbsc? We could discuss this in the general discussion of the workshop if necessary.


Discussion:
tomaslovato: @schlunma a couple of quick comments:
1) I think the test code for ocean variables Z axis 
https://github.com/ESMValGroup/ESMValCore/blob/b08d3d1b28c2bd9cb6801280305c26ff7877f16b/tests/integration/cmor/_fixes/test_shared.py#L508-L509
should better refer to the `lev` coordinate as in the following 
```python
 z_coord = iris.coords.DimCoord(0.0, var_name='lev', 
                                attributes={'positive': 'down'}) 
```

2) Did you tested already the new fix? As I remember that this comparison
https://github.com/ESMValGroup/ESMValCore/blob/65dc7d0f1449ed34eb20e7cab48dc4aec14ac761/esmvalcore/cmor/_fixes/cmip6/cesm2.py#L270
was not working until I changed to the `Units` instance of `cf_units` like in 
https://github.com/ESMValGroup/ESMValCore/blob/8b4a790557a3808711e7c25add7b6ec40150bfd5/esmvalcore/cmor/_fixes/cmip6/cesm2.py#L271

schlunma: 1. I did that on purpose, so that we really make sure that also the `var_name` and the `positive` attribute get fixed. As you can see further down the line, they do :smile:
2. Hmm...for me that works just fine. I tested it with real data and it looks good. Maybe it was a bug in an old version of `cf_units`?

```py
>>> import iris
>>> path = '/mnt/lustre02/work/bd0854/DATA/ESMValTool2/CMIP6_DKRZ/CMIP/NCAR/CESM2/historical/r1i1p1f1/Omon/thetao/gn/v20190308/thetao_Omon_CESM2_historical_r1i1p1f1_gn_185001-201412.nc'
>>> cube = iris.load_cube(path)
>>> lev_coord = cube.coord('lev')
>>> lev_coord.units == 'cm'
True
```

tomaslovato: @schlunma do you think that it might be ok to import also the ocean depth coordinate fixes for the other datasets I mentioned in #1069 ? If so I could also test the whole ce2coast recipe ...
schlunma: I think for this PR it's better to include only the actual support of `mip`-level fixes and one example of it for testing purposes. This way we keep the PR small and easier to review.

Once this is merged I'm happy to assist you with the remaining fixes. I noticed that there are some tests missing in the `ce2coast` branch. Could you maybe already include the to the `ce2coast` so we open the PR right after this is merged? If not, no problem, I can also add them later :+1: 
tomaslovato: @schlunma I see your point about the missing tests ..  so probably it would be better to do as you said by first closing this PR.
After this one has been merged, I'll open a new PR to collect the fixes for the other 3 datasets reported in #1069 and add the related tests.
At this point it will be probably better to solve also #1018 in a separated PR.

schlunma: Perfect, sounds great :+1: 
tomaslovato: @schlunma I was trying to run a recipe to process CESM2 dataset ([recipe_ocean_ce2coast_test.txt](https://github.com/ESMValGroup/ESMValCore/files/6425968/recipe_ocean_ce2coast_test.txt)) to get a 'working test' of the whole machinery, but I got swamped in the 'area statistics' computation with a message generated from here 
https://github.com/ESMValGroup/ESMValCore/blob/65dc7d0f1449ed34eb20e7cab48dc4aec14ac761/esmvalcore/preprocessor/_area.py#L211-L214
as the `area` cell_measure is not available from the cube core_data, although the fx_file is correctly read in the previous step.
 



jvegasbsc: Ping @sloosvel, as she was the one that implemented the handling of fx variables. I doubt it is due to the changes in this PR, though
tomaslovato: @schlunma @jvegasbsc I remove the obs data from the recipe and reduce only to one variable analysis ... 
Looking into the [main_log_debug.txt](https://github.com/ESMValGroup/ESMValCore/files/6426034/main_log_debug.txt) I realized the `cell_measure` attribute added from the fx_file survives until it reaches the `climate_statistics` preprocessor (around line 744 in the log).
`cell_measure` property is removed from the cube when the time statistics is computed with the `aggregate_by()` method in https://github.com/ESMValGroup/ESMValCore/blob/65dc7d0f1449ed34eb20e7cab48dc4aec14ac761/esmvalcore/preprocessor/_time.py#L561



when it got removed 

sloosvel: I checked the preprocessed files, and indeed after the `climate_statistics` step the cell measures are gone. I think it's a matter of modifying those functions so that they preserve the added variables.
schlunma: Thanks for looking into this @sloosvel. So this error is not caused by changes in this PR. Can you maybe try to run your recipe without using special preprocessor and diagnostic (by using `preprocessor: default` and `scripts: null`) and check if the depth coordinate of the dataset in the `preproc` is correct?
sloosvel: It does look fine to me!
schlunma: > It does look fine to me!

Awesome! Would you also have time to briefly check the code for the technical review @sloosvel? Thanks!!
schlunma: @tomaslovato I think your recipe should work if you remove the `custom_order: true` from the preprocessor 👍 
tomaslovato: @schlunma The implementation works fine for 3D data with simple recipe. LookinG forward to continue with remaining problematic datasets!

Patch:
diff --git a/esmvalcore/cmor/_fixes/cmip6/cesm2.py b/esmvalcore/cmor/_fixes/cmip6/cesm2.py
index 9b4bc499f..4ab1f0a4b 100644
--- a/esmvalcore/cmor/_fixes/cmip6/cesm2.py
+++ b/esmvalcore/cmor/_fixes/cmip6/cesm2.py
@@ -10,6 +10,7 @@ from ..shared import (
     add_scalar_height_coord,
     add_scalar_typeland_coord,
     add_scalar_typesea_coord,
+    fix_ocean_depth_coord,
 )
 from .gfdl_esm4 import Siconc as Addtypesi
 
@@ -243,3 +244,34 @@ class Tos(Fix):
 
 
 Siconc = Addtypesi
+
+
+class Omon(Fix):
+    """Fixes for ocean variables."""
+
+    def fix_metadata(self, cubes):
+        """Fix ocean depth coordinate.
+
+        Parameters
+        ----------
+        cubes: iris CubeList
+            List of cubes to fix
+
+        Returns
+        -------
+        iris.cube.CubeList
+
+        """
+        for cube in cubes:
+            if cube.coords(axis='Z'):
+                z_coord = cube.coord(axis='Z')
+
+                # Only points need to be fixed, not bounds
+                if z_coord.units == 'cm':
+                    z_coord.points = z_coord.core_points() / 100.0
+                    z_coord.units = 'm'
+
+                # Fix depth metadata
+                if z_coord.standard_name is None:
+                    fix_ocean_depth_coord(cube)
+        return cubes
diff --git a/esmvalcore/cmor/_fixes/fix.py b/esmvalcore/cmor/_fixes/fix.py
index 12aee7b9a..ba4175813 100644
--- a/esmvalcore/cmor/_fixes/fix.py
+++ b/esmvalcore/cmor/_fixes/fix.py
@@ -162,7 +162,7 @@ class Fix:
 
             classes = inspect.getmembers(fixes_module, inspect.isclass)
             classes = dict((name.lower(), value) for name, value in classes)
-            for fix_name in (short_name, 'allvars'):
+            for fix_name in (short_name, mip.lower(), 'allvars'):
                 try:
                     fixes.append(classes[fix_name](vardef))
                 except KeyError:
diff --git a/esmvalcore/cmor/_fixes/shared.py b/esmvalcore/cmor/_fixes/shared.py
index 62b9c76c8..de03d7885 100644
--- a/esmvalcore/cmor/_fixes/shared.py
+++ b/esmvalcore/cmor/_fixes/shared.py
@@ -517,3 +517,20 @@ def round_coordinates(cubes, decimals=5, coord_names=None):
                 coord.bounds = da.round(da.asarray(coord.core_bounds()),
                                         decimals)
     return cubes
+
+
+def fix_ocean_depth_coord(cube):
+    """Fix attributes of ocean vertical level coordinate.
+
+    Parameters
+    ----------
+    cube : iris.cube.Cube
+        Input cube.
+
+    """
+    depth_coord = cube.coord(axis='Z')
+    depth_coord.standard_name = 'depth'
+    depth_coord.var_name = 'lev'
+    depth_coord.units = 'm'
+    depth_coord.long_name = 'ocean depth coordinate'
+    depth_coord.attributes = {'positive': 'down'}


Test patch:
diff --git a/tests/integration/cmor/_fixes/cmip6/test_cesm2.py b/tests/integration/cmor/_fixes/cmip6/test_cesm2.py
index fec9c8c05..dcddc43f3 100644
--- a/tests/integration/cmor/_fixes/cmip6/test_cesm2.py
+++ b/tests/integration/cmor/_fixes/cmip6/test_cesm2.py
@@ -8,7 +8,7 @@ import numpy as np
 import pytest
 from cf_units import Unit
 
-from esmvalcore.cmor._fixes.cmip6.cesm2 import Cl, Cli, Clw, Tas, Tos
+from esmvalcore.cmor._fixes.cmip6.cesm2 import Cl, Cli, Clw, Omon, Tas, Tos
 from esmvalcore.cmor.fix import Fix
 from esmvalcore.cmor.table import get_var_info
 
@@ -225,6 +225,34 @@ def tos_cubes():
     return iris.cube.CubeList([tos_cube])
 
 
+@pytest.fixture
+def thetao_cubes():
+    """Cubes to test fixes for ``thetao``."""
+    time_coord = iris.coords.DimCoord(
+        [0.0004, 1.09776], var_name='time', standard_name='time',
+        units='days since 1850-01-01 00:00:00')
+    lat_coord = iris.coords.DimCoord(
+        [0.0, 1.0], var_name='lat', standard_name='latitude', units='degrees')
+    lon_coord = iris.coords.DimCoord(
+        [0.0, 1.0], var_name='lon', standard_name='longitude', units='degrees')
+    lev_coord = iris.coords.DimCoord(
+        [500.0, 1000.0], bounds=[[2.5, 7.5], [7.5, 12.5]],
+        var_name='lev', standard_name=None, units='cm',
+        attributes={'positive': 'up'})
+    coord_specs = [
+        (time_coord, 0),
+        (lev_coord, 1),
+        (lat_coord, 2),
+        (lon_coord, 3),
+    ]
+    thetao_cube = iris.cube.Cube(
+        np.ones((2, 2, 2, 2)),
+        var_name='thetao',
+        dim_coords_and_dims=coord_specs,
+    )
+    return iris.cube.CubeList([thetao_cube])
+
+
 def test_get_tas_fix():
     """Test getting of fix."""
     fix = Fix.get_fixes('CMIP6', 'CESM2', 'Amon', 'tas')
@@ -233,8 +261,14 @@ def test_get_tas_fix():
 
 def test_get_tos_fix():
     """Test getting of fix."""
-    fix = Fix.get_fixes('CMIP6', 'CESM2', 'Amon', 'tos')
-    assert fix == [Tos(None)]
+    fix = Fix.get_fixes('CMIP6', 'CESM2', 'Omon', 'tos')
+    assert fix == [Tos(None), Omon(None)]
+
+
+def test_get_thetao_fix():
+    """Test getting of fix."""
+    fix = Fix.get_fixes('CMIP6', 'CESM2', 'Omon', 'thetao')
+    assert fix == [Omon(None)]
 
 
 def test_tas_fix_metadata(tas_cubes):
@@ -271,3 +305,25 @@ def test_tos_fix_metadata(tos_cubes):
     assert out_cubes is tos_cubes
     for cube in out_cubes:
         np.testing.assert_equal(cube.coord("time").points, [0., 1.1])
+
+
+def test_thetao_fix_metadata(thetao_cubes):
+    """Test ``fix_metadata`` for ``thetao``."""
+    vardef = get_var_info('CMIP6', 'Omon', 'thetao')
+    fix = Omon(vardef)
+    out_cubes = fix.fix_metadata(thetao_cubes)
+    assert out_cubes is thetao_cubes
+    assert len(out_cubes) == 1
+    out_cube = out_cubes[0]
+
+    # Check metadata of depth coordinate
+    depth_coord = out_cube.coord('depth')
+    assert depth_coord.standard_name == 'depth'
+    assert depth_coord.var_name == 'lev'
+    assert depth_coord.long_name == 'ocean depth coordinate'
+    assert depth_coord.units == 'm'
+    assert depth_coord.attributes == {'positive': 'down'}
+
+    # Check values of depth coordinate
+    np.testing.assert_allclose(depth_coord.points, [5.0, 10.0])
+    np.testing.assert_allclose(depth_coord.bounds, [[2.5, 7.5], [7.5, 12.5]])
diff --git a/tests/integration/cmor/_fixes/test_fix.py b/tests/integration/cmor/_fixes/test_fix.py
index 019991150..5ffc4ca86 100644
--- a/tests/integration/cmor/_fixes/test_fix.py
+++ b/tests/integration/cmor/_fixes/test_fix.py
@@ -1,3 +1,5 @@
+"""Integration tests for fixes."""
+
 import os
 import shutil
 import tempfile
@@ -6,35 +8,35 @@ import unittest
 import pytest
 from iris.cube import Cube
 
+from esmvalcore.cmor._fixes.cmip5.bnu_esm import Ch4
+from esmvalcore.cmor._fixes.cmip5.canesm2 import FgCo2
+from esmvalcore.cmor._fixes.cmip5.cesm1_bgc import Gpp
+from esmvalcore.cmor._fixes.cmip6.cesm2 import Omon, Tos
 from esmvalcore.cmor.fix import Fix
 
 
 class TestFix(unittest.TestCase):
     def setUp(self):
-        """Set up temp folder"""
+        """Set up temp folder."""
         self.temp_folder = tempfile.mkdtemp()
 
     def tearDown(self):
-        """Remove temp folder"""
+        """Remove temp folder."""
         shutil.rmtree(self.temp_folder)
 
     def test_get_fix(self):
-        from esmvalcore.cmor._fixes.cmip5.canesm2 import FgCo2
         self.assertListEqual(
             Fix.get_fixes('CMIP5', 'CanESM2', 'Amon', 'fgco2'), [FgCo2(None)])
 
     def test_get_fix_case_insensitive(self):
-        from esmvalcore.cmor._fixes.cmip5.canesm2 import FgCo2
         self.assertListEqual(
             Fix.get_fixes('CMIP5', 'CanESM2', 'Amon', 'fgCo2'), [FgCo2(None)])
 
     def test_get_fixes_with_replace(self):
-        from esmvalcore.cmor._fixes.cmip5.bnu_esm import Ch4
         self.assertListEqual(Fix.get_fixes('CMIP5', 'BNU-ESM', 'Amon', 'ch4'),
                              [Ch4(None)])
 
     def test_get_fixes_with_generic(self):
-        from esmvalcore.cmor._fixes.cmip5.cesm1_bgc import Gpp
         self.assertListEqual(
             Fix.get_fixes('CMIP5', 'CESM1-BGC', 'Amon', 'gpp'), [Gpp(None)])
 
@@ -50,6 +52,19 @@ class TestFix(unittest.TestCase):
         self.assertListEqual(
             Fix.get_fixes('CMIP5', 'BNU-ESM', 'Amon', 'BAD_VAR'), [])
 
+    def test_get_fix_only_mip(self):
+        self.assertListEqual(
+            Fix.get_fixes('CMIP6', 'CESM2', 'Omon', 'thetao'), [Omon(None)])
+
+    def test_get_fix_only_mip_case_insensitive(self):
+        self.assertListEqual(
+            Fix.get_fixes('CMIP6', 'CESM2', 'omOn', 'thetao'), [Omon(None)])
+
+    def test_get_fix_mip_and_var(self):
+        self.assertListEqual(
+            Fix.get_fixes('CMIP6', 'CESM2', 'Omon', 'tos'),
+            [Tos(None), Omon(None)])
+
     def test_fix_metadata(self):
         cube = Cube([0])
         reference = Cube([0])
diff --git a/tests/integration/cmor/_fixes/test_shared.py b/tests/integration/cmor/_fixes/test_shared.py
index 687bfd4f1..b185c7c92 100644
--- a/tests/integration/cmor/_fixes/test_shared.py
+++ b/tests/integration/cmor/_fixes/test_shared.py
@@ -15,6 +15,7 @@ from esmvalcore.cmor._fixes.shared import (
     add_sigma_factory,
     cube_to_aux_coord,
     fix_bounds,
+    fix_ocean_depth_coord,
     get_altitude_to_pressure_func,
     get_bounds_cube,
     get_pressure_to_altitude_func,
@@ -500,3 +501,18 @@ def test_round_coordinates_single_coord():
     assert cubes[0].coord('longitude') is out[0].coord('longitude')
     np.testing.assert_allclose(out[0].coord('latitude').points, [10])
     np.testing.assert_allclose(out[0].coord('latitude').bounds, [[9, 11]])
+
+
+def test_fix_ocean_depth_coord():
+    """Test `fix_ocean_depth_coord`."""
+    z_coord = iris.coords.DimCoord(0.0, var_name='alt',
+                                   attributes={'positive': 'up'})
+    cube = iris.cube.Cube([0.0], var_name='x',
+                          dim_coords_and_dims=[(z_coord, 0)])
+    fix_ocean_depth_coord(cube)
+    depth_coord = cube.coord('depth')
+    assert depth_coord.standard_name == 'depth'
+    assert depth_coord.var_name == 'lev'
+    assert depth_coord.units == 'm'
+    assert depth_coord.long_name == 'ocean depth coordinate'
+    assert depth_coord.attributes == {'positive': 'down'}