From 1e23ed8808c04aab5e117944af6f566bd4aa812a Mon Sep 17 00:00:00 2001 From: Sorcha Owens Date: Tue, 2 Jun 2026 17:29:27 +0200 Subject: [PATCH 1/2] WIP: cf-compliant gribs --- packages/evaluate/pyproject.toml | 5 ++++ .../weathergen/evaluate/export/cf_utils.py | 4 ++- .../evaluate/export/export_inference.py | 5 ++-- .../evaluate/export/parser_factory.py | 3 +- .../evaluate/export/parsers/netcdf_parser.py | 29 ++++++++++++++++++- 5 files changed, 41 insertions(+), 5 deletions(-) diff --git a/packages/evaluate/pyproject.toml b/packages/evaluate/pyproject.toml index a2d37fcae..0446a1479 100644 --- a/packages/evaluate/pyproject.toml +++ b/packages/evaluate/pyproject.toml @@ -21,8 +21,13 @@ dependencies = [ "earthkit-data==0.17.0", "earthkit-utils==0.1.2", "imageio[ffmpeg]>=2.37.2", + "cfgrib" ] +[tool.uv.sources] +cfgrib = { git = "https://github.com/ecmwf/cfgrib", rev = "322f38818bdd236a4fc4b70bf999212399bf79cf"} + + [dependency-groups] dev = [ "pytest>=8.3", diff --git a/packages/evaluate/src/weathergen/evaluate/export/cf_utils.py b/packages/evaluate/src/weathergen/evaluate/export/cf_utils.py index 3a9e7b2e5..5a5611c96 100644 --- a/packages/evaluate/src/weathergen/evaluate/export/cf_utils.py +++ b/packages/evaluate/src/weathergen/evaluate/export/cf_utils.py @@ -96,10 +96,12 @@ def _get_file_extension(output_format: str) -> str: return "nc" if output_format == "verif": return "nc" + if output_format == "grib": + return "grib" elif output_format == "quaver": return "grib" else: raise ValueError( f"Unsupported output format: {output_format}," - "supported formats are ['netcdf', 'verif', 'quaver']" + "supported formats are ['netcdf', 'verif', 'grib', 'quaver']" ) diff --git a/packages/evaluate/src/weathergen/evaluate/export/export_inference.py b/packages/evaluate/src/weathergen/evaluate/export/export_inference.py index 22a674a86..93503a8d3 100755 --- a/packages/evaluate/src/weathergen/evaluate/export/export_inference.py +++ b/packages/evaluate/src/weathergen/evaluate/export/export_inference.py @@ -95,9 +95,10 @@ def parse_args(args: list) -> argparse.Namespace: "--format", dest="output_format", type=str, - choices=["netcdf", "verif", "quaver"], + choices=["netcdf", "verif", "quaver", "grib"], help="Output file format; netcdf (CF-compliant netcdfs), \ - verif (netcdf compatible with MetNor verif tool), quaver (GRIB files for Quaver tool)", + verif (netcdf compatible with MetNor verif tool), quaver (GRIB files for Quaver tool), \ + grib (CF-compliant GRIB files)", required=True, ) diff --git a/packages/evaluate/src/weathergen/evaluate/export/parser_factory.py b/packages/evaluate/src/weathergen/evaluate/export/parser_factory.py index 1d3f7e26b..7d8592851 100644 --- a/packages/evaluate/src/weathergen/evaluate/export/parser_factory.py +++ b/packages/evaluate/src/weathergen/evaluate/export/parser_factory.py @@ -1,7 +1,7 @@ from omegaconf import OmegaConf from weathergen.evaluate.export.cf_utils import CfParser -from weathergen.evaluate.export.parsers.netcdf_parser import NetcdfParser +from weathergen.evaluate.export.parsers.netcdf_parser import NetcdfParser, GribParser from weathergen.evaluate.export.parsers.quaver_parser import QuaverParser from weathergen.evaluate.export.parsers.verif_parser import VerifParser @@ -32,6 +32,7 @@ def get_parser(config: OmegaConf, **kwargs) -> CfParser: "netcdf": (NetcdfParser, ["grid_type"]), "quaver": (QuaverParser, ["grid_type", "channels", "template"]), "verif": (VerifParser, ["obs", "method", "verif_template"]), + "grib": (GribParser, ["grid_type"]) } fmt = kwargs.get("output_format") diff --git a/packages/evaluate/src/weathergen/evaluate/export/parsers/netcdf_parser.py b/packages/evaluate/src/weathergen/evaluate/export/parsers/netcdf_parser.py index fc30b05f3..4f9edc3f9 100644 --- a/packages/evaluate/src/weathergen/evaluate/export/parsers/netcdf_parser.py +++ b/packages/evaluate/src/weathergen/evaluate/export/parsers/netcdf_parser.py @@ -7,6 +7,7 @@ import numpy as np import xarray as xr from omegaconf import OmegaConf +from cfgrib.xarray_to_grib import to_grib from weathergen.evaluate.export.cf_utils import CfParser from weathergen.evaluate.export.reshape import Regridder, find_pl, get_grid_points @@ -21,7 +22,8 @@ --output-dir ./test_output1 --format netcdf --samples 1 2 --fsteps 1 2 3 """ - +import cfgrib +print(cfgrib.__version__) class NetcdfParser(CfParser): """ @@ -550,3 +552,28 @@ def save(self, ds: xr.Dataset, forecast_ref_time: np.datetime64) -> None: _logger.info(f"Saving to {out_fname}.") ds.to_netcdf(out_fname) _logger.info(f"Saved NetCDF file to {out_fname}.") + + +class GribParser(NetcdfParser): + """ + Child class for handling GRIB output format. + """ + + def save(self, ds: xr.Dataset, forecast_ref_time: np.datetime64) -> None: + """ + Save the dataset to a GRIB file. + + Parameters + ---------- + ds : xarray Dataset to save. + data_type : Type of data ('pred' or 'targ') to include in the filename. + forecast_ref_time : Forecast reference time to include in the filename. + + Returns + ------- + None + """ + out_fname = self.get_output_filename(forecast_ref_time) + _logger.info(f"Saving to {out_fname}.") + to_grib(ds, out_fname) + _logger.info(f"Saved GRIB file to {out_fname}.") \ No newline at end of file From af4326cddab37f4a23ac120aee1890d0216122b2 Mon Sep 17 00:00:00 2001 From: Sorcha Owens Date: Tue, 9 Jun 2026 17:53:16 +0200 Subject: [PATCH 2/2] convertin gpressure_levels, still an issue w/grib names --- .../evaluate/export/parsers/netcdf_parser.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/packages/evaluate/src/weathergen/evaluate/export/parsers/netcdf_parser.py b/packages/evaluate/src/weathergen/evaluate/export/parsers/netcdf_parser.py index 4f9edc3f9..ec6cbc555 100644 --- a/packages/evaluate/src/weathergen/evaluate/export/parsers/netcdf_parser.py +++ b/packages/evaluate/src/weathergen/evaluate/export/parsers/netcdf_parser.py @@ -558,6 +558,18 @@ class GribParser(NetcdfParser): """ Child class for handling GRIB output format. """ + def gribify(self, ds: xr.Dataset) -> xr.Dataset: + """ + Convert dataset to use GRIB data variable names. + """ + # change pressure to isobaricInhPa for GRIB compliance + if "pressure" in ds.coords: + ds = ds.rename({"pressure": "isobaricInhPa"}) + if "valid_time" in ds.coords: + ds = ds.rename({"valid_time": "time"}) + if "forecast_period" in ds.coords: + ds = ds.rename({"forecast_period": "step"}) + return ds def save(self, ds: xr.Dataset, forecast_ref_time: np.datetime64) -> None: """ @@ -575,5 +587,6 @@ def save(self, ds: xr.Dataset, forecast_ref_time: np.datetime64) -> None: """ out_fname = self.get_output_filename(forecast_ref_time) _logger.info(f"Saving to {out_fname}.") + ds = self.gribify(ds) to_grib(ds, out_fname) _logger.info(f"Saved GRIB file to {out_fname}.") \ No newline at end of file