diff --git a/notebooks/benchmark_rasi_virtualzarr_vs_netcdf.ipynb b/notebooks/benchmark_rasi_virtualzarr_vs_netcdf.ipynb
new file mode 100644
index 0000000..a1d2f74
--- /dev/null
+++ b/notebooks/benchmark_rasi_virtualzarr_vs_netcdf.ipynb
@@ -0,0 +1,1998 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "id": "68130a2c",
+ "metadata": {},
+ "source": [
+ "# End-to-End Benchmark: RASI Virtual Zarr vs Original NetCDF\n",
+ "\n",
+ "This notebook benchmarks point time-series extraction for the **RASI ROUTING** dataset from:\n",
+ "\n",
+ "- **Virtual Zarr** via Icechunk\n",
+ "- **Original NetCDF files on S3**\n",
+ "\n",
+ "It follows the same benchmark structure as the earlier notebooks and measures:\n",
+ "\n",
+ "1. **Access time** \n",
+ " Time to initialize repository access or discover files on S3\n",
+ "\n",
+ "2. **Open time** \n",
+ " Time to open the dataset with xarray and subset the benchmark time range\n",
+ "\n",
+ "3. **Selection time** \n",
+ " Time to select nearest-neighbor point(s)\n",
+ "\n",
+ "4. **Load time** \n",
+ " Time to load the selected data into memory\n",
+ "\n",
+ "5. **End-to-end time** \n",
+ " Access + open + selection + load\n",
+ "\n",
+ "## Benchmark setup\n",
+ "- Benchmark cases: **1 point** and **5 points**\n",
+ "- Repeats: **1** by default\n",
+ "- Experiment: configurable (`HISTORICAL`, `SSP245`, `SSP585`)\n",
+ "- Variable: `RZSM_Percentiles`\n",
+ "- Percentile: **50** (median streamflow)\n",
+ "- Separate start/end dates for Virtual Zarr and NetCDF\n",
+ "\n",
+ "**Author:** Siddharth Chaudhary\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 1,
+ "id": "f04ae333",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import time\n",
+ "import datetime as dt\n",
+ "import numpy as np\n",
+ "import pandas as pd\n",
+ "import xarray as xr\n",
+ "import matplotlib.pyplot as plt\n",
+ "from dask.diagnostics import ProgressBar\n",
+ "from IPython.display import display, Markdown\n",
+ "\n",
+ "import icechunk\n",
+ "import s3fs\n",
+ "import fsspec\n",
+ "import netCDF4 as nc\n",
+ "\n",
+ "pd.set_option(\"display.max_columns\", None)\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "b71c1f8a",
+ "metadata": {},
+ "source": [
+ "## 1. Benchmark configuration"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 2,
+ "id": "c443216f",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Experiment: HISTORICAL\n",
+ "Variable: RZSM_Percentiles\n",
+ "Percentile: 50\n",
+ "Repeats: 1\n",
+ "Virtual Zarr period: 1960-01-01 to 2014-12-31\n",
+ "Original NetCDF period: 1960-01-01 to 2014-12-31\n"
+ ]
+ }
+ ],
+ "source": [
+ "experiment_ids = [\"HISTORICAL\", \"SSP245\", \"SSP585\"]\n",
+ "experiment_id = experiment_ids[0] # change as needed\n",
+ "\n",
+ "benchmark_cases = {\n",
+ " \"1 point\": {\n",
+ " \"lon\": [-100.0],\n",
+ " \"lat\": [30.0],\n",
+ " },\n",
+ " \"5 points\": {\n",
+ " \"lon\": [-100.0, -120.0, -90.0, -105.0, -75.0],\n",
+ " \"lat\": [30.0, 45.0, 35.0, 40.0, 42.0],\n",
+ " },\n",
+ "}\n",
+ "\n",
+ "variable_name = \"RZSM_Percentiles\"\n",
+ "percentile_value = 50\n",
+ "repeats = 1\n",
+ "\n",
+ "# Set benchmark windows independently if needed\n",
+ "vz_start = \"1960-01-01\"\n",
+ "vz_end = \"2014-12-31\"\n",
+ "\n",
+ "nc_start = \"1960-01-01\"\n",
+ "nc_end = \"2014-12-31\"\n",
+ "\n",
+ "print(\"Experiment:\", experiment_id)\n",
+ "print(\"Variable:\", variable_name)\n",
+ "print(\"Percentile:\", percentile_value)\n",
+ "print(\"Repeats:\", repeats)\n",
+ "print(\"Virtual Zarr period:\", vz_start, \"to\", vz_end)\n",
+ "print(\"Original NetCDF period:\", nc_start, \"to\", nc_end)\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "1507e2ea",
+ "metadata": {},
+ "source": [
+ "## 2. Helper functions"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 3,
+ "id": "8153c694",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "def build_month_strings(start_date, end_date):\n",
+ " months = pd.period_range(\n",
+ " start=pd.Period(start_date, freq=\"M\"),\n",
+ " end=pd.Period(end_date, freq=\"M\"),\n",
+ " freq=\"M\"\n",
+ " )\n",
+ " return [p.strftime(\"%Y%m\") for p in months]\n",
+ "\n",
+ "\n",
+ "def list_rasi_netcdf_files_s3(start_date, end_date,\n",
+ " experiment_id,\n",
+ " bucket=\"nasa-waterinsight\",\n",
+ " prefix_root=\"RASI/LSM\"):\n",
+ " fs = s3fs.S3FileSystem(anon=True)\n",
+ "\n",
+ " months = build_month_strings(start_date, end_date)\n",
+ "\n",
+ " files = []\n",
+ " for month in months:\n",
+ " pattern = (\n",
+ " f\"{bucket}/{prefix_root}/{experiment_id}/PERCENTILES/\"\n",
+ " f\"RZSM_percentiles_{month}.nc\"\n",
+ " )\n",
+ " matches = fs.glob(pattern)\n",
+ " for m in matches:\n",
+ " files.append(f\"s3://{m}\")\n",
+ "\n",
+ " files = sorted(files)\n",
+ "\n",
+ " if not files:\n",
+ " raise FileNotFoundError(\"No RASI NetCDF files found for the requested date range.\")\n",
+ "\n",
+ " return files\n",
+ "\n",
+ "\n",
+ "def open_rasi_virtual_zarr_with_timing(experiment_id, start_date=None, end_date=None):\n",
+ " t0 = time.perf_counter()\n",
+ "\n",
+ " storage = icechunk.s3_storage(\n",
+ " bucket=\"veda-data-store-staging\",\n",
+ " prefix=f\"water-insight/virtual-zarr-store/RASI-icechunk-v2/LSM/{experiment_id}/\",\n",
+ " from_env=True,\n",
+ " )\n",
+ "\n",
+ " chunk_url = \"s3://veda-data-store-staging/water-insight/RASI/\"\n",
+ " virtual_credentials = icechunk.containers_credentials({\n",
+ " chunk_url: icechunk.s3_credentials()\n",
+ " })\n",
+ "\n",
+ " repo = icechunk.Repository.open(\n",
+ " storage=storage,\n",
+ " authorize_virtual_chunk_access=virtual_credentials,\n",
+ " )\n",
+ "\n",
+ " session = repo.readonly_session(\"main\")\n",
+ " store = session.store\n",
+ " t1 = time.perf_counter()\n",
+ "\n",
+ " ds = xr.open_zarr(store, consolidated=False, zarr_version=3, chunks={})\n",
+ "\n",
+ " if start_date is not None and end_date is not None:\n",
+ " ds = ds.sel(time=slice(start_date, end_date))\n",
+ "\n",
+ " t2 = time.perf_counter()\n",
+ "\n",
+ " timing = {\n",
+ " \"source\": \"RASI Virtual Zarr\",\n",
+ " \"access_time_sec\": t1 - t0,\n",
+ " \"open_time_sec\": t2 - t1,\n",
+ " \"open_total_time_sec\": t2 - t0,\n",
+ " \"n_files\": np.nan,\n",
+ " }\n",
+ "\n",
+ " return ds, timing\n",
+ "\n",
+ "\n",
+ "def _normalize_rasi_netcdf_dataset(xr_ds):\n",
+ " \"\"\"\n",
+ " Convert RASI monthly NetCDF into a dataset with a proper time dimension.\n",
+ " \"\"\"\n",
+ "\n",
+ " # Extract scalar date value\n",
+ " raw_date = xr_ds[\"date\"].values[0]\n",
+ "\n",
+ " # Convert integer YYYYMM → datetime\n",
+ " time_value = pd.to_datetime(str(raw_date), format=\"%Y%m\")\n",
+ "\n",
+ " # Drop useless dimensions/variables\n",
+ " xr_ds = xr_ds.drop_vars(\"date\")\n",
+ "\n",
+ " if \"ncl0\" in xr_ds.dims:\n",
+ " xr_ds = xr_ds.squeeze(\"ncl0\", drop=True)\n",
+ "\n",
+ " # Add real time dimension\n",
+ " xr_ds = xr_ds.expand_dims(time=[time_value])\n",
+ "\n",
+ " return xr_ds\n",
+ "\n",
+ "def open_rasi_netcdf_s3_with_timing(start_date, end_date, experiment_id):\n",
+ " t0 = time.perf_counter()\n",
+ " files = list_rasi_netcdf_files_s3(start_date, end_date, experiment_id)\n",
+ " t1 = time.perf_counter()\n",
+ "\n",
+ " fs = fsspec.filesystem(\"s3\", anon=True)\n",
+ "\n",
+ " datasets = []\n",
+ " nc_handles = []\n",
+ "\n",
+ " for f in files:\n",
+ " s3_path = f.replace(\"s3://\", \"\")\n",
+ " with fs.open(s3_path, \"rb\") as fp:\n",
+ " raw = fp.read()\n",
+ "\n",
+ " # This pattern matches the approach that worked in prior exploration\n",
+ " ds_nc4 = nc.Dataset(\"dummy\", mode=\"r\", memory=raw)\n",
+ " nc_handles.append(ds_nc4)\n",
+ "\n",
+ " xr_ds = xr.open_dataset(xr.backends.NetCDF4DataStore(ds_nc4))\n",
+ " xr_ds = _normalize_rasi_netcdf_dataset(xr_ds)\n",
+ " datasets.append(xr_ds)\n",
+ "\n",
+ " # Sort monthly datasets by first timestamp, then concatenate along time\n",
+ " \n",
+ " datasets_sorted = sorted(\n",
+ " datasets,\n",
+ " key=lambda ds: pd.to_datetime(ds[\"time\"].values[0])\n",
+ " )\n",
+ "\n",
+ " ds_nc = xr.concat(\n",
+ " datasets_sorted,\n",
+ " dim=\"time\",\n",
+ " data_vars=\"all\",\n",
+ " coords=\"all\",\n",
+ " compat=\"override\"\n",
+ " )\n",
+ "\n",
+ " ds_nc = ds_nc.sortby(\"time\")\n",
+ " ds_nc = ds_nc.sel(time=slice(start_date, end_date))\n",
+ " t2 = time.perf_counter()\n",
+ "\n",
+ " timing = {\n",
+ " \"source\": \"Original NetCDF\",\n",
+ " \"access_time_sec\": t1 - t0,\n",
+ " \"open_time_sec\": t2 - t1,\n",
+ " \"open_total_time_sec\": t2 - t0,\n",
+ " \"n_files\": len(files),\n",
+ " }\n",
+ "\n",
+ " return ds_nc, files, timing, datasets, nc_handles\n",
+ "\n",
+ "\n",
+ "def benchmark_timeseries(\n",
+ " ds,\n",
+ " lons,\n",
+ " lats,\n",
+ " var=\"RZSM_Percentiles\",\n",
+ " percentile=50,\n",
+ " repeats=1,\n",
+ "):\n",
+ " results = []\n",
+ " loaded = None\n",
+ "\n",
+ " base = ds[var].sel(percentile=percentile)\n",
+ "\n",
+ " selection = {\n",
+ " \"lon\": xr.DataArray(lons, dims=[\"points\"]),\n",
+ " \"lat\": xr.DataArray(lats, dims=[\"points\"]),\n",
+ " }\n",
+ "\n",
+ " for run in range(repeats):\n",
+ " t0 = time.perf_counter()\n",
+ " subset = base.sel(**selection, method=\"nearest\")\n",
+ " t1 = time.perf_counter()\n",
+ "\n",
+ " with ProgressBar():\n",
+ " loaded = subset.load()\n",
+ " t2 = time.perf_counter()\n",
+ "\n",
+ " results.append({\n",
+ " \"run\": run + 1,\n",
+ " \"n_points\": len(lons),\n",
+ " \"select_time_sec\": t1 - t0,\n",
+ " \"load_time_sec\": t2 - t1,\n",
+ " \"query_total_time_sec\": t2 - t0,\n",
+ " \"n_time_steps\": loaded.sizes.get(\"time\", np.nan),\n",
+ " })\n",
+ "\n",
+ " return pd.DataFrame(results), loaded\n",
+ "\n",
+ "\n",
+ "def benchmark_cases_for_source(\n",
+ " ds,\n",
+ " benchmark_cases,\n",
+ " source_name,\n",
+ " open_timing,\n",
+ " var=\"RZSM_Percentiles\",\n",
+ " percentile=50,\n",
+ " repeats=1,\n",
+ "):\n",
+ " all_results = []\n",
+ " loaded_outputs = {}\n",
+ "\n",
+ " for case_name, coords in benchmark_cases.items():\n",
+ " print(f\"Running {source_name} benchmark for: {case_name}\")\n",
+ "\n",
+ " result_df, loaded = benchmark_timeseries(\n",
+ " ds=ds,\n",
+ " lons=coords[\"lon\"],\n",
+ " lats=coords[\"lat\"],\n",
+ " var=var,\n",
+ " percentile=percentile,\n",
+ " repeats=repeats,\n",
+ " )\n",
+ "\n",
+ " result_df[\"case\"] = case_name\n",
+ " result_df[\"source\"] = source_name\n",
+ " result_df[\"access_time_sec\"] = open_timing[\"access_time_sec\"]\n",
+ " result_df[\"open_time_sec\"] = open_timing[\"open_time_sec\"]\n",
+ " result_df[\"open_total_time_sec\"] = open_timing[\"open_total_time_sec\"]\n",
+ " result_df[\"end_to_end_time_sec\"] = (\n",
+ " result_df[\"open_total_time_sec\"] + result_df[\"query_total_time_sec\"]\n",
+ " )\n",
+ " result_df[\"n_files\"] = open_timing.get(\"n_files\", np.nan)\n",
+ "\n",
+ " all_results.append(result_df)\n",
+ " loaded_outputs[(source_name, case_name)] = loaded\n",
+ "\n",
+ " benchmark_df = pd.concat(all_results, ignore_index=True)\n",
+ " return benchmark_df, loaded_outputs\n",
+ "\n",
+ "\n",
+ "def summarize_results(benchmark_df):\n",
+ " return (\n",
+ " benchmark_df\n",
+ " .groupby([\"source\", \"case\", \"n_points\"], as_index=False)\n",
+ " .agg(\n",
+ " n_files=(\"n_files\", \"first\"),\n",
+ " access_time_sec=(\"access_time_sec\", \"first\"),\n",
+ " open_time_sec=(\"open_time_sec\", \"first\"),\n",
+ " open_total_time_sec=(\"open_total_time_sec\", \"first\"),\n",
+ " mean_select_time_sec=(\"select_time_sec\", \"mean\"),\n",
+ " mean_load_time_sec=(\"load_time_sec\", \"mean\"),\n",
+ " mean_query_time_sec=(\"query_total_time_sec\", \"mean\"),\n",
+ " mean_end_to_end_time_sec=(\"end_to_end_time_sec\", \"mean\"),\n",
+ " n_time_steps=(\"n_time_steps\", \"first\"),\n",
+ " )\n",
+ " )\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "ad058fe9",
+ "metadata": {},
+ "source": [
+ "## 3. Open datasets"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 4,
+ "id": "c80bbf4f",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "/tmp/ipykernel_137/2002258477.py:59: FutureWarning: zarr_version is deprecated, use zarr_format\n",
+ " ds = xr.open_zarr(store, consolidated=False, zarr_version=3, chunks={})\n"
+ ]
+ },
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "RASI Virtual Zarr timing:\n"
+ ]
+ },
+ {
+ "data": {
+ "text/html": [
+ "\n",
+ "\n",
+ "
\n",
+ " \n",
+ " \n",
+ " | \n",
+ " source | \n",
+ " access_time_sec | \n",
+ " open_time_sec | \n",
+ " open_total_time_sec | \n",
+ " n_files | \n",
+ "
\n",
+ " \n",
+ " \n",
+ " \n",
+ " | 0 | \n",
+ " RASI Virtual Zarr | \n",
+ " 0.204865 | \n",
+ " 3.016286 | \n",
+ " 3.221151 | \n",
+ " NaN | \n",
+ "
\n",
+ " \n",
+ "
\n",
+ "
\n",
+ "\n",
+ "
\n",
+ " \n",
+ " \n",
+ " | \n",
+ " source | \n",
+ " access_time_sec | \n",
+ " open_time_sec | \n",
+ " open_total_time_sec | \n",
+ " n_files | \n",
+ "
\n",
+ " \n",
+ " \n",
+ " \n",
+ " | 0 | \n",
+ " Original NetCDF | \n",
+ " 8.845641 | \n",
+ " 81.366121 | \n",
+ " 90.211763 | \n",
+ " 660 | \n",
+ "
\n",
+ " \n",
+ "
\n",
+ "
\n",
+ "\n",
+ "
\n",
+ " \n",
+ " \n",
+ " | \n",
+ " run | \n",
+ " n_points | \n",
+ " select_time_sec | \n",
+ " load_time_sec | \n",
+ " query_total_time_sec | \n",
+ " n_time_steps | \n",
+ " case | \n",
+ " source | \n",
+ " access_time_sec | \n",
+ " open_time_sec | \n",
+ " open_total_time_sec | \n",
+ " end_to_end_time_sec | \n",
+ " n_files | \n",
+ "
\n",
+ " \n",
+ " \n",
+ " \n",
+ " | 0 | \n",
+ " 1 | \n",
+ " 1 | \n",
+ " 0.018154 | \n",
+ " 1.686319 | \n",
+ " 1.704473 | \n",
+ " 660 | \n",
+ " 1 point | \n",
+ " RASI Virtual Zarr | \n",
+ " 0.204865 | \n",
+ " 3.016286 | \n",
+ " 3.221151 | \n",
+ " 4.925624 | \n",
+ " NaN | \n",
+ "
\n",
+ " \n",
+ " | 1 | \n",
+ " 1 | \n",
+ " 5 | \n",
+ " 0.018313 | \n",
+ " 1.503183 | \n",
+ " 1.521496 | \n",
+ " 660 | \n",
+ " 5 points | \n",
+ " RASI Virtual Zarr | \n",
+ " 0.204865 | \n",
+ " 3.016286 | \n",
+ " 3.221151 | \n",
+ " 4.742647 | \n",
+ " NaN | \n",
+ "
\n",
+ " \n",
+ " | 2 | \n",
+ " 1 | \n",
+ " 1 | \n",
+ " 0.001255 | \n",
+ " 0.000151 | \n",
+ " 0.001406 | \n",
+ " 660 | \n",
+ " 1 point | \n",
+ " Original NetCDF | \n",
+ " 8.845641 | \n",
+ " 81.366121 | \n",
+ " 90.211763 | \n",
+ " 90.213169 | \n",
+ " 660.0 | \n",
+ "
\n",
+ " \n",
+ " | 3 | \n",
+ " 1 | \n",
+ " 5 | \n",
+ " 0.001294 | \n",
+ " 0.000154 | \n",
+ " 0.001448 | \n",
+ " 660 | \n",
+ " 5 points | \n",
+ " Original NetCDF | \n",
+ " 8.845641 | \n",
+ " 81.366121 | \n",
+ " 90.211763 | \n",
+ " 90.213211 | \n",
+ " 660.0 | \n",
+ "
\n",
+ " \n",
+ "
\n",
+ "
\n",
+ "\n",
+ "
\n",
+ " \n",
+ " \n",
+ " | \n",
+ " source | \n",
+ " case | \n",
+ " n_points | \n",
+ " n_files | \n",
+ " access_time_sec | \n",
+ " open_time_sec | \n",
+ " open_total_time_sec | \n",
+ " mean_select_time_sec | \n",
+ " mean_load_time_sec | \n",
+ " mean_query_time_sec | \n",
+ " mean_end_to_end_time_sec | \n",
+ " n_time_steps | \n",
+ "
\n",
+ " \n",
+ " \n",
+ " \n",
+ " | 0 | \n",
+ " Original NetCDF | \n",
+ " 1 point | \n",
+ " 1 | \n",
+ " 660.0 | \n",
+ " 8.845641 | \n",
+ " 81.366121 | \n",
+ " 90.211763 | \n",
+ " 0.001255 | \n",
+ " 0.000151 | \n",
+ " 0.001406 | \n",
+ " 90.213169 | \n",
+ " 660 | \n",
+ "
\n",
+ " \n",
+ " | 1 | \n",
+ " Original NetCDF | \n",
+ " 5 points | \n",
+ " 5 | \n",
+ " 660.0 | \n",
+ " 8.845641 | \n",
+ " 81.366121 | \n",
+ " 90.211763 | \n",
+ " 0.001294 | \n",
+ " 0.000154 | \n",
+ " 0.001448 | \n",
+ " 90.213211 | \n",
+ " 660 | \n",
+ "
\n",
+ " \n",
+ " | 2 | \n",
+ " RASI Virtual Zarr | \n",
+ " 1 point | \n",
+ " 1 | \n",
+ " NaN | \n",
+ " 0.204865 | \n",
+ " 3.016286 | \n",
+ " 3.221151 | \n",
+ " 0.018154 | \n",
+ " 1.686319 | \n",
+ " 1.704473 | \n",
+ " 4.925624 | \n",
+ " 660 | \n",
+ "
\n",
+ " \n",
+ " | 3 | \n",
+ " RASI Virtual Zarr | \n",
+ " 5 points | \n",
+ " 5 | \n",
+ " NaN | \n",
+ " 0.204865 | \n",
+ " 3.016286 | \n",
+ " 3.221151 | \n",
+ " 0.018313 | \n",
+ " 1.503183 | \n",
+ " 1.521496 | \n",
+ " 4.742647 | \n",
+ " 660 | \n",
+ "
\n",
+ " \n",
+ "
\n",
+ "