octavi parser

packages/evaluate/pyproject.toml

@@ -21,9 +21,14 @@ dependencies = [
     "earthkit-data==0.17.0",
     "earthkit-utils==0.1.2",
     "imageio[ffmpeg]>=2.37.2",
+    "cfgrib"
     "scipy>=1.12",
     ]
 
+[tool.uv.sources]
+cfgrib = { git = "https://github.com/ecmwf/cfgrib", rev = "322f38818bdd236a4fc4b70bf999212399bf79cf"}
+
+
 [dependency-groups]
 dev = [
  "pytest>=8.3",

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']"
         )

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,
     )
 

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")

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,41 @@ 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 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:
+        """
+        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}.")
+        ds = self.gribify(ds)
+        to_grib(ds, out_fname)
+        _logger.info(f"Saved GRIB file to {out_fname}.")