virtualizarr_gridmet.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### GRIDMET + Virtualizarr\n",
    "Notebook example on how to use VirtualiZarr to create a virtual dataset from the GRIDMET NetCDF files stored on www.northwestknowledge.net/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Imports\n",
    "Note: This uses parquet as the storage format instead of icechunk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np \n",
    "import xarray as xr \n",
    "from virtualizarr import open_virtual_dataset\n",
    "from dask.distributed import Client\n",
    "import dask \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create a list of all the year and variable combinations in GRIDMET"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "variables = [\"sph\", \"vpd\", \"pr\", \"rmin\", \"rmax\", \"srad\", \"tmmn\", \"tmmx\", \"vs\", \"th\", \"pet\", \"etr\", \"erc\", \"bi\", \"fm100\", \"fm1000\"]\n",
    "\n",
    "min_year = 1979\n",
    "max_year = 2020\n",
    "time_list = np.arange(min_year, max_year+1,1)\n",
    "\n",
    "combinations = [f\"https://www.northwestknowledge.net/metdata/data/{var}_{year}.nc\" for year in time_list for var in variables]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Start a dask distributed client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "client = Client(n_workers=8)\n",
    "client"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parallelize the reference generation with dask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def process(filename):\n",
    "    vds = open_virtual_dataset(filename,  indexes={})\n",
    "    return vds\n",
    "\n",
    "\n",
    "delayed_results = [dask.delayed(process)(filename) for filename in combinations]\n",
    "\n",
    "results = dask.compute(*delayed_results)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Combine the virtual datasets with Xarray and write"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "storage_options = dict(anon=True, default_fill_cache=False, default_cache_type=\"none\")\n",
    "# concat virtual datasets\n",
    "combined_vds = xr.concat(list(results), dim=\"time\", coords=\"minimal\", compat=\"override\")\n",
    "\n",
    "combined_vds.virtualize.to_kerchunk(\n",
    "    \"gridmet_1979_2020.parquet\", format=\"parquet\"\n",
    ")\n",
    "\n",
    "client.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Open local reference \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "combined_ds = xr.open_dataset(\"gridmet_1979_2020.parquet\", engine=\"kerchunk\", chunks={})\n",
    "combined_ds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### OR\n",
    "## Open public version on OSN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "combined_ds = xr.open_dataset(\"https://rice1.osn.mghpcc.org/carbonplan/virtual_datasets/gridmet/gridmet_1979_2020.parquet\", engine=\"kerchunk\", chunks={})\n",
    "combined_ds"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "virtualizarr",
   "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.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}