Support NetCDF zlib compression serialization of UGRID meshes (SciTools#6728)

* support mesh saving with compression

* add test coverage

* add whatsnew entry

* add review comments

* review actions

* review action

* backout note

Author: bjlittle

docs/src/whatsnew/latest.rst

@@ -34,6 +34,12 @@ This document explains the changes made to Iris for this release
    horizontal grid.
    (:issue:`5770`, :pull:`6581`)
 
+#. `@bjlittle`_ extended ``zlib`` compression of :class:`~iris.cube.Cube` data
+   payload when saving to NetCDF to also include any attached `CF-UGRID`_
+   :class:`~iris.mesh.components.MeshXY`. Additionally,
+   :func:`~iris.fileformats.netcdf.saver.save_mesh` also supports ``zlib``
+   compression. (:issue:`6565`, :pull:`6728`)
+
 
 🐛 Bugs Fixed
 =============

lib/iris/fileformats/netcdf/saver.py

@@ -587,11 +587,6 @@ def write(
         all_dimensions = mesh_dimensions + nonmesh_dimensions
         self._create_cf_dimensions(cube, all_dimensions, unlimited_dimensions)
 
-        # Create the mesh components, if there is a mesh.
-        # We do this before creating the data-var, so that mesh vars precede
-        # data-vars in the file.
-        cf_mesh_name = self._add_mesh(cube)
-
         # Group the generic compression keyword arguments together for
         # convenience, as they will be applied to other cube metadata
         # as well as the cube data payload.
@@ -602,6 +597,11 @@ def write(
             "zlib": zlib,
         }
 
+        # Create the mesh components, if there is a mesh.
+        # We do this before creating the data-var, so that mesh vars precede
+        # data-vars in the file.
+        cf_mesh_name = self._add_mesh(cube, compression_kwargs=compression_kwargs)
+
         # Create the associated cube CF-netCDF data variable.
         cf_var_cube = self._create_cf_data_variable(
             cube,
@@ -772,7 +772,7 @@ def _create_cf_dimensions(self, cube, dimension_names, unlimited_dimensions=None
                     size = self._existing_dim[dim_name]
                 self._dataset.createDimension(dim_name, size)
 
-    def _add_mesh(self, cube_or_mesh):
+    def _add_mesh(self, cube_or_mesh, /, *, compression_kwargs=None):
         """Add the cube's mesh, and all related variables to the dataset.
 
         Add the cube's mesh, and all related variables to the dataset.
@@ -787,6 +787,8 @@ def _add_mesh(self, cube_or_mesh):
         ----------
         cube_or_mesh : :class:`iris.cube.Cube` or :class:`iris.mesh.MeshXY`
             The Cube or Mesh being saved to the netCDF file.
+        compression_kwargs : dict, optional
+            NetCDF data compression keyword arguments.
 
         Returns
         -------
@@ -801,10 +803,8 @@ def _add_mesh(self, cube_or_mesh):
         from iris.cube import Cube
 
         if isinstance(cube_or_mesh, Cube):
-            cube = cube_or_mesh
-            mesh = cube.mesh
+            mesh = cube_or_mesh.mesh
         else:
-            cube = None  # The underlying routines must support this !
             mesh = cube_or_mesh
 
         if mesh:
@@ -830,10 +830,11 @@ def _add_mesh(self, cube_or_mesh):
                             if coord is None:
                                 continue  # an awkward thing that mesh.coords does
                             coord_name = self._create_generic_cf_array_var(
-                                cube_or_mesh,
+                                mesh,
                                 [],
                                 coord,
                                 element_dims=(mesh_dims[location],),
+                                compression_kwargs=compression_kwargs,
                             )
                             # Only created once per file, but need to fetch the
                             #  name later in _add_inner_related_vars().
@@ -878,11 +879,12 @@ def _add_mesh(self, cube_or_mesh):
                     else:
                         fill_value = None
                     cf_conn_name = self._create_generic_cf_array_var(
-                        cube_or_mesh,
+                        mesh,
                         [],
                         conn,
                         element_dims=conn_dims,
                         fill_value=fill_value,
+                        compression_kwargs=compression_kwargs,
                     )
                     # Add essential attributes to the Connectivity variable.
                     cf_conn_var = self._dataset.variables[cf_conn_name]
@@ -1757,7 +1759,7 @@ def _create_generic_cf_array_var(
             An Iris :class:`iris.coords._DimensionalMetadata`, belonging to the
             cube.  Provides data, units and standard/long/var names.
             Not used if 'element_dims' is not None.
-        element_dims : list of str, optionsl
+        element_dims : list of str, optional
             If set, contains the variable dimension (names),
             otherwise these are taken from `element.cube_dims[cube]`.
             For Mesh components (element coordinates and connectivities), this
@@ -1793,7 +1795,7 @@ def _create_generic_cf_array_var(
         while cf_name in self._dataset.variables:
             cf_name = self._increment_name(cf_name)
 
-        if element_dims is None:
+        if cube and element_dims is None:
             # Get the list of file-dimensions (names), to create the variable.
             element_dims = [
                 cube_dim_names[dim] for dim in element.cube_dims(cube)
@@ -1804,7 +1806,8 @@ def _create_generic_cf_array_var(
         # (e.g. =points if a coord, =data if an ancillary, etc)
         data = element._core_values()
 
-        if cube is None or cube.shape != data.shape:
+        # This compression contract is *not* applicable to a mesh.
+        if cube and cube.shape != data.shape:
             compression_kwargs = {}
 
         if np.issubdtype(data.dtype, np.str_):
@@ -2993,7 +2996,17 @@ def is_valid_packspec(p):
     return result
 
 
-def save_mesh(mesh, filename, netcdf_format="NETCDF4"):
+def save_mesh(
+    mesh,
+    filename,
+    /,
+    *,
+    complevel=4,
+    fletcher32=False,
+    netcdf_format="NETCDF4",
+    shuffle=True,
+    zlib=False,
+):
     """Save mesh(es) to a netCDF file.
 
     Parameters
@@ -3002,16 +3015,37 @@ def save_mesh(mesh, filename, netcdf_format="NETCDF4"):
         Mesh(es) to save.
     filename : str
         Name of the netCDF file to create.
+    complevel : int, default=4
+        An integer between 1 and 9 describing the level of compression
+        desired. Ignored if ``zlib=False``.
+    fletcher32 : bool, default=False
+        If ``True``, the Fletcher32 HDF5 checksum algorithm is activated to
+        detect errors.
     netcdf_format : str, default="NETCDF4"
-        Underlying netCDF file format, one of 'NETCDF4', 'NETCDF4_CLASSIC',
-        'NETCDF3_CLASSIC' or 'NETCDF3_64BIT'. Default is 'NETCDF4' format.
+        Underlying netCDF file format, one of ``NETCDF4``, ``NETCDF4_CLASSIC``,
+        ``NETCDF3_CLASSIC`` or ``NETCDF3_64BIT``. Default is ``NETCDF4`` format.
+    shuffle : bool, default=True
+        If ``True``, the HDF5 shuffle filter will be applied before
+        compressing the data. This significantly improves compression.
+        Ignored if ``zlib=False``.
+    zlib : bool, default=False
+        If ``True``, the data will be compressed in the netCDF file using
+        gzip compression.
 
     """
     if isinstance(mesh, typing.Iterable):
         meshes = mesh
     else:
         meshes = [mesh]
 
+    # Group the generic compression keyword arguments together.
+    compression_kwargs = {
+        "complevel": complevel,
+        "fletcher32": fletcher32,
+        "shuffle": shuffle,
+        "zlib": zlib,
+    }
+
     # Initialise Manager for saving
     with Saver(filename, netcdf_format) as sman:
         # Iterate through the list.
@@ -3023,7 +3057,7 @@ def save_mesh(mesh, filename, netcdf_format="NETCDF4"):
             sman._create_cf_dimensions(cube=None, dimension_names=mesh_dimensions)
 
             # Create the mesh components.
-            sman._add_mesh(mesh)
+            sman._add_mesh(mesh, compression_kwargs=compression_kwargs)
 
         # Add a conventions attribute.
         # TODO: add 'UGRID' to conventions, when this is agreed with CF ?

lib/iris/tests/unit/fileformats/netcdf/saver/test_Saver__ugrid.py

@@ -329,6 +329,35 @@ def vars_meshdim(vars, location, mesh_name=None):
     return single_location_dim
 
 
+def filter_compression_calls(patch, compression_kwargs, mismatch=False):
+    """Pre-process the mock patch calls for compression kwargs.
+
+    Parameters
+    ----------
+    patch : mock
+        The mock patch instance to be inspected.
+    compression_kwargs : dict
+        The compression kwargs to match as parameters in calls.
+    mismatch : bool, default=False
+        Whether to check for calls that do not match the compression kwargs.
+
+    Returns
+    -------
+    set
+        The netCDF variables serialized with compression kwargs (or not).
+
+    """
+    result = set()
+    for call in patch.call_args_list:
+        kwargs = call.kwargs
+        if all(kwargs.get(k) == v for k, v in compression_kwargs.items()):
+            if not mismatch:
+                result.add(call.args[0])
+        elif mismatch:
+            result.add(call.args[0])
+    return result
+
+
 class TestSaveUgrid__cube(tests.IrisTest):
     """Test for saving cubes which have meshes."""
 
@@ -340,7 +369,7 @@ def setUpClass(cls):
     def tearDownClass(cls):
         shutil.rmtree(cls.temp_dir)
 
-    def check_save_cubes(self, cube_or_cubes):
+    def check_save_cubes(self, cube_or_cubes, compression_kwargs=None):
         """Write cubes to a new file in the common temporary directory.
 
         Use a name unique to this testcase, to avoid any clashes.
@@ -352,11 +381,52 @@ def check_save_cubes(self, cube_or_cubes):
         # in the common temporary directory.
         tempfile_path = self.temp_dir / Path(tempfile_path).name
 
+        if compression_kwargs is None:
+            compression_kwargs = {}
+
         # Save data to the file.
-        save(cube_or_cubes, tempfile_path)
+        save(cube_or_cubes, tempfile_path, **compression_kwargs)
 
         return tempfile_path
 
+    def test_compression(self):
+        """Test NetCDF serialization of a cube with attached mesh using compression.
+
+        NetCDF data compression keyword arguments include "complevel",
+        "fletcher32", "shuffle" and "zlib". Note that "complevel" and "shuffle"
+        are only applicable when "zlib=True".
+
+        """
+        # Note that the patch location is "_thread_safe_nc" when it is imported
+        # into the iris.fileformats.netcdf.saver. Also we want to check that the
+        # compression kwargs are passed into the NetCDF4 createVariable method
+        patch = self.patch(
+            "iris.fileformats.netcdf.saver._thread_safe_nc.DatasetWrapper.createVariable",
+        )
+        # No need to patch this NetCDF4 variable to compensate for the previous patch
+        # on createVariable, which doesn't actually create the variable.
+        self.patch(
+            "iris.fileformats.netcdf.saver._thread_safe_nc.DatasetWrapper.variables"
+        )
+        cube = make_cube(var_name=(var_name := "a"))
+        compression_kwargs = {
+            "complevel": 9,
+            "fletcher32": True,
+            "shuffle": True,
+            "zlib": True,
+        }
+
+        _ = self.check_save_cubes(cube, compression_kwargs=compression_kwargs)
+
+        # The following mesh components and cube should be compressed on serialization.
+        result = filter_compression_calls(patch, compression_kwargs)
+        expected = {"node_x", "node_y", "face_x", "face_y", "mesh2d_faces", var_name}
+        assert result == expected
+        # The primary mesh variable (no payload) is never compressed.
+        result = filter_compression_calls(patch, compression_kwargs, mismatch=True)
+        expected = {"Mesh2d"}
+        assert result == expected
+
     def test_basic_mesh(self):
         # Save a small mesh example and check aspects of the resulting file.
         cube = make_cube()  # A simple face-mapped data example.
@@ -686,7 +756,7 @@ def setUpClass(cls):
     def tearDownClass(cls):
         shutil.rmtree(cls.temp_dir)
 
-    def check_save_mesh(self, mesh):
+    def check_save_mesh(self, mesh, compression_kwargs=None):
         """Write a mesh to a new file in the common temporary directory.
 
         Use a name unique to this testcase, to avoid any clashes.
@@ -698,11 +768,47 @@ def check_save_mesh(self, mesh):
         # in the common temporary directory.
         tempfile_path = self.temp_dir / Path(tempfile_path).name
 
+        if compression_kwargs is None:
+            compression_kwargs = {}
+
         # Save data to the file.
-        save_mesh(mesh, tempfile_path)
+        save_mesh(mesh, tempfile_path, **compression_kwargs)
 
         return tempfile_path
 
+    def test_compression(self):
+        """Test NetCDF serialization of a mesh using compression.
+
+        NetCDF data compression keyword arguments include "complevel",
+        "fletcher32", "shuffle" and "zlib". Note that "complevel" and "shuffle"
+        are only applicable when "zlib=True".
+
+        """
+        patch = self.patch(
+            "iris.fileformats.netcdf.saver._thread_safe_nc.DatasetWrapper.createVariable",
+        )
+        self.patch(
+            "iris.fileformats.netcdf.saver._thread_safe_nc.DatasetWrapper.variables"
+        )
+        mesh = make_mesh()
+        compression_kwargs = {
+            "complevel": 9,
+            "fletcher32": True,
+            "shuffle": True,
+            "zlib": True,
+        }
+
+        _ = self.check_save_mesh(mesh, compression_kwargs=compression_kwargs)
+
+        # The following mesh components should be compressed on serialization.
+        result = filter_compression_calls(patch, compression_kwargs)
+        expected = {"node_x", "node_y", "face_x", "face_y", "mesh2d_faces"}
+        assert result == expected
+        # The primary mesh variable (no payload) is never compressed.
+        result = filter_compression_calls(patch, compression_kwargs, mismatch=True)
+        expected = {"Mesh2d"}
+        assert result == expected
+
     def test_connectivity_dim_order(self):
         """Test a mesh with some connectivities in the 'other' order.