2_hycom_access_refs.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "80447fc4-36b8-42b4-9ea2-42e3edc67322",
   "metadata": {},
   "source": [
    "# Access HYCOM Reanalysis data on AWS Open Data\n",
    "Access the  63,341 NetCDF 64-bit offset files from the HYbrid Coordinate Ocean Model (HYCOM) Global Ocean Forecast System Reanalysis (1994-2015) on [AWS Open Data](https://registry.opendata.aws/hycom-gofs-3pt1-reanalysis/) as a single Kerchunk-generated virtual dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d143b938-7bb1-4cf3-99de-f52c9596cbac",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import fsspec\n",
    "import xarray as xr\n",
    "import hvplot.xarray\n",
    "import intake"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "818b36c4-3b7a-49d7-b9c5-a1507f99fce1",
   "metadata": {},
   "source": [
    "## Open the virtual dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98280b1c-6d0c-488a-adb4-3002f039e15f",
   "metadata": {},
   "source": [
    "Open using Intake:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b32ab21-68fd-4cde-89ff-a14718b865f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "cat = intake.open_catalog('https://ncsa.osn.xsede.org/esip/rsignell/hycom.yaml')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92b04b36-bd23-44c5-9cc4-6dd3887282ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "ds = cat['gofs-3pt1'].read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74357618-9418-4afa-99f1-ed4a7bb350ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The much more verbose way non-Intake way to open the dataset using the xarray kerchunk engine:\n",
    "# combined_parquet_aws = 's3://esip/rsignell/hycom.parq'\n",
    "\n",
    "# so = dict(anon=True)    # data stored on AWS Open Data S3\n",
    "# to = dict(anon=True,    # refs stored on OSN \n",
    "#           client_kwargs={'endpoint_url': 'https://ncsa.osn.xsede.org'})\n",
    "\n",
    "# ds = xr.open_dataset(combined_parquet_aws, engine='kerchunk', chunks={},\n",
    "#                     backend_kwargs=dict(storage_options=dict(target_options=to,\n",
    "#                     remote_protocol='s3', lazy=True, remote_options=so)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02657266-596e-491e-869f-b9aebe620c4e",
   "metadata": {},
   "source": [
    "## Read Data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5db3458f-656d-41d8-a5af-496e6922a64a",
   "metadata": {},
   "source": [
    "<u> Case 1: Read a 3D field at a specific time step.  </u>\n",
    "\n",
    "We don't need a cluster for this since just one chunk of data (24MB) is actually loaded"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27d09f2f-7924-4850-8c76-d39ecd36adde",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds['water_temp'].sel(depth=0, time='2012-10-29 17:00', method='nearest')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7408862b-20d8-4130-ab33-de7b93755075",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "da = ds['water_temp'].sel(depth=0, time='2012-10-29 17:00', method='nearest').load()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa1f300a-ef2c-4806-bcda-61c62a2196c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "da.hvplot.quadmesh(x='lon', y='lat', geo=True, global_extent=True, tiles='ESRI', cmap='viridis', rasterize=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1084dfe4-5017-451b-91f4-b773df04c85b",
   "metadata": {},
   "source": [
    "<u> Case 2: Load a time series at a specific location. </u>\n",
    "\n",
    "Because each chunk only contains one time value, we want to read chunks in parallel using a Dask cluster.  Here we use coiled.io to generate a cluster on AWS in the same region as the data.   We also specify a cheap ARM instance type to keep the cost low. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9caf6e7b-fb7e-4016-a960-22ebfd9aa3df",
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster_type = 'Coiled'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59c1c3ad-c4af-4097-8dad-3fbe67787384",
   "metadata": {},
   "outputs": [],
   "source": [
    "if cluster_type == 'Local':\n",
    "    from dask.distributed import Client\n",
    "\n",
    "    client = Client()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04c3e4c9-85ea-4a76-8305-7b40b391733b",
   "metadata": {},
   "outputs": [],
   "source": [
    "if cluster_type == 'Coiled':\n",
    "    import coiled\n",
    "    cluster = coiled.Cluster(\n",
    "        region=\"us-west-2\",\n",
    "        arm=True,\n",
    "        worker_vm_types=[\"t4g.small\"],  # cheap, small ARM instances, 2cpus, 2GB RAM\n",
    "        worker_options={'nthreads':2},\n",
    "        n_workers=100,\n",
    "        wait_for_workers=True,\n",
    "        compute_purchase_option=\"spot_with_fallback\",\n",
    "        name='hycom',\n",
    "        software='esip-pangeo-arm',\n",
    "        workspace='esip-lab',\n",
    "        timeout=300   # leave cluster running for 5 min in case we want to use it again\n",
    "    )\n",
    "\n",
    "    client = cluster.get_client()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b49caf7-bfb3-4889-927c-facbd75b8133",
   "metadata": {},
   "source": [
    "We open the dataset again and tell Dask to load 20 time values (20 chunks) for each task.  Loading multiple time steps means we only incur object storage latency once for each task.  Each task will use more memory, however.  We picked 20 because it fits within memory on the 2GB ARM `t4g.small` instance types (and larger than 20 doesn't give significant performance benefit). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e02a3b62-0b52-4751-a36a-6de43d002730",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = cat['gofs-3pt1'].read(chunks={\"time\": 20})   # Intake way to re-open the dataset with different Dask chunks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c6962ff-4598-490d-be15-1ecdae60a89b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# non-Intake way to re-open the dataset with different Dask chunks:\n",
    "# ds = xr.open_dataset(combined_parquet_aws, engine='kerchunk', chunks={'time':20},\n",
    "#                    backend_kwargs=dict(storage_options=dict(target_options=to,\n",
    "#                    remote_protocol='s3', lazy=True, remote_options=so)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f68145cc-3a88-4ee8-9402-086680dec171",
   "metadata": {},
   "source": [
    "Extract the time series:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1f88ea2-309c-43a6-aef1-5193cad56840",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "da = ds['water_temp'].isel(depth=0).sel(lon=-69.6, lat=42.5, method='nearest').load(retries=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "895ba87c-48eb-4a4f-b1bf-5e055293ce41",
   "metadata": {},
   "source": [
    "Interactive visualization of the time series:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b775977-3bd2-4502-8b76-29ca52b1f5f9",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "da.hvplot(x='time', grid=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "global-global-pangeo",
   "language": "python",
   "name": "conda-env-global-global-pangeo-py"
  },
  "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}