Save GDT20 - Polar Stereographic - take 2 (#405)

* Add saver for GDT20

* Save GDT20 (Polar Stereographic)

* License headers!

* Fix bits missed in branch update.

* Address @pp-mo comment on 0x80.

* Remove comment.

* Fix flake8.

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

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

* Fix missing bracket.

* Correct test value after ff8fc0d.

* Further fixes.

* Correct expected projectionCentreFlag to 128 following ff8fc0d.

* Handle both Stereographic and PolarStereographic appropriately.

* Add a GDT20 roundtrip test.

* Fix tests.

* Don't flake8 type hints.

* Mypy - remove six import.

* Revert "Don't flake8 type hints."

This reverts commit 32f6806.

* Exclude test_grid_definition_section.py from pep8 - E701 type hints.

* More translation errors.

* Fix tests.

* Remove PEP8 exception file.

* Add release notes entry.

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

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

* Approximate equality check.

* Allow for scale factor of 1.0 .

* Common non-standard scale factor error.

* Better TranslationError for non-polar.

* Fix test.

---------

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

docs/ref/release_notes.rst

@@ -23,6 +23,11 @@ Features
   `(ISSUE#330) <https://github.com/SciTools/iris-grib/issues/330>`_,
   `(PR#402) <https://github.com/SciTools/iris-grib/pull/402>`_
 
+* `@DPeterK <https://github.com/DPeterK>`_ and
+  `@trexfeathers <https://github.com/trexfeathers>`_ added saving support for
+  grid definition template 20 - polar stereographic.
+  `(PR#405) <https://github.com/SciTools/iris-grib/pull/405>`_
+
 Documentation
 ^^^^^^^^^^^^^
 * `@pp-mo <https://github.com/pp-mo>`_ reworked the main docs page to :

iris_grib/__init__.py

@@ -413,8 +413,11 @@ def _compute_extra_keys(self):
             else:
                 raise TranslationError("Unhandled projectionCentreFlag")
 
+            # Always load PolarStereographic - never Stereographic.
+            #  Stereographic is a CF/Iris concept and not something described
+            #  in GRIB.
             # Note: I think the grib api defaults LaDInDegrees to 60 for grib1.
-            self.extra_keys["_coord_system"] = coord_systems.Stereographic(
+            self.extra_keys["_coord_system"] = coord_systems.PolarStereographic(
                 pole_lat,
                 self.orientationOfTheGridInDegrees,
                 0,

iris_grib/_load_convert.py

@@ -912,7 +912,10 @@ def grid_definition_template_20(section, metadata):
         central_lat = 90.0
     central_lon = section["orientationOfTheGrid"] * _GRID_ACCURACY_IN_DEGREES
     true_scale_lat = section["LaD"] * _GRID_ACCURACY_IN_DEGREES
-    cs = icoord_systems.Stereographic(
+    # Always load PolarStereographic - never Stereographic.
+    #  Stereographic is a CF/Iris concept and not something described
+    #  in GRIB.
+    cs = icoord_systems.PolarStereographic(
         central_lat=central_lat,
         central_lon=central_lon,
         true_scale_lat=true_scale_lat,

iris_grib/_save_rules.py

@@ -27,6 +27,7 @@
     TransverseMercator,
     LambertConformal,
     LambertAzimuthalEqualArea,
+    Stereographic,
 )
 from iris.exceptions import TranslationError
 
@@ -544,43 +545,55 @@ def m_to_cm(value):
     eccodes.codes_set(grib, "scaleFactorAtReferencePoint", value)
 
 
-def grid_definition_template_30(cube, grib):
+def _perspective_projection_common(cube, grib):
     """
-    Set keys within the provided grib message based on
-    Grid Definition Template 3.30.
-
-    Template 3.30 is used to represent a Lambert Conformal grid.
+    Common functionality for setting grid keys for the perspective projection
+    grid definition templates (20 and 30; Polar Stereographic and Lambert
+    Conformal.
 
     """
-
-    eccodes.codes_set(grib, "gridDefinitionTemplateNumber", 30)
-
     # Retrieve some information from the cube.
     y_coord = cube.coord(dimensions=[0])
     x_coord = cube.coord(dimensions=[1])
     cs = y_coord.coord_system
+    cs_name = cs.grid_mapping_name.replace("_", " ").title()
+
+    both_falses_zero = all(
+        np.isclose(f, 0.0, atol=1e-6) for f in (cs.false_easting, cs.false_northing)
+    )
+    if not both_falses_zero:
+        message = (
+            f"{cs.false_easting=}, {cs.false_northing=} . Non-zero "
+            f"unsupported for {cs_name}."
+        )
+        raise TranslationError(message)
 
     # Normalise the coordinate values to millimetres - the resolution
     # used in the GRIB message.
     y_mm = points_in_unit(y_coord, "mm")
     x_mm = points_in_unit(x_coord, "mm")
 
     # Encode the horizontal points.
-
     # NB. Since we're already in millimetres, our tolerance for
     # discrepancy in the differences is 1.
     try:
         x_step = step(x_mm, atol=1)
         y_step = step(y_mm, atol=1)
     except ValueError:
-        msg = "Irregular coordinates not supported for Lambert " "Conformal."
-        raise TranslationError(msg)
+        msg = "Irregular coordinates not supported for {!r}."
+        raise TranslationError(msg.format(cs_name))
     eccodes.codes_set(grib, "Dx", abs(x_step))
     eccodes.codes_set(grib, "Dy", abs(y_step))
 
     horizontal_grid_common(cube, grib, xy=True)
 
-    # Transform first point into geographic CS
+    eccodes.codes_set(
+        grib,
+        "resolutionAndComponentFlags",
+        0x1 << _RESOLUTION_AND_COMPONENTS_GRID_WINDS_BIT,
+    )
+
+    # Transform first point into geographic CS.
     geog = cs.ellipsoid if cs.ellipsoid is not None else GeogCS(1)
     first_x, first_y = geog.as_cartopy_crs().transform_point(
         x_coord.points[0], y_coord.points[0], cs.as_cartopy_crs()
@@ -600,20 +613,85 @@ def grid_definition_template_30(cube, grib):
     )
     eccodes.codes_set(grib, "LaD", cs.central_lat / _DEFAULT_DEGREES_UNITS)
     eccodes.codes_set(grib, "LoV", central_lon / _DEFAULT_DEGREES_UNITS)
+
+
+def grid_definition_template_20(cube, grib):
+    """
+    Set keys within the provided GRIB message based on
+    Grid Definition Template 3.20.
+
+    Template 3.20 is used to represent a Polar Stereographic grid.
+
+    """
+    eccodes.codes_set(grib, "gridDefinitionTemplateNumber", 20)
+
+    # Retrieve the cube coord system.
+    cs = cube.coord(dimensions=[0]).coord_system
+
+    _perspective_projection_common(cube, grib)
+
+    # Is this a north or south polar stereographic projection?
+    if np.isclose(cs.central_lat, -90.0):
+        centre_flag = 0x80
+    elif np.isclose(cs.central_lat, 90.0):
+        centre_flag = 0x0
+    else:
+        message = (
+            f"{cs.central_lat=} . GRIB Template 3.20 only supports the polar "
+            "stereographic projection; central_lat must be 90.0 or -90.0."
+        )
+        raise TranslationError(message)
+
+    def non_standard_scale_factor_error(message: str):
+        message = f"Non-standard scale factor specified. {message}"
+        raise TranslationError(message)
+
+    if cs.true_scale_lat and cs.true_scale_lat != cs.central_lat:
+        non_standard_scale_factor_error(
+            f"{cs.true_scale_lat=}, {cs.central_lat=} . iris_grib can "
+            "only write a GRIB Template 3.20 file where these are identical."
+        )
+
+    if cs.scale_factor_at_projection_origin and not np.isclose(
+        cs.scale_factor_at_projection_origin, 1.0
+    ):
+        non_standard_scale_factor_error(
+            f"{cs.scale_factor_at_projection_origin=} . iris_grib cannot "
+            "write scale_factor_at_projection_origin to a GRIB Template 3.20 "
+            "file."
+        )
+
+    eccodes.codes_set(grib, "projectionCentreFlag", centre_flag)
+
+
+def grid_definition_template_30(cube, grib):
+    """
+    Set keys within the provided grib message based on
+    Grid Definition Template 3.30.
+
+    Template 3.30 is used to represent a Lambert Conformal grid.
+
+    """
+
+    eccodes.codes_set(grib, "gridDefinitionTemplateNumber", 30)
+
+    # Retrieve the cube coord system.
+    cs = cube.coord(dimensions=[0]).coord_system
+
+    _perspective_projection_common(cube, grib)
+
+    # Which pole are the parallels closest to? That is the direction
+    # that the cone converges.
     latin1, latin2 = cs.secant_latitudes
     eccodes.codes_set(grib, "Latin1", latin1 / _DEFAULT_DEGREES_UNITS)
     eccodes.codes_set(grib, "Latin2", latin2 / _DEFAULT_DEGREES_UNITS)
-    eccodes.codes_set(
-        grib,
-        "resolutionAndComponentFlags",
-        0x1 << _RESOLUTION_AND_COMPONENTS_GRID_WINDS_BIT,
-    )
 
     # Which pole are the parallels closest to? That is the direction
     # that the cone converges.
-    poliest_sec = latin1 if abs(latin1) > abs(latin2) else latin2
-    centre_flag = 0x0 if poliest_sec > 0 else 0x1
+    poliest_secant = latin1 if abs(latin1) > abs(latin2) else latin2
+    centre_flag = 0x0 if poliest_secant > 0 else 0x80
     eccodes.codes_set(grib, "projectionCentreFlag", centre_flag)
+
     eccodes.codes_set(grib, "latitudeOfSouthernPole", 0)
     eccodes.codes_set(grib, "longitudeOfSouthernPole", 0)
 
@@ -729,6 +807,11 @@ def grid_definition_section(cube, grib):
         # Transverse Mercator coordinate system (template 3.12).
         grid_definition_template_12(cube, grib)
 
+    elif isinstance(cs, Stereographic):
+        # Stereographic coordinate system (template 3.20).
+        # This will also handle the PolarStereographic subclass.
+        grid_definition_template_20(cube, grib)
+
     elif isinstance(cs, LambertConformal):
         # Lambert Conformal coordinate system (template 3.30).
         grid_definition_template_30(cube, grib)

iris_grib/tests/integration/load_convert/test_sample_file_loads.py

@@ -75,7 +75,10 @@ def test_polar_stereo_grib2_grid_definition(self):
         self.assertAlmostEqual(pyc.points.max(), -216.1459, places=4)
         self.assertAlmostEqual(pyc.points.min(), -5970216.1459, places=4)
         self.assertEqual(pyc.coord_system, pxc.coord_system)
-        self.assertEqual(pyc.coord_system.grid_mapping_name, "stereographic")
+        self.assertEqual(
+            pyc.coord_system.grid_mapping_name,
+            "polar_stereographic",
+        )
         self.assertEqual(pyc.coord_system.central_lat, 90.0)
         self.assertEqual(pyc.coord_system.central_lon, 249.0)
         self.assertEqual(pyc.coord_system.false_easting, 0.0)

iris_grib/tests/integration/round_trip/test_grid_definition_section.py

@@ -11,8 +11,14 @@
 # before importing anything else.
 import iris_grib.tests as tests
 
+from typing import Literal
+
+import numpy as np
+
 from iris import load_cube, save
-from iris.coord_systems import RotatedGeogCS
+from iris.coord_systems import GeogCS, PolarStereographic, RotatedGeogCS, Stereographic
+from iris.coords import DimCoord
+from iris.cube import Cube
 from iris.fileformats.pp import EARTH_RADIUS as UM_DEFAULT_EARTH_RADIUS
 from iris.util import is_regular
 
@@ -124,5 +130,106 @@ def test_save_load(self):
         self.assertArrayAllClose(cube.data, cube_loaded_from_saved.data)
 
 
+class GDT20Common:
+    """'Wrapper' to avoid class being detected during test runs."""
+
+    class GDT20Common(tests.TestGribMessage):
+        """Roundtrip testing that attributes save+load correctly."""
+
+        # Enable subclassing to test different permutations.
+        coord_system_class: type[Stereographic]
+        pole: Literal["north", "south"]
+
+        def setUp(self):
+            # Create a Cube to save, inspired by
+            #  iris-sample-data toa_brightness_stereographic.nc
+            if self.pole == "north":
+                central_lat = 90
+            elif self.pole == "south":
+                central_lat = -90
+            else:
+                raise NotImplementedError(f"Invalid pole: {self.pole}")
+
+            self.coord_system_kwargs = dict(
+                central_lat=central_lat,
+                central_lon=325,
+                true_scale_lat=central_lat,
+                ellipsoid=GeogCS(6378169.0),
+            )
+            coord_system = self.coord_system_class(**self.coord_system_kwargs)
+
+            coord_kwargs = dict(
+                units="m",
+                coord_system=coord_system,
+            )
+            coord_x = DimCoord(
+                np.linspace(-2250000, 6750192, 256, endpoint=False),
+                standard_name="projection_x_coordinate",
+                **coord_kwargs,
+            )
+            coord_y = DimCoord(
+                np.linspace(-980000, -6600000, 160, endpoint=False),
+                standard_name="projection_y_coordinate",
+                **coord_kwargs,
+            )
+            coord_t = DimCoord(
+                0, standard_name="time", units="hours since 1970-01-01 00:00:00"
+            )
+            coord_fp = DimCoord(0, standard_name="forecast_period", units="hours")
+            coord_frt = DimCoord(
+                0, standard_name="forecast_reference_time", units=coord_t.units
+            )
+            shape = (coord_y.shape[0], coord_x.shape[0])
+            self.cube = Cube(
+                np.arange(np.prod(shape), dtype=float).reshape(shape),
+                dim_coords_and_dims=[(coord_y, 0), (coord_x, 1)],
+                aux_coords_and_dims=[
+                    (coord_t, None),
+                    (coord_fp, None),
+                    (coord_frt, None),
+                ],
+            )
+
+        def test_save_load(self):
+            with self.temp_filename("polar_stereo.grib2") as temp_file_path:
+                save(self.cube, temp_file_path)
+                cube_reloaded = load_cube(temp_file_path)
+                # Touch the data before destroying the file.
+                _ = cube_reloaded.data
+
+            cube_expected = self.cube.copy()
+            for coord in cube_expected.dim_coords:
+                # GRIB only describes PolarStereographic, so we always expect
+                #  that system even when we started with Stereographic.
+                coord.coord_system = PolarStereographic(**self.coord_system_kwargs)
+
+            # Modifications to remove irrelevant inequalities.
+            del cube_reloaded.attributes["GRIB_PARAM"]
+            for coord in cube_reloaded.dim_coords:
+                coord.points = np.round(coord.points)
+
+            self.assertEqual(cube_expected, cube_reloaded)
+
+
+class TestGDT20StereoNorth(GDT20Common.GDT20Common):
+    coord_system_class = Stereographic
+    pole = "north"
+
+
+class TestGDT20StereoSouth(GDT20Common.GDT20Common):
+    coord_system_class = Stereographic
+    pole = "south"
+
+
+class TestGDT20PolarNorth(GDT20Common.GDT20Common):
+    coord_system_class = PolarStereographic
+    pole = "north"
+
+
+class TestGDT20PolarSouth(GDT20Common.GDT20Common):
+    coord_system_class = PolarStereographic
+    pole = "south"
+
+
 if __name__ == "__main__":
     tests.main()

iris_grib/tests/integration/save_rules/test_grid_definition_section.py

@@ -20,6 +20,8 @@
     LambertConformal,
     AlbersEqualArea,
     LambertAzimuthalEqualArea,
+    Stereographic,
+    PolarStereographic,
 )
 import numpy as np
 
@@ -92,6 +94,26 @@ def test_grid_definition_template_12(self):
         grid_definition_section(test_cube, self.mock_grib)
         self._check_key("gridDefinitionTemplateNumber", 12)
 
+    def grid_definition_template_20_common(self, coord_system: type[Stereographic]):
+        # Stereographic grid.
+        # Common code to allow testing PolarStereographic and Stereographic.
+        if not issubclass(coord_system, Stereographic):
+            raise ValueError("coord_system must be Stereographic type.")
+
+        x_points = np.arange(3)
+        y_points = np.arange(3)
+        coord_units = "m"
+        cs = coord_system(90.0, 0, ellipsoid=self.ellipsoid)
+        test_cube = self._make_test_cube(cs, x_points, y_points, coord_units)
+        grid_definition_section(test_cube, self.mock_grib)
+        self._check_key("gridDefinitionTemplateNumber", 20)
+
+    def test_grid_definition_template_20_s(self):
+        self.grid_definition_template_20_common(Stereographic)
+
+    def test_grid_definition_template_20_ps(self):
+        self.grid_definition_template_20_common(PolarStereographic)
+
     def test_grid_definition_template_30(self):
         # Lambert Conformal grid.
         x_points = np.arange(3)