From 7fcdf5a63f2ab0788c417bf66448be261af89b67 Mon Sep 17 00:00:00 2001 From: Sandor Kertesz Date: Tue, 9 Jan 2024 20:02:18 +0000 Subject: [PATCH] Add option to set generatingProcessIdentifier when writing GRIB data --- earthkit/data/readers/grib/codes.py | 7 +------ earthkit/data/readers/grib/output.py | 5 ++++- earthkit/data/writers/grib.py | 8 ++------ tests/numpy_fs/test_numpy_fs_write.py | 19 +++++-------------- 4 files changed, 12 insertions(+), 27 deletions(-) diff --git a/earthkit/data/readers/grib/codes.py b/earthkit/data/readers/grib/codes.py index 844e503a..a6fa0022 100644 --- a/earthkit/data/readers/grib/codes.py +++ b/earthkit/data/readers/grib/codes.py @@ -219,15 +219,10 @@ def get_longitudes(self, dtype=None): def get_data_points(self): return eccodes.codes_grib_get_data(self._handle) - def set_values(self, values, generating_proc_id=None): + def set_values(self, values): try: assert self.path is None, "Only cloned handles can have values changed" eccodes.codes_set_values(self._handle, values.flatten()) - # For ECMWF gribs: indicates that something has been modified (255=unknown) - if generating_proc_id is not None: - eccodes.codes_set_long( - self._handle, "generatingProcessIdentifier", generating_proc_id - ) except Exception as e: LOG.error("Error setting values") LOG.exception(e) diff --git a/earthkit/data/readers/grib/output.py b/earthkit/data/readers/grib/output.py index f365f8cc..bf582367 100644 --- a/earthkit/data/readers/grib/output.py +++ b/earthkit/data/readers/grib/output.py @@ -123,12 +123,15 @@ def write( k: v for k, v in sorted(metadata.items(), key=lambda x: order(x[0])) } + if "generatingProcessIdentifier" not in metadata: + metadata["generatingProcessIdentifier"] = 255 + LOG.debug("GribOutput.metadata %s", metadata) for k, v in metadata.items(): handle.set(k, v) - handle.set_values(values, generating_proc_id=255) + handle.set_values(values) file, path = self.f(handle) handle.write(file) diff --git a/earthkit/data/writers/grib.py b/earthkit/data/writers/grib.py index a456ca74..32091109 100644 --- a/earthkit/data/writers/grib.py +++ b/earthkit/data/writers/grib.py @@ -13,7 +13,7 @@ class GribWriter(Writer): DATA_FORMAT = "grib" - def write(self, f, values, metadata, check_nans=True, generating_proc_id=None): + def write(self, f, values, metadata, check_nans=True): r"""Write a GRIB field to a file object. Parameters @@ -26,10 +26,6 @@ def write(self, f, values, metadata, check_nans=True, generating_proc_id=None): Metadata of the GRIB field/message. check_nans: bool Replace nans in ``values`` with GRIB missing values when writing to``f``. - generating_proc_id: int - Set the ``generatingProcessIdentifier`` ecCodes GRIB key - to the specified value when writing to ``f``. When - ``None`` no changes are made. """ handle = metadata._handle if check_nans: @@ -41,7 +37,7 @@ def write(self, f, values, metadata, check_nans=True, generating_proc_id=None): handle.set_double("missingValue", missing_value) handle.set_long("bitmapPresent", 1) - handle.set_values(values, generating_proc_id=generating_proc_id) + handle.set_values(values) handle.write(f) diff --git a/tests/numpy_fs/test_numpy_fs_write.py b/tests/numpy_fs/test_numpy_fs_write.py index 4ac932e0..5dbefe91 100644 --- a/tests/numpy_fs/test_numpy_fs_write.py +++ b/tests/numpy_fs/test_numpy_fs_write.py @@ -113,16 +113,7 @@ def test_numpy_fs_grib_write_append(): assert r_tmp.metadata("shortName") == ["msl", "2d"] -@pytest.mark.parametrize( - "_kwargs,expected_value", - [ - ({}, 150), - ({"generating_proc_id": None}, 150), - ({"generating_proc_id": 255}, 255), - ({"generating_proc_id": 144}, 144), - ], -) -def test_numpy_fs_grib_write_generating_proc_id(_kwargs, expected_value): +def test_numpy_fs_grib_write_generating_proc_id(): ds = from_source("file", earthkit_examples_file("test.grib")) assert ds[0].metadata("shortName") == "2t" @@ -132,21 +123,21 @@ def test_numpy_fs_grib_write_generating_proc_id(_kwargs, expected_value): v2 = v + 2 md = ds[0].metadata() - md1 = md.override(shortName="msl") + md1 = md.override(shortName="msl", generatingProcessIdentifier=255) md2 = md.override(shortName="2d") r1 = FieldList.from_numpy([v1, v2], [md1, md2]) # save to disk: using generatingProcessIdentifier=255 (default) with temp_file() as tmp: - r1.save(tmp, **_kwargs) + r1.save(tmp) assert os.path.exists(tmp) r_tmp = from_source("file", tmp) assert len(r_tmp) == 2 assert r_tmp.metadata("shortName") == ["msl", "2d"] assert r_tmp.metadata("generatingProcessIdentifier") == [ - expected_value, - expected_value, + 255, + 150, ]