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

@@ -22,5 +22,6 @@ dependencies:
   - 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:
     # add pip requirements here if needed.

setup_env.sh

@@ -106,3 +106,14 @@ else
         fi
     fi
 fi
+
+
+# setting ECCODES_DEFINITION_PATH:
+git clone --recursive https://github.com/COSMO-ORG/eccodes-cosmo-resources.git
+
+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}

templates/ICON.jinja

@@ -10,6 +10,7 @@
             "member_num": "[{member_num}]",
             "member_type": "{{member_type}}",
             "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

@@ -4,6 +4,7 @@
 import numpy as np
 import pandas as pd
 import xarray
+import earthkit.data
 
 from util.constants import compute_statistics
 from util.log_handler import logger
@@ -60,6 +61,80 @@ def parse_netcdf(file_id, filename, specification):
     ds.close()
     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.where(~np.isin(short_name, short_name_excl))].tolist()
+
+    level_type = np.unique(ds_grib.metadata("typeOfLevel")).tolist()
+    level_type.remove("unknown")
+
+
+    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 = get_ds(ds_grib, pid, lev)
+            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):
+    try:
+        ds = ds_grib.sel(paramId=pid, typeOfLevel=lev).to_xarray()
+    except:
+        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()
+            except:
+                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()
+                    except:
+                        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()
+                            except:
+                                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()
+                                    except:
+                                        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()
+                                            except:
+                                                print("not working!")
+
+    return ds
+
 
 def __get_variables(data, time_dim, horizontal_dims):
     # return a list of variable names from the dataset data that have a time dimension
@@ -228,4 +303,5 @@ def parse_csv(file_id, filename, specification):
 model_output_parser = {  # global lookup dict
     "netcdf": parse_netcdf,
     "csv": parse_csv,
+    "grib": parse_grib,
 }