docs: add example notebook showing how to work with NNJA data locally

example_notebooks/local_nnja.ipynb

@@ -0,0 +1,154 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Using NNJA-AI data locally\n",
+    "\n",
+    "Some users might not want to rely on the data in the official NNJA repositories and want data closer to their compute (maybe a local copy is preferred, or a copy in a lower latency cloud bucket). This is supported! You can make a copy of the dataset wherever you want. You can either continue to use the NNJA-AI SDK or treat your copy as regular parquet files. Here we'll show an example of both approaches with a local copy of a subset of the data. \n",
+    "\n",
+    "Here we've completely downloaded the conventional NNJA data (surface station and rawinsondes) to a local directory. This is a rich dataset but thankfully not the highest volume, about 90 GiB of data.  "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "! gcloud storage cp -r gs://nnja-ai/data/v1/conv/ ~/data/nnja-local-demo/\n",
+    "#this will take a long time and download 90 GiB of data"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We'll show three ways loading the data for this sample:\n",
+    "1) Using a modified Catalog\n",
+    "2) skipping the Catalog and directly instantiating \n",
+    "3) skipping NNJA entirely and just loading the parquets\n",
+    "\n",
+    "## 1) Use the Catalog\n",
+    "\n",
+    "The catalog is a json that points to the various root directories of different datasets, and the DataCatalog class takes in a json relative path and base path for the entire projecct. Assuming we stick with the relative organization of data in the existing NNJA-AI cloud buckets, we can just download the existing catalog, delete all the references to other datasets, and change the base path. "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "! gcloud storage cp gs://nnja-ai/data/v1/catalog.json ~/data/nnja-local-demo/\n",
+    "# Make sure to prune this for all datasets except the ones downloaded; in this case we keep all the \"conv\" datasets"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from nnja_ai import DataCatalog\n",
+    "import pandas as pd\n",
+    "local_dir = '~/data/nnja-local-demo'\n",
+    "catalog = DataCatalog(mirror=None, base_path=local_dir, catalog_json='catalog.json')\n",
+    "date = pd.to_datetime(\"2021-01-01\").tz_localize(\"UTC\")\n",
+    "metar_dataset = catalog[\"conv-adpsfc-NC000001\"]\n",
+    "print(metar_dataset)\n",
+    "ds = metar_dataset.sel(time=date).load_dataset(backend=\"pandas\")\n",
+    "print(ds.head(2))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 2) Skip the Catalog and directly instantiate a NNJADataset\n",
+    "\n",
+    "The Catalog isn't that useful if you have already identifed and downloaded the datasets of interest. If you still find it helpful to have other features of the NNJADataset metadata (column descriptors, etc.), you can directly create a NNJADataset from the downloaded directory's metadata file:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from nnja_ai import NNJADataset\n",
+    "metar_dataset = NNJADataset(json_uri=f'{local_dir}/conv/adpsfc/NC000001/.pmetadata', base_path=local_dir)\n",
+    "print(metar_dataset)\n",
+    "ds = metar_dataset.sel(time=date).load_dataset(backend=\"pandas\")\n",
+    "ds.head(2)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 3) Skip the NNJA-AI SDK and go read files directly\n",
+    "\n",
+    "There is no requirement to use this SDK; the NNJA-AI files are published in parquet for maximum flexibility. You can directly read the files yourself using your preferred backend:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# reading a single file with pandas\n",
+    "parquet_file = f'{local_dir}/conv/adpsfc/NC000001/OBS_DATE=2021-01-01/gdas.20210101.NC000001.parquet'\n",
+    "ds = pd.read_parquet([parquet_file], engine=\"pyarrow\")\n",
+    "ds.head(2)\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Or using polars and filtering:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import glob\n",
+    "import os\n",
+    "import polars as pl\n",
+    "import datetime\n",
+    "parquet_dir = os.path.expanduser(f'{local_dir}/conv/adpsfc/NC000001')\n",
+    "parquet_files = glob.glob(f'{parquet_dir}/*/*.parquet')\n",
+    "\n",
+    "ds = pl.scan_parquet(parquet_files, hive_partitioning=True)\n",
+    "ds.filter(pl.col(\"OBS_DATE\") == datetime.date(2021, 1, 1)).head(2).collect().to_pandas()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "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.15"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}