diff --git a/iris_grib/_load_convert.py b/iris_grib/_load_convert.py index 1185f222..b99546a0 100644 --- a/iris_grib/_load_convert.py +++ b/iris_grib/_load_convert.py @@ -26,13 +26,13 @@ from iris.coords import AuxCoord, DimCoord, CellMethod from iris.exceptions import TranslationError from . import grib_phenom_translation as itranslation +from .grib_phenom_translation import GRIBCode from iris.fileformats.rules import ConversionMetadata, Factory, Reference, \ ReferenceTarget from iris.util import _is_circular from ._iris_mercator_support import confirm_extended_mercator_supported from ._grib1_load_rules import grib1_convert -from .message import GribMessage # Restrict the names imported from this namespace. @@ -1378,6 +1378,13 @@ def translate_phenomenon(metadata, discipline, parameterCategory, metadata['long_name'] = long_name metadata['units'] = Unit(1) + # Add a standard attribute recording the grib phenomenon identity. + metadata['attributes']['GRIB_PARAM'] = GRIBCode( + edition_or_string=2, + discipline=discipline, + category=parameterCategory, + number=parameterNumber) + # Identify hybrid height and pressure reference fields. # Look for fields at surface level first. if (typeOfFirstFixedSurface == 1 and diff --git a/iris_grib/_save_rules.py b/iris_grib/_save_rules.py index 563c8f8c..e3c35293 100644 --- a/iris_grib/_save_rules.py +++ b/iris_grib/_save_rules.py @@ -30,6 +30,7 @@ from . import grib_phenom_translation as gptx from ._load_convert import (_STATISTIC_TYPE_NAMES, _TIME_RANGE_UNITS, _SPATIAL_PROCESSING_TYPES) +from .grib_phenom_translation import GRIBCode from iris.util import is_regular, regular_step @@ -690,22 +691,45 @@ def grid_definition_section(cube, grib): ############################################################################### def set_discipline_and_parameter(cube, grib): - # NOTE: for now, can match by *either* standard_name or long_name. - # This allows workarounds for data with no identified standard_name. - grib2_info = gptx.cf_phenom_to_grib2_info(cube.standard_name, - cube.long_name) - if grib2_info is not None: - gribapi.grib_set(grib, "discipline", grib2_info.discipline) - gribapi.grib_set(grib, "parameterCategory", grib2_info.category) - gribapi.grib_set(grib, "parameterNumber", grib2_info.number) - else: - gribapi.grib_set(grib, "discipline", 255) - gribapi.grib_set(grib, "parameterCategory", 255) - gribapi.grib_set(grib, "parameterNumber", 255) + # Default values for parameter identity keys = effectively "MISSING". + discipline, category, number = 255, 255, 255 + identity_found = False + + # First, see if we can find and interpret a 'GRIB_PARAM' attribute. + attr = cube.attributes.get('GRIB_PARAM', None) + if attr: + try: + # Convert to standard tuple-derived form. + gc = GRIBCode(attr) + if gc.edition == 2: + discipline = gc.discipline + category = gc.category + number = gc.number + identity_found = True + except: + pass + + if not identity_found: + # Else, translate a cube phenomenon, if possible. + # NOTE: for now, can match by *either* standard_name or long_name. + # This allows workarounds for data with no identified standard_name. + grib2_info = gptx.cf_phenom_to_grib2_info(cube.standard_name, + cube.long_name) + if grib2_info is not None: + discipline = grib2_info.discipline + category = grib2_info.category + number = grib2_info.number + identity_found = True + + if not identity_found: warnings.warn('Unable to determine Grib2 parameter code for cube.\n' 'discipline, parameterCategory and parameterNumber ' 'have been set to "missing".') + gribapi.grib_set(grib, "discipline", discipline) + gribapi.grib_set(grib, "parameterCategory", category) + gribapi.grib_set(grib, "parameterNumber", number) + def _non_missing_forecast_period(cube): # Calculate "model start time" to use as the reference time. diff --git a/iris_grib/grib_phenom_translation.py b/iris_grib/grib_phenom_translation.py index cee9cb29..988a1464 100644 --- a/iris_grib/grib_phenom_translation.py +++ b/iris_grib/grib_phenom_translation.py @@ -16,8 +16,8 @@ * cf --> grib2 ''' - -import collections +from collections import namedtuple +import re import warnings import cf_units @@ -54,12 +54,12 @@ def __setitem__(self, key, value): # Define namedtuples for keys+values of the Grib1 lookup table. -_Grib1ToCfKeyClass = collections.namedtuple( +_Grib1ToCfKeyClass = namedtuple( 'Grib1CfKey', ('table2_version', 'centre_number', 'param_number')) # NOTE: this form is currently used for both Grib1 *and* Grib2 -_GribToCfDataClass = collections.namedtuple( +_GribToCfDataClass = namedtuple( 'Grib1CfData', ('standard_name', 'long_name', 'units', 'set_height')) @@ -146,7 +146,7 @@ def _make_grib1_cf_entry(table2_version, centre_number, param_number, # Define a namedtuple for the keys of the Grib2 lookup table. -_Grib2ToCfKeyClass = collections.namedtuple( +_Grib2ToCfKeyClass = namedtuple( 'Grib2CfKey', ('param_discipline', 'param_category', 'param_number')) @@ -205,11 +205,11 @@ def _make_grib2_cf_entry(param_discipline, param_category, param_number, # Define namedtuples for key+values of the cf-to-grib2 lookup table. -_CfToGrib2KeyClass = collections.namedtuple( +_CfToGrib2KeyClass = namedtuple( 'CfGrib2Key', ('standard_name', 'long_name')) -_CfToGrib2DataClass = collections.namedtuple( +_CfToGrib2DataClass = namedtuple( 'CfGrib2Data', ('discipline', 'category', 'number', 'units')) @@ -316,3 +316,70 @@ def cf_phenom_to_grib2_info(standard_name, long_name=None): if standard_name is not None: long_name = None return _CF_GRIB2_TABLE[(standard_name, long_name)] + + +class GRIBCode(namedtuple('GRIBCode', + 'edition discipline category number')): + """ + An object representing a specific Grib phenomenon identity. + + Basically a namedtuple of (edition, discipline, category, number). + + Also provides a string representation, and supports creation from: another + similar object; a tuple of numbers; or any string with 4 separate decimal + numbers in it. + + """ + __slots__ = () + + def __new__(cls, edition_or_string, + discipline=None, category=None, number=None): + args = (edition_or_string, discipline, category, number) + nargs = sum(arg is not None for arg in args) + if nargs == 1: + # Single argument: convert to a string and extract 4 integers. + # NOTE: this also allows input from a GRIBCode, or a plain tuple. + edition_or_string = str(edition_or_string) + edition, discipline, category, number = \ + cls._fournums_from_gribcode_string(edition_or_string) + elif nargs == 4: + edition = edition_or_string + edition, discipline, category, number = [ + int(arg) + for arg in (edition, discipline, category, number)] + else: + msg = ('Cannot create GRIBCode from {} arguments, ' + '"GRIBCode{!r}" : ' + 'expected either 1 or 4 non-None arguments.') + raise ValueError(msg.format(nargs, args)) + + return super(GRIBCode, cls).__new__( + cls, edition, discipline, category, number) + + RE_PARSE_FOURNUMS = re.compile(4 * r'[^\d]*(\d*)') + + @classmethod + def _fournums_from_gribcode_string(cls, edcn_string): + parsed_ok = False + nums_match = cls.RE_PARSE_FOURNUMS.match(edcn_string).groups() + if nums_match is not None: + try: + nums = [int(grp) for grp in nums_match] + parsed_ok = True + except ValueError: + pass + + if not parsed_ok: + msg = ('Invalid argument for GRIBCode creation, ' + '"GRIBCode({!r})" : ' + 'requires 4 numbers, separated by non-numerals.') + raise ValueError(msg.format(edcn_string)) + + return nums + + PRINT_FORMAT = 'GRIB{:1d}:d{:03d}c{:03d}n{:03d}' + + def __str__(self): + result = self.PRINT_FORMAT.format( + self.edition, self.discipline, self.category, self.number) + return result diff --git a/iris_grib/tests/results/unit/load_cubes/load_cubes/reduced_raw.cml b/iris_grib/tests/results/unit/load_cubes/load_cubes/reduced_raw.cml index b4f27991..6846382f 100644 --- a/iris_grib/tests/results/unit/load_cubes/load_cubes/reduced_raw.cml +++ b/iris_grib/tests/results/unit/load_cubes/load_cubes/reduced_raw.cml @@ -2,6 +2,7 @@ + diff --git a/iris_grib/tests/unit/grib_phenom_translation/test_grib_phenom_translation.py b/iris_grib/tests/unit/grib_phenom_translation/test_grib_phenom_translation.py index ae3cb715..88a6e627 100644 --- a/iris_grib/tests/unit/grib_phenom_translation/test_grib_phenom_translation.py +++ b/iris_grib/tests/unit/grib_phenom_translation/test_grib_phenom_translation.py @@ -18,6 +18,7 @@ import cf_units import iris_grib.grib_phenom_translation as gptx +from iris_grib.grib_phenom_translation import GRIBCode class TestGribLookupTableType(tests.IrisTest): @@ -151,5 +152,113 @@ def check_cf_grib2(standard_name, long_name, expect_none=True) +class TestGRIBcode(tests.IrisTest): + # GRIBCode is basically a namedtuple, so not all behaviour needs testing. + # However, creation is a bit special so exercise all those cases. + def test_create_from_keys(self): + gribcode = GRIBCode( + edition_or_string=5, + discipline=7, + category=4, + number=199) + self.assertEqual(gribcode.edition, 5) + self.assertEqual(gribcode.discipline, 7) + self.assertEqual(gribcode.category, 4) + self.assertEqual(gribcode.number, 199) + + def test_create_from_args(self): + gribcode = GRIBCode(7, 3, 12, 99) + self.assertEqual(gribcode.edition, 7) + self.assertEqual(gribcode.discipline, 3) + self.assertEqual(gribcode.category, 12) + self.assertEqual(gribcode.number, 99) + + def test_create_is_copy(self): + gribcode1 = GRIBCode(7, 3, 12, 99) + gribcode2 = GRIBCode(7, 3, 12, 99) + self.assertEqual(gribcode1, gribcode2) + self.assertIsNot(gribcode1, gribcode2) + + def test_create_from_gribcode(self): + gribcode1 = GRIBCode((4, 3, 2, 1)) + gribcode2 = GRIBCode(gribcode1) + self.assertEqual(gribcode1, gribcode2) + # NOTE: *not* passthrough : it creates a copy + # (though maybe not too significant, as it is immutable anyway?) + self.assertIsNot(gribcode1, gribcode2) + + def test_create_from_string(self): + gribcode = GRIBCode('xxx12xs-34 -5,678qqqq') + # NOTE: args 2 and 3 are *not* negative. + self.assertEqual(gribcode, GRIBCode(12, 34, 5, 678)) + + def test_create_from_own_string(self): + # Check that GRIBCode string reprs are valid as create arguments. + gribcode = GRIBCode( + edition_or_string=2, + discipline=17, + category=94, + number=231) + grib_param_string = str(gribcode) + newcode = GRIBCode(grib_param_string) + self.assertEqual(newcode, gribcode) + + def test_create_from_tuple(self): + gribcode = GRIBCode((4, 3, 2, 1)) + self.assertEqual(gribcode, GRIBCode(4, 3, 2, 1)) + + def test_create_bad_nargs(self): + # Between 1 and 4 args is not invalid call syntax, but it should fail. + with self.assertRaisesRegex( + ValueError, + 'Cannot create GRIBCode from 2 arguments'): + GRIBCode(1, 2) + + def test_create_bad_single_arg_None(self): + with self.assertRaisesRegex( + ValueError, + 'Cannot create GRIBCode from 0 arguments'): + GRIBCode(None) + + def test_create_bad_single_arg_empty_string(self): + with self.assertRaisesRegex( + ValueError, + 'Invalid argument for GRIBCode creation'): + GRIBCode('') + + def test_create_bad_single_arg_nonums(self): + with self.assertRaisesRegex( + ValueError, + 'Invalid argument for GRIBCode creation'): + GRIBCode('saas- dsa- ') + + def test_create_bad_single_arg_less_than_4_nums(self): + with self.assertRaisesRegex( + ValueError, + 'Invalid argument for GRIBCode creation'): + GRIBCode('1,2,3') + + def test_create_bad_single_arg_number(self): + with self.assertRaisesRegex( + ValueError, + 'Invalid argument for GRIBCode creation'): + GRIBCode(4) + + def test_create_bad_single_arg_single_numeric(self): + with self.assertRaisesRegex( + ValueError, + 'Invalid argument for GRIBCode creation'): + GRIBCode('44') + + def test_create_string_more_than_4_nums(self): + # Note: does not error, just discards the extra. + gribcode = GRIBCode('1,2,3,4,5,6,7,8') + self.assertEqual(gribcode, GRIBCode(1, 2, 3, 4)) + + def test__str__(self): + result = str(GRIBCode(2, 17, 3, 123)) + self.assertEqual(result, 'GRIB2:d017c003n123') + + if __name__ == '__main__': tests.main() diff --git a/iris_grib/tests/unit/load_convert/test_translate_phenomenon.py b/iris_grib/tests/unit/load_convert/test_translate_phenomenon.py index dd1f22c1..6f79474f 100644 --- a/iris_grib/tests/unit/load_convert/test_translate_phenomenon.py +++ b/iris_grib/tests/unit/load_convert/test_translate_phenomenon.py @@ -18,7 +18,7 @@ from iris.coords import DimCoord from iris_grib._load_convert import Probability, translate_phenomenon -from iris_grib.grib_phenom_translation import _GribToCfDataClass +from iris_grib.grib_phenom_translation import _GribToCfDataClass, GRIBCode class Test_probability(tests.IrisGribTest): @@ -30,11 +30,11 @@ def setUp(self): return_value=_GribToCfDataClass('air_temperature', '', 'K', None)) # Construct dummy call arguments self.probability = Probability('', 22.0) - self.metadata = {'aux_coords_and_dims': []} + self.metadata = {'aux_coords_and_dims': [], 'attributes': {}} def test_basic(self): - result = translate_phenomenon(self.metadata, None, None, None, None, - None, None, probability=self.probability) + translate_phenomenon(self.metadata, 7, 8, 9, None, + None, None, probability=self.probability) # Check metadata. thresh_coord = DimCoord([22.0], standard_name='air_temperature', @@ -43,14 +43,23 @@ def test_basic(self): 'standard_name': None, 'long_name': 'probability_of_air_temperature_', 'units': Unit(1), - 'aux_coords_and_dims': [(thresh_coord, None)]}) + 'aux_coords_and_dims': [(thresh_coord, None)], + 'attributes': {'GRIB_PARAM': GRIBCode(2, 7, 8, 9)}}) def test_no_phenomenon(self): - original_metadata = deepcopy(self.metadata) self.phenom_lookup_patch.return_value = None - result = translate_phenomenon(self.metadata, None, None, None, None, - None, None, probability=self.probability) - self.assertEqual(self.metadata, original_metadata) + expected_metadata = self.metadata.copy() + translate_phenomenon(self.metadata, + discipline=7, + parameterCategory=77, + parameterNumber=777, + typeOfFirstFixedSurface=None, + scaledValueOfFirstFixedSurface=None, + typeOfSecondFixedSurface=None, + probability=self.probability) + expected_metadata['attributes']['GRIB_PARAM'] = \ + GRIBCode(2, 7, 77, 777) + self.assertEqual(self.metadata, expected_metadata) if __name__ == '__main__': diff --git a/iris_grib/tests/unit/save_rules/test_set_discipline_and_parameter.py b/iris_grib/tests/unit/save_rules/test_set_discipline_and_parameter.py new file mode 100644 index 00000000..b8071254 --- /dev/null +++ b/iris_grib/tests/unit/save_rules/test_set_discipline_and_parameter.py @@ -0,0 +1,80 @@ +# Copyright iris-grib contributors +# +# This file is part of iris-grib and is released under the LGPL license. +# See COPYING and COPYING.LESSER in the root of the repository for full +# licensing details. +"""Unit tests for `iris_grib.grib_save_rules.set_discipline_and_parameter`.""" +# Import iris_grib.tests first so that some things can be initialised before +# importing anything else. +import iris_grib.tests as tests +from unittest import mock + +from iris.cube import Cube + +from iris_grib.grib_phenom_translation import GRIBCode + +from iris_grib._save_rules import set_discipline_and_parameter + + +class TestPhenomenonCoding(tests.IrisGribTest): + def setUp(self): + # A mock cube with empty phenomenon-specifying metadata. + self.mock_cube = mock.Mock( + spec=Cube, + standard_name=None, + long_name=None, + attributes={}) + + def _check_coding(self, cube, discipline, paramCategory, paramNumber): + # Check that encoding 'cube' writes the expected phenomenon keys. + grib_set_patch = self.patch( + 'iris_grib._save_rules.gribapi.grib_set') + mock_message = mock.sentinel.grib2_message + + set_discipline_and_parameter(cube, mock_message) + + expected_calls = [ + mock.call(mock_message, "discipline", discipline), + mock.call(mock_message, "parameterCategory", paramCategory), + mock.call(mock_message, "parameterNumber", paramNumber)] + + self.assertEqual(grib_set_patch.call_args_list, expected_calls) + + def test_unknown_phenomenon(self): + cube = self.mock_cube + self._check_coding(cube, 255, 255, 255) + + def test_known_standard_name(self): + cube = self.mock_cube + cube.standard_name = 'sea_water_y_velocity' + self._check_coding(cube, 10, 1, 3) # as seen in _grib_cf_map.py + + def test_gribcode_attribute_object(self): + cube = self.mock_cube + cube.attributes = {'GRIB_PARAM': GRIBCode(2, 7, 12, 99)} + self._check_coding(cube, 7, 12, 99) + + def test_gribcode_attribute_string(self): + cube = self.mock_cube + cube.attributes = {'GRIB_PARAM': '2, 9, 33, 177'} + self._check_coding(cube, 9, 33, 177) + + def test_gribcode_attribute_tuple(self): + cube = self.mock_cube + cube.attributes = {'GRIB_PARAM': (2, 33, 4, 12)} + self._check_coding(cube, 33, 4, 12) + + def test_gribcode_attribute_not_edition_2(self): + cube = self.mock_cube + cube.attributes = {'GRIB_PARAM': GRIBCode(1, 7, 12, 99)} + self._check_coding(cube, 255, 255, 255) + + def test_gribcode_attribute_overrides_phenomenon(self): + cube = self.mock_cube + cube.standard_name = 'sea_water_y_velocity' + cube.attributes = {'GRIB_PARAM': '2, 9, 33, 177'} + self._check_coding(cube, 9, 33, 177) + + +if __name__ == "__main__": + tests.main()