Extend Cube Metadata NetCDF zlib compression (#6552)

* Extend Cube Metadata NetCDF zlib compression

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* add whatsnew entry

* review actions

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: bjlittle

docs/src/whatsnew/latest.rst

@@ -55,6 +55,10 @@ This document explains the changes made to Iris for this release
    the :data:`iris.FUTURE` object, to avoid breaking existing code.
    (:issue:`3678`, :pull:`6481`, :pull:`6540`)
 
+#. `@bjlittle`_ extended ``zlib`` compression of :class:`~iris.cube.Cube` data payload when saving to NetCDF
+   to also include any auxiliary coordinates and ancillary variables with the same ``shape``.
+   (:issue:`6539`, :pull:`6552`)
+
 
 🐛 Bugs Fixed
 =============

lib/iris/fileformats/netcdf/saver.py

@@ -586,21 +586,28 @@ def write(
         # 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.
+        compression_kwargs = {
+            "complevel": complevel,
+            "fletcher32": fletcher32,
+            "shuffle": shuffle,
+            "zlib": zlib,
+        }
+
         # Create the associated cube CF-netCDF data variable.
         cf_var_cube = self._create_cf_data_variable(
             cube,
             cube_dimensions,
-            local_keys,
-            zlib=zlib,
-            complevel=complevel,
-            shuffle=shuffle,
-            fletcher32=fletcher32,
+            local_keys=local_keys,
+            packing=packing,
+            fill_value=fill_value,
             contiguous=contiguous,
             chunksizes=chunksizes,
             endian=endian,
             least_significant_digit=least_significant_digit,
-            packing=packing,
-            fill_value=fill_value,
+            **compression_kwargs,
         )
 
         # Associate any mesh with the data-variable.
@@ -615,15 +622,19 @@ def write(
 
         # Add the auxiliary coordinate variables and associate the data
         # variable to them
-        self._add_aux_coords(cube, cf_var_cube, cube_dimensions)
+        self._add_aux_coords(
+            cube, cf_var_cube, cube_dimensions, compression_kwargs=compression_kwargs
+        )
 
         # Add the cell_measures variables and associate the data
         # variable to them
         self._add_cell_measures(cube, cf_var_cube, cube_dimensions)
 
         # Add the ancillary_variables variables and associate the data variable
         # to them
-        self._add_ancillary_variables(cube, cf_var_cube, cube_dimensions)
+        self._add_ancillary_variables(
+            cube, cf_var_cube, cube_dimensions, compression_kwargs=compression_kwargs
+        )
 
         # Add the formula terms to the appropriate cf variables for each
         # aux factory in the cube.
@@ -883,7 +894,14 @@ def _add_mesh(self, cube_or_mesh):
         return cf_mesh_name
 
     def _add_inner_related_vars(
-        self, cube, cf_var_cube, dimension_names, coordlike_elements
+        self,
+        cube,
+        cf_var_cube,
+        dimension_names,
+        coordlike_elements,
+        /,
+        *,
+        compression_kwargs=None,
     ):
         """Create a set of variables for aux-coords, ancillaries or cell-measures.
 
@@ -913,7 +931,10 @@ def _add_inner_related_vars(
                 if cf_name is None:
                     # Not already present : create it
                     cf_name = self._create_generic_cf_array_var(
-                        cube, dimension_names, element
+                        cube,
+                        dimension_names,
+                        element,
+                        compression_kwargs=compression_kwargs,
                     )
                     self._name_coord_map.append(cf_name, element)
 
@@ -929,7 +950,9 @@ def _add_inner_related_vars(
                 variable_names = " ".join(sorted(element_names))
                 _setncattr(cf_var_cube, role_attribute_name, variable_names)
 
-    def _add_aux_coords(self, cube, cf_var_cube, dimension_names):
+    def _add_aux_coords(
+        self, cube, cf_var_cube, dimension_names, /, *, compression_kwargs=None
+    ):
         """Add aux. coordinate to the dataset and associate with the data variable.
 
         Parameters
@@ -940,6 +963,9 @@ def _add_aux_coords(self, cube, cf_var_cube, dimension_names):
             A cf variable cube representation.
         dimension_names : list
             Names associated with the dimensions of the cube.
+        compression_kwargs : dict, optional
+            NetCDF data compression keyword arguments.
+
         """
         from iris.mesh.components import (
             MeshEdgeCoords,
@@ -967,6 +993,7 @@ def _add_aux_coords(self, cube, cf_var_cube, dimension_names):
             cf_var_cube,
             dimension_names,
             coords_to_add,
+            compression_kwargs=compression_kwargs,
         )
 
     def _add_cell_measures(self, cube, cf_var_cube, dimension_names):
@@ -988,7 +1015,9 @@ def _add_cell_measures(self, cube, cf_var_cube, dimension_names):
             cube.cell_measures(),
         )
 
-    def _add_ancillary_variables(self, cube, cf_var_cube, dimension_names):
+    def _add_ancillary_variables(
+        self, cube, cf_var_cube, dimension_names, /, *, compression_kwargs=None
+    ):
         """Add ancillary variables measures to the dataset and associate with the data variable.
 
         Parameters
@@ -999,12 +1028,16 @@ def _add_ancillary_variables(self, cube, cf_var_cube, dimension_names):
             A cf variable cube representation.
         dimension_names : list
             Names associated with the dimensions of the cube.
+        compression_kwargs : dict, optional
+            NetCDF data compression keyword arguments.
+
         """
         return self._add_inner_related_vars(
             cube,
             cf_var_cube,
             dimension_names,
             cube.ancillary_variables(),
+            compression_kwargs=compression_kwargs,
         )
 
     def _add_dim_coords(self, cube, dimension_names):
@@ -1439,7 +1472,7 @@ def _ensure_valid_dtype(self, values, src_name, src_object):
             values = values.astype(np.int32)
         return values
 
-    def _create_cf_bounds(self, coord, cf_var, cf_name):
+    def _create_cf_bounds(self, coord, cf_var, cf_name, /, *, compression_kwargs=None):
         """Create the associated CF-netCDF bounds variable.
 
         Parameters
@@ -1450,13 +1483,18 @@ def _create_cf_bounds(self, coord, cf_var, cf_name):
             CF-netCDF variable.
         cf_name : str
             Name of the CF-NetCDF variable.
+        compression_kwargs : dict, optional
+            NetCDF data compression keyword arguments.
 
         Returns
         -------
         None
 
         """
         if hasattr(coord, "has_bounds") and coord.has_bounds():
+            if compression_kwargs is None:
+                compression_kwargs = {}
+
             # Get the values in a form which is valid for the file format.
             bounds = self._ensure_valid_dtype(
                 coord.core_bounds(), "the bounds of coordinate", coord
@@ -1489,6 +1527,7 @@ def _create_cf_bounds(self, coord, cf_var, cf_name):
                 boundsvar_name,
                 bounds.dtype.newbyteorder("="),
                 cf_var.dimensions + (bounds_dimension_name,),
+                **compression_kwargs,
             )
             self._lazy_stream_data(data=bounds, cf_var=cf_var_bounds)
 
@@ -1685,8 +1724,11 @@ def _create_generic_cf_array_var(
         cube_or_mesh,
         cube_dim_names,
         element,
+        /,
+        *,
         element_dims=None,
         fill_value=None,
+        compression_kwargs=None,
     ):
         """Create theCF-netCDF variable given dimensional_metadata.
 
@@ -1718,6 +1760,8 @@ def _create_generic_cf_array_var(
             If not set, standard netcdf4-python behaviour : the variable has no
             '_FillValue' property, and uses the "standard" fill-value for its
             type.
+        compression_kwargs : dict, optional
+            NetCDF data compression keyword arguments.
 
         Returns
         -------
@@ -1732,6 +1776,9 @@ def _create_generic_cf_array_var(
         else:
             cube = None
 
+        if compression_kwargs is None:
+            compression_kwargs = {}
+
         # Work out the var-name to use.
         # N.B. the only part of this routine that may use a mesh _or_ a cube.
         cf_name = self._get_coord_variable_name(cube_or_mesh, element)
@@ -1749,6 +1796,9 @@ 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:
+            compression_kwargs = {}
+
         if np.issubdtype(data.dtype, np.str_):
             # Deal with string-type variables.
             # Typically CF label variables, but also possibly ancil-vars ?
@@ -1811,6 +1861,7 @@ def _create_generic_cf_array_var(
                 data.dtype.newbyteorder("="),
                 element_dims,
                 fill_value=fill_value,
+                **compression_kwargs,
             )
 
             # Add the axis attribute for spatio-temporal CF-netCDF coordinates.
@@ -1820,7 +1871,9 @@ def _create_generic_cf_array_var(
                     _setncattr(cf_var, "axis", axis.upper())
 
             # Create the associated CF-netCDF bounds variable, if any.
-            self._create_cf_bounds(element, cf_var, cf_name)
+            self._create_cf_bounds(
+                element, cf_var, cf_name, compression_kwargs=compression_kwargs
+            )
 
         # Add the data to the CF-netCDF variable.
         self._lazy_stream_data(data=data, cf_var=cf_var)
@@ -2784,16 +2837,16 @@ def is_valid_packspec(p):
         for cube, packspec, fill_value in zip(cubes, packspecs, fill_values):
             sman.write(
                 cube,
-                local_keys,
-                unlimited_dimensions,
-                zlib,
-                complevel,
-                shuffle,
-                fletcher32,
-                contiguous,
-                chunksizes,
-                endian,
-                least_significant_digit,
+                local_keys=local_keys,
+                unlimited_dimensions=unlimited_dimensions,
+                zlib=zlib,
+                complevel=complevel,
+                shuffle=shuffle,
+                fletcher32=fletcher32,
+                contiguous=contiguous,
+                chunksizes=chunksizes,
+                endian=endian,
+                least_significant_digit=least_significant_digit,
                 packing=packspec,
                 fill_value=fill_value,
             )