Adding fix for ORAS5

esmvalcore/cmor/_fixes/oras5/_base_fixes.py

@@ -0,0 +1,131 @@
+"""Fix base classes for ORAS5 on-the-fly CMORizer."""
+
+import logging
+from pathlib import Path
+
+import iris
+import numpy as np
+import dask.array as da
+import xarray as xr
+from iris import Constraint
+from iris.experimental.ugrid import Connectivity, Mesh
+
+from ..icon.icon import IconFix
+
+logger = logging.getLogger(__name__)
+
+
+class Oras5Fix(IconFix):
+    """Base class for fixes."""
+
+    CACHE_DIR = Path.home() / '.esmvaltool' / 'cache'
+    CACHE_VALIDITY = 7 * 24 * 60 * 60  # [s]; = 1 week
+    TIMEOUT = 5 * 60  # [s]; = 5 min
+    GRID_FILE_ATTR = 'grid_file_uri'
+
+    def __init__(self, *args, **kwargs):
+        """Initialize fix."""
+        super().__init__(*args, **kwargs)
+        self._horizontal_grids = {}
+        self._meshes = {}
+
+
+    def _create_mesh(self, cube):
+        """Create mesh from horizontal grid file.
+
+        Note
+        ----
+        This functions creates a new :class:`iris.experimental.ugrid.Mesh` from
+        the ``clat`` (already present in the cube), ``clon`` (already present
+        in the cube), ``vertex_index``, ``vertex_of_cell``, ``vlat``, and
+        ``vlon`` variables of the horizontal grid file.
+
+        We do not use :func:`iris.experimental.ugrid.Mesh.from_coords` with the
+        existing latitude and longitude coordinates here because this would
+        produce lots of duplicated entries for the node coordinates. The reason
+        for this is that the node coordinates are constructed from the bounds;
+        since each node is contained 6 times in the bounds array (each node is
+        shared by 6 neighboring cells) the number of nodes is 6 times higher
+        with :func:`iris.experimental.ugrid.Mesh.from_coords` compared to using
+        the information already present in the horizontal grid file.
+
+        """
+
+        horizontal_grid = self.get_horizontal_grid(cube)
+        mesh = horizontal_grid.extract_cube(Constraint('cell_area'))
+        face_lon = mesh.coord('longitude').core_points().flatten()
+        face_lat = mesh.coord('latitude').core_points().flatten()
+
+        node_lon = mesh.coord('longitude').core_bounds().flatten()
+        node_lat = mesh.coord('latitude').core_bounds().flatten()
+
+        # Make the node locations a 2D array
+        nodes_flat = np.stack([node_lon, node_lat], axis=1)
+
+        # Find the unique nodes to be able to associate them with the faces
+        # Unfortunately, dask does not support the axis parameter...
+        nodes_unique, indices = np.unique(nodes_flat, return_inverse=True, 
+                                          axis=0)
+
+        node_lon = da.from_array(nodes_unique[:,0])
+        node_lat = da.from_array(nodes_unique[:,1])
+
+        n_faces = len(face_lat)
+        n_vertices = int(len(indices) / n_faces)
+
+        # Reshaping to N_faces x M_nodes array
+        indices = da.reshape(da.from_array(indices), (n_faces, n_vertices))
+
+        # Add the mask, which should not have a True entry for ORAS5
+        mask = da.full(da.shape(indices), False)
+
+        ### Define the connectivity
+        connectivity = Connectivity(
+                    indices=da.ma.masked_array(indices,mask=mask),
+                    cf_role='face_node_connectivity',
+                    start_index=0,
+                    location_axis=0,
+        )
+
+        face_lon = (face_lon + 360) % 360
+        node_lon = (node_lon + 360) % 360
+
+        # Put everything together to get a U-Grid style mesh
+        node_lat = iris.coords.AuxCoord(node_lat, standard_name='latitude',
+                                        var_name='lat', long_name='latitude',
+                                        units='degrees_north')
+        node_lon = iris.coords.AuxCoord(node_lon, standard_name='longitude', 
+                                        var_name='lon', long_name='longitude',
+                                        units='degrees_east')
+        face_lat = iris.coords.AuxCoord(face_lat,  standard_name='latitude',
+                                        var_name='lat', long_name='latitude',
+                                        units='degrees_north')
+        face_lon = iris.coords.AuxCoord(face_lon, standard_name='longitude', 
+                                        var_name='lon', long_name='longitude',
+                                        units='degrees_east')
+
+        mesh = Mesh(
+            topology_dimension=2,
+            node_coords_and_axes=[(node_lat, 'y'), (node_lon, 'x')],
+            connectivities=[connectivity],
+            face_coords_and_axes=[(face_lat, 'y'), (face_lon, 'x')],
+        )
+
+        return mesh
+
+    def _get_grid_from_facet(self):
+        """Get horizontal grid from user-defined facet `horizontal_grid`."""
+        grid_path = self._get_path_from_facet(
+            'horizontal_grid', 'Horizontal grid file'
+        )
+        grid_name = grid_path.name
+
+        # If already loaded, return the horizontal grid
+        if grid_name in self._horizontal_grids:
+            return self._horizontal_grids[grid_name]
+
+        # Load file
+        self._horizontal_grids[grid_name] = iris.load_raw(grid_path)
+        # self._horizontal_grids[grid_name] = xr.open_dataset(grid_path)     
+        logger.debug("Loaded ORAS5 grid file from %s", grid_path)
+        return self._horizontal_grids[grid_name]

esmvalcore/cmor/_fixes/oras5/oras5.py

@@ -0,0 +1,199 @@
+"""On-the-fly CMORizer for ORAS5."""
+
+import logging
+
+import iris
+import iris.util
+import numpy as np
+import dask.array as da
+from iris import Constraint
+from iris.coords import DimCoord
+from iris.cube import CubeList
+
+from ..shared import fix_ocean_depth_coord
+
+from ._base_fixes import Oras5Fix
+from ..icon.icon import AllVars as AllVars_ICON
+
+logger = logging.getLogger(__name__)
+
+
+class AllVars(Oras5Fix, AllVars_ICON):
+    """Fixes for all variables."""
+
+    def fix_metadata(self, cubes):
+
+        """Fix metadata."""
+        cubes = self.add_additional_cubes(cubes)
+        cube = self.get_cube(cubes)
+
+        # This is just a quick solution for other than horizontal coordinates,
+        # needs to be adapted to also deal with depth.
+        time = cube.coord('time')
+
+        # Adding the option to make the irregular (2d) grid unstructured (1d)
+        # to take advantage of UGRID
+        if self.extra_facets.get('make_unstructured', True):
+            # ORAS5 has 1 redundant row and 2 redundant columns that need to be
+            # removed.
+            data = cube.core_data()[...,:-1,1:-1].T.flatten()
+            data = da.reshape(data, (len(time.points), len(data)))
+            lat_points = cube.coord('latitude').core_points()
+            lat_points = lat_points[:-1,1:-1].T.flatten()
+            lon_points = cube.coord('longitude').core_points()
+            lon_points = lon_points[:-1,1:-1].T.flatten()
+
+            lat_coord = iris.coords.AuxCoord(lat_points, 
+                                             standard_name='latitude',
+                                             units=cube.coord('latitude').units)
+            lon_coord = iris.coords.AuxCoord(lon_points, 
+                                             standard_name='longitude',
+                                             units=cube.coord('longitude').units)
+
+            # See above concerning additional coordinates and dimensions
+            new_cube = iris.cube.Cube(data, dim_coords_and_dims=[(time,0)])
+            new_cube.add_aux_coord(lat_coord, 1)
+            new_cube.add_aux_coord(lon_coord, 1)
+
+            new_cube.long_name = cube.long_name
+            cube = new_cube
+
+        else:
+            # ORAS5 has 1 redundant row and 2 redundant columns that need to be
+            # removed.
+            cube = cube[...,:-1,1:-1]
+            lon_shape = cube.coord('longitude').points.shape
+            mesh = self.get_horizontal_grid(cube)
+            mesh = mesh.extract_cube(Constraint('cell_area'))
+            lon_bnds = mesh.coord('longitude').bounds
+            lat_bnds = mesh.coord('latitude').bounds
+            lon_bnds = np.reshape(lon_bnds, (lon_shape[0], lon_shape[1],
+                                             min(lon_bnds.shape)))
+            lat_bnds = np.reshape(lat_bnds, (lon_shape[0], lon_shape[1],
+                                             min(lat_bnds.shape)))
+            cube.coord('longitude').bounds = lon_bnds
+            cube.coord('latitude').bounds = lat_bnds
+
+        # Fix time
+        if self.vardef.has_coord_with_standard_name('time'):
+            cube = self._fix_time(cube, cubes)
+
+        if cube.coords(axis='Z'):
+            fix_ocean_depth_coord(cube)
+
+        # Fix latitude
+        if self.vardef.has_coord_with_standard_name('latitude'):
+            lat_idx = self._fix_lat(cube)
+        else:
+            lat_idx = None
+
+        # Fix longitude
+        if self.vardef.has_coord_with_standard_name('longitude'):
+            lon_idx = self._fix_lon(cube)
+        else:
+            lon_idx = None
+
+        # Fix unstructured mesh of unstructured grid if present
+        if self._is_unstructured_grid(lat_idx, lon_idx):
+            self._fix_mesh(cube, lat_idx)
+
+        # Fix metadata of variable
+        self.fix_var_metadata(cube)
+
+        return CubeList([cube])
+
+    def _add_coord_from_grid_file(self, cube, coord_name):
+        """Add coordinate from grid file to cube.
+
+        Note
+        ----
+        Assumes that the input cube has a single unnamed dimension, which will
+        be used as dimension for the new coordinate.
+
+        Parameters
+        ----------
+        cube: iris.cube.Cube
+            ICON data to which the coordinate from the grid file is added.
+        coord_name: str
+            Name of the coordinate to add from the grid file. Must be one of
+            ``'latitude'``, ``'longitude'``.
+
+        Raises
+        ------
+        ValueError
+            Invalid ``coord_name`` is given; input cube does not contain a
+            single unnamed dimension that can be used to add the new
+            coordinate.
+
+        """
+
+        # Use 'cell_area' as dummy cube to extract desired coordinates
+        # Note: it might be necessary to expand this when more coord_names are
+        # supported
+        horizontal_grid = self.get_horizontal_grid(cube)
+        if type(horizontal_grid) == iris.cube.CubeList:
+            grid_cube = horizontal_grid.extract_cube(
+                Constraint('cell_area'))
+            coord = grid_cube.coord(coord_name)
+        else:
+            if coord_name == 'longitude':
+                coord = iris.coords.AuxCoord(
+                                    points = (horizontal_grid.grid_center_lon
+                                              .values), 
+                                    bounds = (horizontal_grid.grid_corner_lon
+                                              .values),
+                                    standard_name = 'longitude',
+                                    units = 'degrees')
+            elif coord_name == 'latitude':
+                coord = iris.coords.AuxCoord(
+                                    points = (horizontal_grid.grid_center_lat
+                                              .values), 
+                                    bounds = (horizontal_grid.grid_corner_lat
+                                              .values),
+                                    standard_name = 'latitude',
+                                    units = 'degrees')
+
+        # Find index of mesh dimension (= single unnamed dimension)
+        n_unnamed_dimensions = cube.ndim - len(cube.dim_coords)
+        if n_unnamed_dimensions != 1:
+            raise ValueError(
+                f"Cannot determine coordinate dimension for coordinate "
+                f"'{coord_name}', cube does not contain a single unnamed "
+                f"dimension:\n{cube}")
+        coord_dims = ()
+        for idx in range(cube.ndim):
+            if not cube.coords(dimensions=idx, dim_coords=True):
+                coord_dims = (idx,)
+                break
+
+        # Adapt coordinate names so that the coordinate can be referenced with
+        # 'cube.coord(coord_name)'; the exact name will be set at a later stage
+        coord.standard_name = None
+        coord.long_name = coord_name
+        cube.add_aux_coord(coord, coord_dims)
+
+    def _fix_mesh(self, cube, mesh_idx):
+        """Fix mesh."""
+        # Remove any already-present dimensional coordinate describing the mesh
+        # dimension
+        if cube.coords(dimensions=mesh_idx, dim_coords=True):
+            cube.remove_coord(cube.coord(dimensions=mesh_idx, dim_coords=True))
+
+        # Add dimensional coordinate that describes the mesh dimension
+        index_coord = DimCoord(
+            np.arange(cube.shape[mesh_idx[0]]),
+            var_name='i',
+            long_name=('first spatial index for variables stored on an '
+                       'unstructured grid'),
+            units='1',
+        )
+        cube.add_dim_coord(index_coord, mesh_idx)
+
+        # If desired, get mesh and replace the original latitude and longitude
+        # coordinates with their new mesh versions
+        if self.extra_facets.get('ugrid', True):
+            mesh = self.get_mesh(cube)
+            cube.remove_coord('latitude')
+            cube.remove_coord('longitude')
+            for mesh_coord in mesh.to_MeshCoords('face'):
+                cube.add_aux_coord(mesh_coord, mesh_idx)

esmvalcore/config-developer.yml

@@ -194,3 +194,13 @@ CESM:
   output_file: '{project}_{dataset}_{case}_{gcomp}_{scomp}_{type}_{mip}_{short_name}'
   cmor_type: 'CMIP6'
   cmor_default_table_prefix: 'CMIP6_'
+
+ORAS5:
+  cmor_strict: false
+  input_dir:
+    default: '/'
+  input_file:
+    default: '*{raw_name}*{version}*.nc'
+  output_file: '{project}_{dataset}_{version}_{mip}_{short_name}'
+  cmor_type: 'CMIP6'
+  cmor_default_table_prefix: 'CMIP6_'