coawst_open_data_create_refs.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6a3a5c6e-97c3-44c0-9224-3756b6774479",
   "metadata": {},
   "source": [
    "# Create Parquet references for the COAWST forecast archive on AWS Open Data\n",
    "We use [kerchunk](https://fsspec.github.io/kerchunk/) to first create individual JSON reference files for each weekly NetCDF file, \n",
    "then create Parquet references that allow access as a single dataset in Xarray"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "outside-mayor",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import fsspec\n",
    "import xarray as xr\n",
    "\n",
    "from kerchunk.hdf import SingleHdf5ToZarr\n",
    "from kerchunk.combine import MultiZarrToZarr\n",
    "from fsspec.implementations.reference import LazyReferenceMapper\n",
    "\n",
    "from dask.distributed import Client\n",
    "import dask.bag as db\n",
    "from pathlib import Path\n",
    "import numpy as np\n",
    "import ujson"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64131bd5-8de0-4e83-aa51-c2dd3ac584e1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import kerchunk\n",
    "kerchunk.__version__"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f6e8365-aa8a-4d06-8661-129855593fda",
   "metadata": {},
   "source": [
    "We can read from AWS Open Data using `anon=True`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "endangered-therapist",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "fs_read = fsspec.filesystem('s3', anon=True, skip_instance_cache=True, use_listings_cache=False )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b949a69-7304-425b-876a-74d5f4cd0d92",
   "metadata": {},
   "source": [
    "We can't *write* to AWS Open Data without credentials, which we will specify through environment variables.  Because we are going to use environment variables instead of referencing an AWS profile, we don't specify `profile=` here in fs_write, but use `anon=False`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "native-serial",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "fs_write = fsspec.filesystem('s3', anon=False, skip_instance_cache=True, use_listings_cache=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "packed-lightning",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "flist = fs_read.glob('s3://usgs-coawst/useast-archive/*.nc')\n",
    "json_dir = 's3://usgs-coawst/useast-archive/individual_jsons'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "destroyed-abortion",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "flist = [f's3://{f}' for f in flist]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "charitable-logan",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(len(flist))\n",
    "print(flist[0])\n",
    "print(flist[-1])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cdc01098-ca89-498f-9e80-c2b523525811",
   "metadata": {},
   "source": [
    "#### Create references for each NetCDF file in parallel "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "separated-timber",
   "metadata": {},
   "outputs": [],
   "source": [
    "so = dict(mode='rb', anon=True, skip_instance_cache=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "close-assignment",
   "metadata": {},
   "outputs": [],
   "source": [
    "def gen_json(u):\n",
    "    with fs_read.open(u, **so) as infile:\n",
    "        fname = Path(u).stem\n",
    "        h5chunks = SingleHdf5ToZarr(infile, u, inline_threshold=300)\n",
    "        outf = f'{json_dir}/{fname}.json'\n",
    "        with fs_write.open(outf, 'wb') as f:\n",
    "            f.write(ujson.dumps(h5chunks.translate()).encode());\n",
    "    return outf"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7828088a-6fcb-455e-b536-72e69968cd37",
   "metadata": {},
   "source": [
    "Import Nebari helper functions and set credentials"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "987fcf4f-048b-45d2-85af-7beb9d748893",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import sys, os\n",
    "sys.path.append(os.path.join(os.environ['HOME'],'shared','users','lib'))\n",
    "import nebari_tools as nbt\n",
    "\n",
    "aws_profile = 'coawst_open_data'\n",
    "aws_region = 'us-west-2'\n",
    "endpoint_url = f's3.{aws_region}.amazonaws.com'\n",
    "nbt.set_credentials(profile=aws_profile, region=aws_region, endpoint_url=endpoint_url)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6568f677-71ce-4fc6-bb88-d395cbf5cb04",
   "metadata": {},
   "source": [
    "Create a Dask Gateway cluster on Nebari"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc3a7cd1-7c6f-4ea4-9279-9d6f2ac50297",
   "metadata": {},
   "outputs": [],
   "source": [
    "worker_max = 30\n",
    "client,cluster = nbt.start_dask_cluster(profile=aws_profile, worker_max=worker_max, \n",
    "                                      region=aws_region, use_existing_cluster=False,\n",
    "                                      adaptive_scaling=False, wait_for_cluster=False, \n",
    "                                      worker_profile='Small Worker', \n",
    "                                      propagate_env=True) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "567c63ca-e667-4830-b4cf-8104009a9282",
   "metadata": {},
   "source": [
    "Create the references for each file in parallel using Dask Bag.  Dask Bag splits a large list of tasks into partitions, and here we specify splitting the list into four partitions for each Dask worker.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "collectible-satellite",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "bag = db.from_sequence(flist, npartitions=worker_max*4).map(gen_json)\n",
    "bag.compute(retries=20)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "hourly-tucson",
   "metadata": {},
   "source": [
    "#### Create combined references and store in Parquet files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "636a01dd-334b-46fa-a7ec-52b2c3d567e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "json_list = fs_read.glob(f'{json_dir}/*.json')\n",
    "json_list = [f's3://{j}' for j in json_list]\n",
    "print(len(json_list))\n",
    "print(json_list[0])\n",
    "print(json_list[-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "383368a2-efac-4f77-ba26-96d9f5f76637",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# fs_write.rm(json_list)  # use this if you need to start over"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f7af26c-eb5c-42fa-95cc-b204c2763dde",
   "metadata": {},
   "source": [
    "Create a list of variables not to combine along the time dimension"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "511c0144-b32b-4236-a442-d39189a97d14",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = xr.open_dataset(fs_read.open(flist[0]))\n",
    "identical_dims = []\n",
    "for v in ds.variables.keys():\n",
    "    if 'ocean_time' not in ds[v].dims:\n",
    "        identical_dims.append(v)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fcbc1aa8-1d52-4ddf-8f95-259b0ced8c0a",
   "metadata": {},
   "source": [
    "Define pre and post-processing functions for the combined dataset references"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "protecting-recall",
   "metadata": {},
   "outputs": [],
   "source": [
    "import zarr\n",
    "\n",
    "def postprocess(refs):\n",
    "    tmp = zarr.open(refs)\n",
    "    tmp['ocean_time'].attrs['standard_name'] = 'time' \n",
    "    return refs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb1efa98-f24a-43b3-bcd7-c97fcfb42c98",
   "metadata": {},
   "source": [
    "Create a Lazy Reference Mapper object to be written locally with a large record_size so we don't get the references for each variable split into a bunch of parts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2110c3e4-24a1-4c85-9508-c12d655f8976",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "combined_parquet = 'combined.parq'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8db5e233-30cb-4b1b-99a2-d6fbe3798527",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "out = LazyReferenceMapper.create(combined_parquet, fs=None, record_size=100000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "western-sunset",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "_ = MultiZarrToZarr(\n",
    "        json_list,\n",
    "        remote_protocol=\"s3\",\n",
    "        concat_dims=[\"ocean_time\"],\n",
    "        coo_map={\"ocean_time\": \"cf:ocean_time\"},\n",
    "        identical_dims=identical_dims,\n",
    "        preprocess=kerchunk.combine.drop(\"dstart\"),\n",
    "        postprocess=postprocess,\n",
    "        out=out).translate()\n",
    "out.flush()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67874b84-873f-4f43-8653-eaa957eb0a54",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "combined_parquet_aws = 's3://usgs-coawst/useast-archive/combined.parq'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2595ab1-c8ec-4b5b-8977-017c7ed06ae6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "fs_write.rm(combined_parquet_aws, recursive=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa649096-e992-4752-96c2-955a0eb16bc1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "_ = fs_write.upload(combined_parquet, combined_parquet_aws, recursive=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0303aaa4-cb0a-4b62-a476-918546385ac3",
   "metadata": {},
   "source": [
    "#### Test opening combined dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70e09b6e-8620-4c98-9956-2d0b3ccf6d93",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "ds = xr.open_dataset(combined_parquet, engine='kerchunk', chunks={},\n",
    "                    backend_kwargs=dict(storage_options=dict(target_options=so,\n",
    "                    remote_protocol='s3', lazy=True, remote_options=so)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66dc5a8c-faa8-4482-a198-0886a66b0d36",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "ds.ocean_time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "radical-application",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "fs_ref = fsspec.implementations.reference.ReferenceFileSystem(\n",
    "            combined_parquet_aws, remote_protocol=\"s3\", target_protocol=\"s3\", lazy=True)\n",
    "\n",
    "ds = xr.open_dataset(fs_ref.get_mapper(), engine=\"zarr\", backend_kwargs={\"consolidated\": False}, chunks={}, drop_variables=['dstart'])\n",
    "\n",
    "ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d3e6f9b-696f-4337-9c1c-df4e2041252e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "ds.Vtransform.data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2faf3423-1051-4f67-b038-de107ab0e43f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import fsspec\n",
    "fs = fsspec.filesystem('s3',anon=True)\n",
    "xr.open_dataset(fs.open(flist[0])).Vtransform.data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8fb75393-a7be-43a8-bb31-2dddd752bcec",
   "metadata": {},
   "source": [
    "#### Write Intake Catalog"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec9ef534-44da-4969-8934-aed5787cec37",
   "metadata": {},
   "outputs": [],
   "source": [
    "cat_file = \"coawst_intake.yml\"\n",
    "\n",
    "dataset = 'COAWST-USEAST' \n",
    "    \n",
    "ymlentry = f\"\"\"# THIS FILE AUTO-GENERATED\n",
    "# This 'header' information needed to make this yml valid.  \n",
    "description: 'intake catalog example'\n",
    "metadata:\n",
    "  version: 1\n",
    "  description: \"Intake Catalog for data from USGS-COAWST model\"\n",
    "\n",
    "sources:\n",
    "# If you are copy/paste-ing this entry to another catalog, you only need the \n",
    "# lines from here down in your 'sources' section (be sure to indent as appropriate).\n",
    "  {dataset}:\n",
    "    driver: intake_xarray.xzarr.ZarrSource\n",
    "    description: 'USGS COAWST US East Coast and Gulf Coast Forecast Archive'\n",
    "    args:\n",
    "      urlpath: ['reference://']\n",
    "      consolidated: false\n",
    "      drop_variables: dstart\n",
    "      storage_options:\n",
    "        target_options:\n",
    "          anon: true\n",
    "          skip_instance_cache: true\n",
    "        fo: {combined_parquet_aws}\n",
    "        lazy: 'true'\n",
    "        remote_options:\n",
    "          anon: true\n",
    "          skip_instance_cache: true\n",
    "        remote_protocol: 's3'\n",
    "\"\"\"\n",
    "\n",
    "with open(cat_file, 'w') as ymlfile:\n",
    "   ymlfile.write(ymlentry)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce24238c-2866-42af-a606-49964a9e57d8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "cat_file_s3 = f's3://usgs-coawst/useast-archive/{cat_file}'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95ab5703-777c-426b-b07e-f158724e8835",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "_ = fs_write.upload(cat_file, cat_file_s3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aaa8ed5d-72e4-434b-8cc7-3abc481b3659",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "cat_file_s3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "665ee5f3-3997-45bb-9ad0-4f747bbcb331",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "cat_file_https = 'https://usgs-coawst.s3.amazonaws.com/useast-archive/coawst_intake.yml'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d318460b-3038-418f-a8b0-2cf9ef5779e7",
   "metadata": {},
   "source": [
    "#### Test Intake dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6ad76fc-9e4e-4c55-b673-1ff21e6ccd8d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import intake"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b52f3daa-ed72-4e71-a8c7-22576ef52c44",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "cat = intake.open_catalog(cat_file_https)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df1c1ad1-ec75-41e7-972a-42ac2ffcd770",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "list(cat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6807501a-ecb9-4b99-8688-8b92952a33ed",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "ds = cat['COAWST-USEAST'].to_dask()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19a7f0ea-e49f-4aa6-998f-45b498247d09",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "ds"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "75d9d822-7380-4baf-b9d1-a06e5de544a4",
   "metadata": {},
   "source": [
    "#### Shut down cluster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb8d94af-4613-4317-80a0-386f9f299f15",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "client.close();   cluster.shutdown()"
   ]
  }
 ],
 "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.6"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}