dynamicgribchunking.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "n8nQijCi9wVu"
   },
   "source": [
    "# GRIB2 NWP Archive aggretation using .idx files\n",
    "\n",
    "Work with NODD GFS data on GCP (theoretically works with AWS as well)\n",
    "\n",
    "1) Demonstrate building a DataTree, Index File mapping and Static Metdata data file. Use the mapping to rapidly build a on month dataset.\n",
    "2) Demonstrate using zarr chunk_getitems hack to parallelize getting a timeseries\n",
    "   "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "326f3vGe_HHb"
   },
   "source": [
    "## Import some stuff..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3DQi8Hvg_GSD"
   },
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import logging\n",
    "import importlib\n",
    "\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\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "NzO2IOsm_MfL"
   },
   "outputs": [],
   "source": [
    "import datetime\n",
    "import copy\n",
    "import xarray as xr\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import fsspec\n",
    "import kerchunk\n",
    "from kerchunk.grib2 import (\n",
    "    grib_tree, scan_grib, extract_datatree_chunk_index, strip_datavar_chunks, \n",
    "    reinflate_grib_store, AggregationType, read_store, write_store, parse_grib_idx,\n",
    "    build_idx_grib_mapping, map_from_index\n",
    ")\n",
    "import gcsfs\n",
    "\n",
    "pd.set_option('display.max_columns', None)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ueTGv6_g6plE"
   },
   "source": [
    "## Extract the zarr store metadata\n",
    "\n",
    "\n",
    "Pick a file, any file... Must be on GCS so that coords use the same file store as the data vars"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "TvOmUkr-6plE",
    "outputId": "e61f1f20-0c43-40bb-ca6a-4c5076e76a9f"
   },
   "outputs": [],
   "source": [
    "# Pick two files to build a grib_tree with the correct dimensions\n",
    "gfs_files = [\n",
    "    \"gs://global-forecast-system/gfs.20230928/00/atmos/gfs.t00z.pgrb2.0p25.f000\",\n",
    "    \"gs://global-forecast-system/gfs.20230928/00/atmos/gfs.t00z.pgrb2.0p25.f001\"\n",
    "]\n",
    "\n",
    "# This operation reads two of the large Grib2 files from GCS\n",
    "# scan_grib extracts the zarr kerchunk metadata for each individual grib message\n",
    "# grib_tree builds a zarr/xarray compatible hierarchical view of the dataset\n",
    "gfs_grib_tree_store = grib_tree([group for f in gfs_files for group in scan_grib(f)])\n",
    "# it is slow even in parallel because it requires a huge amount of IO\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 186
    },
    "id": "IEOG_v8O5KVC",
    "outputId": "4b8e9b14-efc2-440e-9953-bd15df58bcc3"
   },
   "outputs": [],
   "source": [
    "# The grib_tree can be opened directly using either zarr or xarray datatree\n",
    "# But this is too slow to build big aggregations\n",
    "gfs_dt = xr.open_datatree(fsspec.filesystem(\"reference\", fo=gfs_grib_tree_store).get_mapper(\"\"), engine=\"zarr\", consolidated=False)\n",
    "gfs_dt\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "R0Ew4zOIEoIw"
   },
   "source": [
    "## Separating static metadata from the chunk indexes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 423
    },
    "id": "ukcrJ9fi5NSS",
    "outputId": "0c9b6d89-883c-4fad-d891-6c62973c115f"
   },
   "outputs": [],
   "source": [
    "# The key metadata associated with each grib message can be extracted into a table\n",
    "gfs_kind = extract_datatree_chunk_index(gfs_dt, gfs_grib_tree_store, grib=True)\n",
    "gfs_kind"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "hX7Hh3aB6plF",
    "outputId": "ede04c84-cc23-4260-f8c2-6973f2ee56bc"
   },
   "outputs": [],
   "source": [
    "# While the static zarr metadata associated with the dataset can be seperated - created once.\n",
    "deflated_gfs_grib_tree_store = copy.deepcopy(gfs_grib_tree_store)\n",
    "strip_datavar_chunks(deflated_gfs_grib_tree_store)\n",
    "\n",
    "write_store(\"./\", deflated_gfs_grib_tree_store)\n",
    "\n",
    "print(\"Original references: \", len(gfs_grib_tree_store[\"refs\"]))\n",
    "print(\"Stripped references: \", len(deflated_gfs_grib_tree_store[\"refs\"]))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WlbvPP4G6plF"
   },
   "source": [
    "## Building it faster\n",
    "\n",
    "Okay that was fun - I promise you can recombine these pieces but it still takes the same amount of time to run scan_grib.\n",
    "\n",
    "The k(erchunk) index data looks a lot like the idx files that are present for every grib file in NODD's GCS archive though...\n",
    "\n",
    "\n",
    "```\n",
    "1:0:d=2023092800:PRMSL:mean sea level:1 hour fcst:\n",
    "2:990253:d=2023092800:CLWMR:1 hybrid level:1 hour fcst:\n",
    "3:1079774:d=2023092800:ICMR:1 hybrid level:1 hour fcst:\n",
    "4:1332540:d=2023092800:RWMR:1 hybrid level:1 hour fcst:\n",
    "5:1558027:d=2023092800:SNMR:1 hybrid level:1 hour fcst:\n",
    "6:1638489:d=2023092800:GRLE:1 hybrid level:1 hour fcst:\n",
    "7:1673516:d=2023092800:REFD:1 hybrid level:1 hour fcst:\n",
    "8:2471710:d=2023092800:REFD:2 hybrid level:1 hour fcst:\n",
    "9:3270627:d=2023092800:REFC:entire atmosphere:1 hour fcst:\n",
    "10:4144435:d=2023092800:VIS:surface:1 hour fcst:\n",
    "```\n",
    "\n",
    "But the metadata isn't quiet the same... they have mangled the attributes in the : seperated attributes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 628
    },
    "id": "XGNMVeWrHzJD",
    "outputId": "2c270e66-38ef-4acd-bc8f-62e3b1cd4e2b"
   },
   "outputs": [],
   "source": [
    "# We can pull this out into a dataframe, that starts to look a bit like what we got above extracted from the actual grib files\n",
    "# But this method runs in under a second reading a file that is less than 100k\n",
    "idxdf = parse_grib_idx(\n",
    "    basename=\"gs://global-forecast-system/gfs.20230901/00/atmos/gfs.t00z.pgrb2.0p25.f006\"\n",
    ")\n",
    "idxdf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 275
    },
    "id": "leAb3jXChBVQ",
    "outputId": "d6d842d0-e575-4c29-cd72-a4f999a946ed"
   },
   "outputs": [],
   "source": [
    "# Unfortunately, some accumulation variables have duplicate attributes making them\n",
    "# indesinguishable from the IDX file\n",
    "idxdf.loc[idxdf['attrs'].duplicated(keep=False), :]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 943
    },
    "id": "6ws0cK3gHzMK",
    "outputId": "62e74a5a-86db-4a71-93d8-120515108b2f"
   },
   "outputs": [],
   "source": [
    "# What we need is a mapping from our grib/zarr metadata to the attributes in the idx files\n",
    "# They are unique for each time horizon e.g. you need to build a unique mapping for the 1 hour\n",
    "# forecast, the 2 hour forecast... the 48 hour forecast.\n",
    "\n",
    "# let's make one for the 6 hour horizon. This requires reading both the grib and the idx file,\n",
    "# mapping the data for each grib message in order\n",
    "mapping = build_idx_grib_mapping(\n",
    "    basename=\"gs://global-forecast-system/gfs.20230928/00/atmos/gfs.t00z.pgrb2.0p25.f006\",\n",
    ")\n",
    "mapping"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 614
    },
    "id": "CzU9fWn5HzPS",
    "outputId": "770783d1-846a-47ce-c01f-02eea30939ef"
   },
   "outputs": [],
   "source": [
    "# Now if we parse the RunTime from the idx file name `gfs.20230901/00/`\n",
    "# We can build a fully compatible k_index\n",
    "mapped_index = map_from_index(\n",
    "    pd.Timestamp(\"2023-09-01T00\"),\n",
    "    mapping.loc[~mapping[\"attrs\"].duplicated(keep=\"first\"), :],\n",
    "    idxdf.loc[~idxdf[\"attrs\"].duplicated(keep=\"first\"), :]\n",
    ")\n",
    "mapped_index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mapping.loc[mapping.varname == \"sdswrf\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 770
    },
    "id": "Ny9xKXjP7n57",
    "outputId": "15a5793b-a9df-4065-a403-1a35aa5c7cf8"
   },
   "outputs": [],
   "source": [
    "mapped_index_list = []\n",
    "\n",
    "deduped_mapping = mapping.loc[~mapping[\"attrs\"].duplicated(keep=\"first\"), :]\n",
    "for date in pd.date_range(\"2023-09-01\", \"2023-09-30\"):\n",
    "  for runtime in range(0,24,6):\n",
    "    horizon=6\n",
    "    fname=f\"gs://global-forecast-system/gfs.{date.strftime('%Y%m%d')}/{runtime:02}/atmos/gfs.t{runtime:02}z.pgrb2.0p25.f{horizon:03}\"\n",
    "\n",
    "    idxdf = parse_grib_idx(\n",
    "        basename=fname\n",
    "    )\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",
    "    mapped_index_list.append(mapped_index)\n",
    "\n",
    "gfs_kind = pd.concat(mapped_index_list)\n",
    "gfs_kind\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "VF60ETYemD6N"
   },
   "source": [
    "## We just aggregated a 120 GFS grib files in 18 seconds!\n",
    "\n",
    "Lets build it back into a data tree!\n",
    "\n",
    "The reinflate_grib_store interface is pretty opaque but allows building any slice of an FMRC. A good area for future improvement, but for now, since\n",
    "we have just a single 6 hour horizon slice let's build that..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "IW1M7go56plF",
    "outputId": "17c2fa69-81f8-4d21-f260-4233e902a6ac"
   },
   "outputs": [],
   "source": [
    "axes = [\n",
    "  pd.Index(\n",
    "    [\n",
    "      pd.timedelta_range(start=\"0 hours\", end=\"6 hours\", freq=\"6h\", closed=\"right\", name=\"6 hour\"),\n",
    "    ],\n",
    "    name=\"step\"\n",
    "  ),\n",
    "  pd.date_range(\"2023-09-01T06:00\", \"2023-10T00:00\", freq=\"360min\", name=\"valid_time\")\n",
    "]\n",
    "axes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "_M9oH4TJ6plG",
    "outputId": "c6043655-7c8b-4981-cb5f-f5d6a903315b"
   },
   "outputs": [],
   "source": [
    "# It is fast to rebuild the datatree - but lets pull out two varables to look at...\n",
    "\n",
    "# If you skipped building the deflated store, read it here.\n",
    "#deflated_gfs_grib_tree_store = read_store(\"./\")\n",
    "\n",
    "gfs_store = reinflate_grib_store(\n",
    "    axes=axes,\n",
    "    aggregation_type=AggregationType.HORIZON,\n",
    "    chunk_index=gfs_kind.loc[gfs_kind.varname.isin([\"sdswrf\", \"t2m\"])],\n",
    "    zarr_ref_store=deflated_gfs_grib_tree_store\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "Eeml5x_d6plG",
    "outputId": "19052d52-3ca5-4d1e-9a8a-a0608381b342"
   },
   "outputs": [],
   "source": [
    "gfs_dt = xr.open_datatree(fsspec.filesystem(\"reference\", fo=gfs_store).get_mapper(\"\"), engine=\"zarr\", consolidated=False)\n",
    "gfs_dt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 296
    },
    "id": "MFCPjIJnpKh4",
    "outputId": "98c66ef2-8034-4b5e-f239-65e68c679e97"
   },
   "outputs": [],
   "source": [
    "# Reading the data - especially extracting point time series isn't any faster once you have\n",
    "# the xarray datatree. This is just a much faster way of building the aggregations than\n",
    "# directly running scan_grib over all the data first.\n",
    "gfs_dt.sdswrf.avg.surface.sdswrf[0,0:10,300,400].compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "gfs_dt.t2m.instant.heightAboveGround.t2m[0,0:10,300,400].compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 735
    },
    "id": "PSYb4OasbTKG",
    "outputId": "b4031b98-24fa-451d-e460-452dfad3695a"
   },
   "outputs": [],
   "source": [
    "gfs_dt.sdswrf.avg.surface.sdswrf[0,1,:,:].plot(figsize=(12,8))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 735
    },
    "id": "7zAfk1IshXMH",
    "outputId": "640f0354-d7f0-4a5f-f928-1ad1d0fcd481"
   },
   "outputs": [],
   "source": [
    "gfs_dt.sdswrf.avg.surface.sdswrf[0,2,:,:].plot(figsize=(12,8))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 735
    },
    "id": "T3gM0XLXhoZp",
    "outputId": "c37186b5-dbb1-4c98-c6e1-a787a5b85dae"
   },
   "outputs": [],
   "source": [
    "gfs_dt.sdswrf.avg.surface.sdswrf[0,3,:,:].plot(figsize=(12,8))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Timeseries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from joblib import parallel_config\n",
    "# with parallel_config(n_jobs=8):\n",
    "    #res = gfs_dt.dswrf.avg.surface.dswrf.interp(longitude=[320.5, 300.2], latitude=[20.6, 45.7], method=\"linear\")\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "res = gfs_dt.sdswrf.avg.surface.sdswrf.interp(longitude=[320.5], latitude=[20.6], method=\"linear\")\n",
    "res.plot()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "authorship_tag": "ABX9TyPQShGkdrEd0WnSDM4SoI/N",
   "include_colab_link": true,
   "provenance": []
  },
  "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
}