Add reading grib files for stats

requirements/dev-environment.yml

@@ -255,5 +255,6 @@ dependencies:
   - zipp=3.11.0
   - zlib=1.2.13
   - zstd=1.5.2
+  - earthkit-data
   - pip:
       - flake8-pyproject==1.2.2

requirements/dev-requirements.yml

@@ -21,6 +21,6 @@ dependencies:
   - sphinx
   - sphinx-autobuild
   - toml>=0.10.2
-  - pip
+  - earthkit-data
   - pip:
     - flake8-pyproject

requirements/environment.yml

@@ -175,3 +175,4 @@ dependencies:
   - xz=5.2.6
   - zlib=1.2.13
   - zstd=1.5.2
+  - earthkit-data

requirements/requirements.yml

@@ -15,5 +15,6 @@ dependencies:
   - scipy>=1.4.1
   - xarray>=0.16.1
   - pip
+  - earthkit-data
   - pip:
     - pytest-cov

setup_env.sh

@@ -106,3 +106,16 @@ else
         fi
     fi
 fi
+
+
+# setting ECCODES_DEFINITION_PATH:
+git clone --recursive https://github.com/COSMO-ORG/eccodes-cosmo-resources.git
+
+${CONDA} activate ${DEV_ENV_NAME}
+conda_loc=${CONDA_PREFIX}
+base_dir=$(pwd)
+def_path_default=${conda_loc}/share/eccodes/definitions
+def_path_resources=${base_dir}/eccodes-cosmo-resources/definitions
+
+conda env config vars set ECCODES_DEFINITION_PATH=${def_path_default}:${def_path_resources}
+${CONDA} deactivate

templates/ICON.jinja

@@ -11,6 +11,7 @@
             "member_type": "{{member_type}}",
             "factor": 5,
             "file_specification": [{
+                "GRIB": {"format": "GRIB", "time_dim": "step", "horizontal_dims": ["values"], "var_excl": ["tlon", "tlat", "vlon", "vlat", "ulon", "ulat", "h", "slor", "anor", "isor", "sdor"]},
                 "latlon": { "format": "netcdf", "time_dim": "time", "horizontal_dims": ["lat:lon"] },
                 "meteogram": { "format": "netcdf", "time_dim": "time", "horizontal_dims": ["max_nlevs:nstations", "nstations"] },
                 "dace":{ "format": "netcdf", "time_dim": null, "horizontal_dims": ["d_body"]},

util/model_output_parser.py

@@ -1,6 +1,7 @@
 import sys
 from collections.abc import Iterable
 
+import earthkit.data
 import numpy as np
 import pandas as pd
 import xarray
@@ -61,6 +62,155 @@ def parse_netcdf(file_id, filename, specification):
     return var_dfs
 
 
+def parse_grib(file_id, filename, specification):
+    logger.debug("parse GRIB file {}".format(filename))
+    time_dim = specification["time_dim"]
+    horizontal_dims = specification["horizontal_dims"]
+    fill_value_key = specification.get("fill_value_key", None)
+
+    ds_grib = earthkit.data.from_source("file", filename)
+    short_name_excl = specification["var_excl"]
+
+    short_name = np.unique(ds_grib.metadata("shortName"))
+    short_name = short_name[
+        np.isin(short_name, short_name_excl, invert=True, assume_unique=True)
+    ].tolist()
+
+    level_type = np.unique(ds_grib.metadata("typeOfLevel")).tolist()
+
+    var_dfs = []
+    for lev in level_type:
+        paramId = np.unique(
+            ds_grib.sel(typeOfLevel=lev, shortName=short_name).metadata("paramId")
+        ).tolist()
+        for pid in paramId:
+            ds_temp_list = get_ds(ds_grib, pid, lev)
+            for ds_temp in ds_temp_list:
+                v = list(ds_temp.keys())[0]
+
+                dim_to_squeeze = [
+                    dim
+                    for dim, size in zip(ds_temp[v].dims, ds_temp[v].shape)
+                    if size == 1 and dim != time_dim
+                ]
+                ds = ds_temp.squeeze(dim=dim_to_squeeze)
+
+                sub_df = dataframe_from_ncfile(
+                    file_id=file_id,
+                    filename=filename,
+                    varname=v,
+                    time_dim=time_dim,
+                    horizontal_dims=horizontal_dims,
+                    xarray_ds=ds,
+                    fill_value_key=fill_value_key,
+                )
+                var_dfs.append(sub_df)
+
+    return var_dfs
+
+
+def get_ds(ds_grib, pid, lev):
+    ds_list = []
+    try:
+        ds = ds_grib.sel(paramId=pid, typeOfLevel=lev).to_xarray()
+        ds_list.append(ds)
+    except KeyError:
+        stepType = np.unique(
+            ds_grib.sel(paramId=pid, typeOfLevel=lev).metadata("stepType")
+        ).tolist()
+        for steps in stepType:
+            try:
+                ds = ds_grib.sel(
+                    paramId=pid, typeOfLevel=lev, stepType=steps
+                ).to_xarray()
+                ds_list.append(ds)
+            except KeyError:
+                num_points = np.unique(
+                    ds_grib.sel(paramId=pid, typeOfLevel=lev, stepType=steps).metadata(
+                        "numberOfPoints"
+                    )
+                ).tolist()
+                for points in num_points:
+                    try:
+                        ds = ds_grib.sel(
+                            paramId=pid,
+                            typeOfLevel=lev,
+                            stepType=steps,
+                            numberOfPoints=points,
+                        ).to_xarray()
+                        ds_list.append(ds)
+                    except KeyError:
+                        units = np.unique(
+                            ds_grib.sel(
+                                paramId=pid,
+                                typeOfLevel=lev,
+                                stepType=steps,
+                                numberOfPoints=points,
+                            ).metadata("stepUnits")
+                        ).tolist()
+                        for unit in units:
+                            try:
+                                ds = ds_grib.sel(
+                                    paramId=pid,
+                                    typeOfLevel=lev,
+                                    stepType=steps,
+                                    numberOfPoints=points,
+                                    stepUnits=unit,
+                                ).to_xarray()
+                                ds_list.append(ds)
+                            except KeyError:
+                                dataType = np.unique(
+                                    ds_grib.sel(
+                                        paramId=pid,
+                                        typeOfLevel=lev,
+                                        stepType=steps,
+                                        numberOfPoints=points,
+                                        stepUnits=unit,
+                                    ).metadata("dataType")
+                                ).tolist()
+                                for dtype in dataType:
+                                    try:
+                                        ds = ds_grib.sel(
+                                            paramId=pid,
+                                            typeOfLevel=lev,
+                                            stepType=steps,
+                                            numberOfPoints=points,
+                                            stepUnits=unit,
+                                            dataType=dtype,
+                                        ).to_xarray()
+                                        ds_list.append(ds)
+                                    except KeyError:
+                                        gridType = np.unique(
+                                            ds_grib.sel(
+                                                paramId=pid,
+                                                typeOfLevel=lev,
+                                                stepType=steps,
+                                                numberOfPoints=points,
+                                                stepUnits=unit,
+                                                dataType=dtype,
+                                            ).metadata("gridType")
+                                        ).tolist()
+                                        for gtype in gridType:
+                                            try:
+                                                ds = ds_grib.sel(
+                                                    paramId=pid,
+                                                    typeOfLevel=lev,
+                                                    stepType=steps,
+                                                    numberOfPoints=points,
+                                                    stepUnits=unit,
+                                                    dataType=dtype,
+                                                    gridType=gtype,
+                                                ).to_xarray()
+                                                ds_list.append(ds)
+                                            except KeyError:
+                                                print(
+                                                    f"GRIB file of level {lev} and"
+                                                    "paramId {pid} cannot be read."
+                                                )
+
+    return ds_list
+
+
 def __get_variables(data, time_dim, horizontal_dims):
     # return a list of variable names from the dataset data that have a time dimension
     # and horizontal dimension or in case there is no time dimension just the variables
@@ -228,4 +378,5 @@ def parse_csv(file_id, filename, specification):
 model_output_parser = {  # global lookup dict
     "netcdf": parse_netcdf,
     "csv": parse_csv,
+    "grib": parse_grib,
 }