Jupyter notebook demonstrating how to create custom GRIB definitions to use with eccodes

Custom ECCODES GRIB definitions.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6b03ff89-643d-4afe-a220-13deecc47e7c",
   "metadata": {},
   "source": [
    "## Custom GRIB Definitions for `eccodes`\n",
    "\n",
    "`eccodes` relies on a variety of definition files for name, shortname, and units insides a very specific directory structure to\n",
    "decode custom fields. This notebook provides a helper context manager class to ease this process and uses it to open a GRIB1\n",
    "files for some [NARR data](https://thredds.rda.ucar.edu/thredds/catalog/files/g/d608000/3HRLY/catalog.html) using xarray and the\n",
    "cfgrib engine."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e8435550-8afd-47a2-9ac2-d8fec9a88b6e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from pathlib import Path\n",
    "import tempfile\n",
    "\n",
    "class CustomGribDef:\n",
    "    '''Manage custom GRIB definitions for loading data with eccodes.\n",
    "\n",
    "    Sets up a temporary directory to write custom defintion files and sets the\n",
    "    appropriate environment variable for eccodes to use them when opening a GRIB\n",
    "    file.\n",
    "\n",
    "    Uses a Python context manager to ensure consistent setup and teardown.\n",
    "    '''\n",
    "    def __init__(self, edition, center, version, parameters):\n",
    "        '''Initialize Manager\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        edition: int\n",
    "            The GRIB edition needed, 1 or 2.\n",
    "        center: str or int\n",
    "            Identifier for the center for the parameters. Needs to be a string (e.g. kwbc)\n",
    "            if known to eccodes.\n",
    "        version: int\n",
    "            Table version for the paramters.\n",
    "        params: dict\n",
    "            Mapping of parameter id to a tuple of 3 strings (name, short name, unit). For GRIB2, the\n",
    "            parameter ID should be a tuple of (category, parameter_num).\n",
    "        '''\n",
    "        self.edition = edition\n",
    "        self.center = center\n",
    "        self.version = version\n",
    "        self.params = parameters\n",
    "\n",
    "    def _write_param(self, direc, pid, name, sname, units):\n",
    "        '''Handle writing JSON blocks to the various files for a parameter.'''\n",
    "        if self.edition == 1:\n",
    "            block = f' = {{\\n    table2Version = {self.version} ;\\n    indicatorOfParameter = {pid} ;\\n}}\\n'\n",
    "        elif self.edition == 2:\n",
    "            paramCat, paramNum = pid\n",
    "            block = f' = {{\\n    discipline = {discipline} ;\\n    parameterCategory = {paramCat} ;\\n    parameterNumber = {paramNum} ;\\n}}\\n'\n",
    "        with (direc / 'name.def').open('at') as out:\n",
    "            out.write(f\"'{name}'\" + block)\n",
    "        with (direc / 'shortName.def').open('at') as out:\n",
    "            out.write(f\"'{sname}'\" + block)\n",
    "        with (direc / 'units.def').open('at') as out:\n",
    "            out.write(f\"'{units}'\" + block)\n",
    "        with (direc / 'paramId.def').open('at') as out:\n",
    "            paramid = f'{self.version:03d}{pid:03d}' if self.edition == 1 else f'{paramCat:03d}{paramNum:03d}'\n",
    "            out.write(f\"'{paramid}'\" + block)\n",
    "\n",
    "    def __enter__(self):\n",
    "        '''Set up the information for the custom definitions.'''\n",
    "        self.temp_dir = tempfile.TemporaryDirectory()\n",
    "        self._old_env = os.environ.get('ECCODES_DEFINITION_PATH', '')\n",
    "        os.environ['ECCODES_DEFINITION_PATH'] = os.pathsep.join((str(self.temp_dir), self._old_env))\n",
    "\n",
    "        center_dir = Path(str(self.temp_dir)) / f'grib{self.edition}' / 'localConcepts' / str(self.center)\n",
    "        center_dir.mkdir(parents=True)\n",
    "        for pid, info in self.params.items():\n",
    "            self._write_param(center_dir, pid, *info)\n",
    "\n",
    "    def __exit__(self, exc_type, exc_value, traceback):\n",
    "        '''Clean up when done.'''\n",
    "        os.environ['ECCODES_DEFINITION_PATH'] = self._old_env\n",
    "        del self.temp_dir"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ff1520b5-5a71-4584-9a57-484d591f7a7c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "skipping variable: paramId==131158 shortName='tke'\n",
      "Traceback (most recent call last):\n",
      "  File \"/Users/rmay/miniconda3/envs/py313/lib/python3.13/site-packages/cfgrib/dataset.py\", line 721, in build_dataset_components\n",
      "    dict_merge(variables, coord_vars)\n",
      "    ~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/Users/rmay/miniconda3/envs/py313/lib/python3.13/site-packages/cfgrib/dataset.py\", line 639, in dict_merge\n",
      "    raise DatasetBuildError(\n",
      "    ...<2 lines>...\n",
      "    )\n",
      "cfgrib.dataset.DatasetBuildError: key present and new value is different: key='isobaricInhPa' value=Variable(dimensions=('isobaricInhPa',), data=array([1000.,  975.,  950.,  925.,  900.,  875.,  850.,  825.,  800.,\n",
      "        775.,  750.,  725.,  700.,  650.,  600.,  550.,  500.,  450.,\n",
      "        400.,  350.,  300.,  275.,  250.,  225.,  200.,  175.,  150.,\n",
      "        125.,  100.])) new_value=Variable(dimensions=('isobaricInhPa',), data=array([1000.,  975.,  950.,  925.,  900.,  875.,  850.,  825.,  800.,\n",
      "        775.,  750.,  725.,  700.,  650.,  600.]))\n",
      "skipping variable: paramId==131135 shortName='mconv'\n",
      "Traceback (most recent call last):\n",
      "  File \"/Users/rmay/miniconda3/envs/py313/lib/python3.13/site-packages/cfgrib/dataset.py\", line 721, in build_dataset_components\n",
      "    dict_merge(variables, coord_vars)\n",
      "    ~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/Users/rmay/miniconda3/envs/py313/lib/python3.13/site-packages/cfgrib/dataset.py\", line 639, in dict_merge\n",
      "    raise DatasetBuildError(\n",
      "    ...<2 lines>...\n",
      "    )\n",
      "cfgrib.dataset.DatasetBuildError: key present and new value is different: key='isobaricInhPa' value=Variable(dimensions=('isobaricInhPa',), data=array([1000.,  975.,  950.,  925.,  900.,  875.,  850.,  825.,  800.,\n",
      "        775.,  750.,  725.,  700.,  650.,  600.,  550.,  500.,  450.,\n",
      "        400.,  350.,  300.,  275.,  250.,  225.,  200.,  175.,  150.,\n",
      "        125.,  100.])) new_value=Variable(dimensions=(), data=np.float64(850.0))\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<xarray.Dataset> Size: 91MB\n",
      "Dimensions:        (isobaricInhPa: 29, y: 277, x: 349)\n",
      "Coordinates:\n",
      "    time           datetime64[ns] 8B ...\n",
      "    step           timedelta64[ns] 8B ...\n",
      "  * isobaricInhPa  (isobaricInhPa) float64 232B 1e+03 975.0 ... 125.0 100.0\n",
      "    latitude       (y, x) float64 773kB ...\n",
      "    longitude      (y, x) float64 773kB ...\n",
      "    valid_time     datetime64[ns] 8B ...\n",
      "Dimensions without coordinates: y, x\n",
      "Data variables:\n",
      "    hgt            (isobaricInhPa, y, x) float32 11MB ...\n",
      "    tmp            (isobaricInhPa, y, x) float32 11MB ...\n",
      "    spfh           (isobaricInhPa, y, x) float32 11MB ...\n",
      "    vvel           (isobaricInhPa, y, x) float32 11MB ...\n",
      "    ugrd           (isobaricInhPa, y, x) float32 11MB ...\n",
      "    vgrd           (isobaricInhPa, y, x) float32 11MB ...\n",
      "    clwmr          (isobaricInhPa, y, x) float32 11MB ...\n",
      "    icmr           (isobaricInhPa, y, x) float32 11MB ...\n",
      "Attributes:\n",
      "    GRIB_edition:            1\n",
      "    GRIB_centre:             kwbc\n",
      "    GRIB_centreDescription:  US National Weather Service - NCEP\n",
      "    GRIB_subCentre:          15\n",
      "    Conventions:             CF-1.7\n",
      "    institution:             US National Weather Service - NCEP\n",
      "    history:                 2024-12-30T15:16 GRIB to CDM+CF via cfgrib-0.9.1...\n"
     ]
    }
   ],
   "source": [
    "import eccodes\n",
    "import xarray as xr\n",
    "\n",
    "params = {7: ('Geopotential Height', 'hgt', 'gpm'),\n",
    "          11: ('Temperature', 'tmp', 'K'),\n",
    "          33: ('u-component of wind', 'ugrd', 'm/s'),\n",
    "          34: ('v-component of wind', 'vgrd', 'm/s'),\n",
    "          39: ('Pressure Vertical Velocity', 'vvel', 'Pa/s'),\n",
    "          51: ('Specific Humidity', 'spfh', 'kg/kg'),\n",
    "          135: ('Horizontal Moisture Divergence', 'mconv', 'kg/kg/s'),\n",
    "          153: ('Cloud water', 'clwmr', 'kg/kg'),\n",
    "          158: ('Turbulent Kinetic Energy', 'tke', 'j/kg'),\n",
    "          178: ('Ice Mixing Ratio', 'icmr', 'kg/kg')}\n",
    "\n",
    "with CustomGribDef(1, 'kwbc', 131, params):\n",
    "    # Open without using an index to ensure we are using our custom tables. If an index\n",
    "    # exists without our custom parameters, it will prevent our custom definitions from being used.\n",
    "    ds1 = xr.open_dataset('merged_AWIP32.2024041900.3D', engine='cfgrib', indexpath='')\n",
    "\n",
    "print(ds1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a18b871-fb83-4539-afc3-8e9c4cafacbe",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:py313]",
   "language": "python",
   "name": "conda-env-py313-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}