Fast_aggregation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fast NODD GRIB Aggregations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "In this tutorial we are going to demonstrate building fast kerchunk aggregations of NODD grib2 weather forecasts. This workflow primarily involves [xarray-datatree](https://xarray-datatree.readthedocs.io/en/latest/), [pandas](https://pandas.pydata.org/) and the new `grib_tree` function released in kerchunk v0.2.3.\n",
    "\n",
    "\n",
    "### About the Dataset\n",
    "\n",
    "In this demo we will be looking at GRIB2 files generated by [**NOAA Global Ensemble Forecast System (GEFS)**](https://www.ncei.noaa.gov/products/weather-climate-models/global-ensemble-forecast). This dataset is a weather forecast model made up of 21 separate forecasts, or ensemble members. GEFS has global coverage and is produced four times a day with forecasts going out to 16 days. It is updated 4 times a day, every 6 hours starting at midnight.\n",
    "\n",
    "More information on this dataset can be found [here](https://registry.opendata.aws/noaa-gefs)\n",
    "\n",
    "\n",
    "## Prerequisites\n",
    "| Concepts | Importance | Notes |\n",
    "| --- | --- | --- |\n",
    "| [Kerchunk Basics](../foundations/kerchunk_basics) | Required | Core |\n",
    "| [Pandas Tutorial](https://foundations.projectpythia.org/core/pandas/pandas.html#) | Required | Core |\n",
    "| [Kerchunk and Xarray-Datatree](https://projectpythia.org/kerchunk-cookbook/notebooks/using_references/Datatree.html) | Required | IO |\n",
    "| [Xarray-Datatree Overview](https://xarray-datatree.readthedocs.io/en/latest/quick-overview.html)| Required | IO |\n",
    "\n",
    "- **Time to learn**: 30 minutes\n",
    "\n",
    "## Motivation\n",
    "\n",
    "As we know that **kerchunk**  provides a unified way to represent a variety of chunked, compressed data formats (e.g. NetCDF/HDF5, GRIB2, TIFF, …) by generating *references*. This task flow has ability to build large aggregations from **NODD grib forecasts**\n",
    "in a fraction of the time using the `idx files`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import logging\n",
    "import importlib\n",
    "importlib.reload(logging)\n",
    "logging.basicConfig(\n",
    "    format=\"%(asctime)s.%(msecs)03dZ %(processName)s %(threadName)s %(levelname)s:%(name)s:%(message)s\",\n",
    "    datefmt=\"%Y-%m-%dT%H:%M:%S\",\n",
    "    level=logging.WARNING,\n",
    ")\n",
    "\n",
    "logger = logging.getLogger(\"juypter\")\n",
    "\n",
    "\n",
    "import copy\n",
    "import fsspec\n",
    "import pandas as pd\n",
    "import xarray as xr\n",
    "import datetime\n",
    "from kerchunk.grib2 import (\n",
    "    AggregationType,\n",
    "    build_idx_grib_mapping,\n",
    "    extract_datatree_chunk_index,\n",
    "    grib_tree,\n",
    "    map_from_index,\n",
    "    parse_grib_idx,\n",
    "    reinflate_grib_store,\n",
    "    scan_grib,\n",
    "    strip_datavar_chunks,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building the Aggregation directly from the GRIB files\n",
    "\n",
    "We are using the newly implemented Kerchunk `grib_tree` function to build a hierarchical data model from a set of scanned grib messages.This data model can be opened directly using either zarr or xarray datatree. *This way of building the aggregation is very slow*. Here we're going to use xarray-datatree to open and view it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s3_files = [\n",
    "    \"s3://noaa-gefs-pds/gefs.20170101/06/gec00.t06z.pgrb2af006\",\n",
    "    \"s3://noaa-gefs-pds/gefs.20170101/06/gec00.t06z.pgrb2af012\",\n",
    "    \"s3://noaa-gefs-pds/gefs.20170101/06/gec00.t06z.pgrb2af018\",\n",
    "    \"s3://noaa-gefs-pds/gefs.20170101/06/gec00.t06z.pgrb2af024\",\n",
    "    \"s3://noaa-gefs-pds/gefs.20170101/06/gec00.t06z.pgrb2af030\",\n",
    "    \"s3://noaa-gefs-pds/gefs.20170101/06/gec00.t06z.pgrb2af036\",\n",
    "    \"s3://noaa-gefs-pds/gefs.20170101/06/gec00.t06z.pgrb2af042\",\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# converting the references into the hierarchical datamodel\n",
    "grib_tree_store = grib_tree(\n",
    "    [\n",
    "        group\n",
    "        for f in s3_files\n",
    "        for group in scan_grib(f, storage_options=dict(anon=True))\n",
    "    ],\n",
    "    remote_options=dict(anon=True),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Transforming the output to datatree to view it. This tree model the variables\n",
    "s3_dt = xr.open_datatree(\n",
    "    fsspec.filesystem(\n",
    "        \"reference\",\n",
    "        fo=grib_tree_store,\n",
    "        remote_protocol=\"s3\",\n",
    "        remote_options={\"anon\": True},\n",
    "    ).get_mapper(\"\"),\n",
    "    engine=\"zarr\",\n",
    "    consolidated=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# In this tree model, the variables are organized into hierarchical groups, first by \"stepType\" and then by \"typeOfLevel.\"\n",
    "s3_dt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Note**: If trying out this notebook, the above way of building the aggregation will take few minutes for this `GEFS` S3 repository. But in general, it take more time, based on the size of the grib files"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building the aggregation faster with `idx` files\n",
    "\n",
    "This method of aggregating GRIB files requires that each GRIB file be accompanied by its corresponding `idx` file in text format. This file, also called index files, contain some key metadata of the GRIB messages across the files. These metadata is then used to build an reference index, which helps in aggregation. \n",
    "\n",
    "The advantage is, building aggregation this way doesn't involve scanning the whole archive of GRIB files but only a single file. \n",
    "\n",
    "### Index Dataframe made from a single Grib file\n",
    "\n",
    "Here is what the contents of an `idx` file looks like.\n",
    "\n",
    "```\n",
    "1:0:d=2017010100:HGT:10 mb:12 hour fcst:ENS=low-res ctl\n",
    "2:48163:d=2017010100:TMP:10 mb:12 hour fcst:ENS=low-res ctl\n",
    "3:68112:d=2017010100:RH:10 mb:12 hour fcst:ENS=low-res ctl\n",
    "4:79092:d=2017010100:UGRD:10 mb:12 hour fcst:ENS=low-res ctl\n",
    "5:102125:d=2017010100:VGRD:10 mb:12 hour fcst:ENS=low-res ctl\n",
    "6:122799:d=2017010100:HGT:50 mb:12 hour fcst:ENS=low-res ctl\n",
    "7:178898:d=2017010100:TMP:50 mb:12 hour fcst:ENS=low-res ctl\n",
    "8:201799:d=2017010100:RH:50 mb:12 hour fcst:ENS=low-res ctl\n",
    "9:224321:d=2017010100:UGRD:50 mb:12 hour fcst:ENS=low-res ctl\n",
    "10:272234:d=2017010100:VGRD:50 mb:12 hour fcst:ENS=low-res ctl\n",
    "11:318288:d=2017010100:HGT:100 mb:12 hour fcst:ENS=low-res ctl\n",
    "12:379010:d=2017010100:TMP:100 mb:12 hour fcst:ENS=low-res ctl\n",
    "13:405537:d=2017010100:RH:100 mb:12 hour fcst:ENS=low-res ctl\n",
    "14:441517:d=2017010100:UGRD:100 mb:12 hour fcst:ENS=low-res ctl\n",
    "15:497421:d=2017010100:VGRD:100 mb:12 hour fcst:ENS=low-res ctl\n",
    "```\n",
    "\n",
    "The general format of `idx` data is: **index:offset:date_with_runtime:variable:forecast_time:**.\n",
    "\n",
    "The metadata are separated by \":\" (colon) and we need to convert it into a `Dataframe` for building the index. \n",
    "\n",
    "\n",
    "> **Note**: This way of building the reference index only works for a particular **horizon** file irrespective of the run time of the model, as the  messages from a same time horizon have an identical structure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# converting the idx data into a dataframe\n",
    "idxdf = parse_grib_idx(\n",
    "    #\"s3://noaa-gefs-pds/gefs.20170101/06/gec00.t06z.pgrb2af006\",\n",
    "    \"s3://noaa-gefs-pds/gefs.20230101/00/atmos/pgrb2sp25/geavg.t00z.pgrb2s.0p25.f006\",\n",
    "    storage_options=dict(anon=True),\n",
    "    \n",
    ")\n",
    "idxdf.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Building a mapping between the index data and grib metadata\n",
    "\n",
    "Creation of the reference index or **k_index** first involves generating a single mapping from the grib/zarr metadata(the `grib_tree` output) to the attributes in the idx files. They are unique for each time horizon, so while building the index, we need to build a unique mapping for the 1 hour forecast, the 2 hour forecast and so on. \n",
    "\n",
    "This mapping can be made from any arbitrary GRIB file, given that the it belongs to the **same time horizon** of files we're trying to index. \n",
    "\n",
    "We'll start by examining the GRIB data. The metadata that we'll be extracting will be static in nature. We're going to use a single node by [accessing](https://projectpythia.org/kerchunk-cookbook/notebooks/using_references/Datatree.html#accessing-the-datatree) it with datatree."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s3_dt.groups"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Parsing the grib metadata from a single datatree node and converting it into a dataframe\n",
    "grib_df = extract_datatree_chunk_index(s3_dt[\"sulwrf/avg/nominalTop\"], grib_tree_store)\n",
    "grib_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Note**: Above process is part of the mapping creation, the function call to `extract_datatree_chunk_index` in handled inside `build_idx_grib_mapping` function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# creating a mapping for a single horizon file which is to be used later\n",
    "mapping = build_idx_grib_mapping(\n",
    "    #\"s3://noaa-gefs-pds/gefs.20170101/00/geavg00.t00z.pgrb2f006\",\n",
    "    \"s3://noaa-gefs-pds/gefs.20230101/00/atmos/pgrb2sp25/geavg.t00z.pgrb2s.0p25.f006\",\n",
    "    storage_options=dict(anon=True),\n",
    "    validate=True,\n",
    ")\n",
    "mapping.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ">**Note**: Reference index can be combined across many horizons but **each horizon must be indexed separately**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Building the index \n",
    "\n",
    "Now if we parse the runtime from the idx file and other attributes, we can build a fully compatible kerchunk index(k_index) for that \n",
    "**particular file**. Before creating the index, we need to clean some of the data in the mapping and index dataframe for the some variables as they tend to contain duplicate values, as demonstrated below. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# cleaning the mapping\n",
    "mapping.loc[mapping[\"attrs\"].duplicated(keep=False), :]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# this step will be performed for every idx dataframe where we will be using the \"mapping\" dataframe which we created previously\n",
    "mapped_index = map_from_index(\n",
    "    pd.Timestamp(\"2017-01-01T06\"),\n",
    "    mapping.loc[~mapping[\"attrs\"].duplicated(keep=\"first\"), :],\n",
    "    idxdf.loc[~idxdf[\"attrs\"].duplicated(keep=\"first\"), :],\n",
    ")\n",
    "mapped_index.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Final step of building of Aggregation\n",
    "\n",
    "For the final step of the aggregation, we will create an index for each GRIB file to cover a two-month period. We will be using the 6-hour horizon file for building the aggregation, this will be from `2017-01-01` to `2017-02-28`. The generated index can then be combined with indexes of other time horizons and stored for later use. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "mapped_index_list = []\n",
    "\n",
    "deduped_mapping = mapping.loc[~mapping[\"attrs\"].duplicated(keep=\"first\"), :]\n",
    "\n",
    "# this process is manually done because the runtime and forecast time will vary between models i.e. GEFS, GFS, HRRR etc.\n",
    "for date in pd.date_range(\"2023-01-01\", \"2023-02-28\"):\n",
    "    for runtime in range(0, 24, 6):\n",
    "        #fname = f\"s3://noaa-gefs-pds/gefs.{date.strftime('%Y%m%d')}/{runtime:02}/gec00.t{runtime:02}z.pgrb2af006\"\n",
    "        fname = f\"s3://noaa-gefs-pds/gefs.{date.strftime('%Y%m%d')}/{runtime:02}/atmos/pgrb2sp25/geavg.t{runtime:02}z.pgrb2s.0p25.f006\"\n",
    "\n",
    "        idxdf = parse_grib_idx(basename=fname, storage_options=dict(anon=True))\n",
    "    \n",
    "        mapped_index = map_from_index(\n",
    "            pd.Timestamp(date + datetime.timedelta(hours=runtime)),\n",
    "            deduped_mapping,\n",
    "            idxdf.loc[~idxdf[\"attrs\"].duplicated(keep=\"first\"), :],\n",
    "        )\n",
    "    \n",
    "        mapped_index_list.append(mapped_index)\n",
    "\n",
    "s3_kind = pd.concat(mapped_index_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Tip**: To confirm the above step, check out this [notebook](https://gist.github.com/Anu-Ra-g/efa01ad1c274c1bd1c14ee01666caa77)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using the aggregation\n",
    "\n",
    "The difference between .idx and Kerchunk index that we built is that the former indexes the GRIB messages and the latter indexes the variables in those messages, in a time horizon. Now we'll need a tree model from grib_tree function to reinflate index (variables in the messages). The important point to note here is that the tree model should be made from the GRIB files of the repository that we are indexing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grib_tree_store = grib_tree(\n",
    "    scan_grib(\n",
    "        #\"s3://noaa-gefs-pds/gefs.20170101/06/gec00.t06z.pgrb2af006\",\n",
    "        \"s3://noaa-gefs-pds/gefs.20230101/00/atmos/pgrb2sp25/geavg.t00z.pgrb2s.0p25.f006\",\n",
    "        storage_options=dict(anon=True),\n",
    "    ),\n",
    "    remote_options=dict(anon=True),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we need to run that tree model through `strip_datavar_chunks` function, which will strip the tree of the variables *in place*. This step is optional"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "deflated_s3_grib_tree_store = copy.deepcopy(\n",
    "    grib_tree_store\n",
    ")  # not affecting the original tree model\n",
    "strip_datavar_chunks(deflated_s3_grib_tree_store)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After stripping the tree model, we need to introduce two new axes: the `runtime` used for mapping and `date_range` for indexing the new reinflated tree. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "axes = [\n",
    "    pd.Index(\n",
    "        [\n",
    "            pd.timedelta_range(\n",
    "                start=\"0 hours\", end=\"6 hours\", freq=\"6h\", closed=\"right\", name=\"6 hour\"\n",
    "            ),\n",
    "        ],\n",
    "        name=\"step\",\n",
    "    ),\n",
    "    pd.date_range(\n",
    "        \"2023-01-01T00:00\", \"2023-02-28T18:00\", freq=\"360min\", name=\"valid_time\"\n",
    "    ),\n",
    "]\n",
    "axes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s3_kind.loc[s3_kind.varname.isin([\"sdswrf\"])].head(20)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now for reinflating, we'll be needing the aggregation types which are: `horizon`, `valid_time`, `run_time` and `best_available`. We will also be needing the variable(s) that we are reinflating. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s3_store = reinflate_grib_store(\n",
    "    axes=axes,\n",
    "    aggregation_type=AggregationType.HORIZON,\n",
    "    chunk_index=s3_kind.loc[s3_kind.varname.isin([\"sulwrf\", \"prmsl\", \"sdswrf\"])],\n",
    "    zarr_ref_store=deflated_s3_grib_tree_store,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Viewing the new subset of the datatree\n",
    "\n",
    "In this step, we can view the new subset as a `datatree` model. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s3_dt_subset = xr.open_datatree(\n",
    "    fsspec.filesystem(\n",
    "        \"reference\", fo=s3_store, remote_protocol=\"s3\", remote_options={\"anon\": True}\n",
    "    ).get_mapper(\"\"),\n",
    "    engine=\"zarr\",\n",
    "    consolidated=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s3_dt_subset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s3_dt_subset.groups"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s3_dt_subset.sdswrf.avg.surface.sdswrf[0,3,:,:].plot(figsize=(12,8))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "da = s3_dt_subset.sulwrf.avg.nominalTop.sulwrf.isel(model_horizons=0)\n",
    "da"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Create an index for time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "da= da.rename({'valid_times':'time'}).drop_vars(['step'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "da = da.set_index(time='time')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "da"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "da.sel(time='2023-01-02', method='nearest')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Try extracting some data.  \n",
    "The results from the cell below show lot's of NaN and what look like integer values. This doesn't seem like what would be expected for \"Surface upward long-wave radiation flux\", right?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "da[:,90,180].load()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}