0_hycom_generate_refs.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6a2ae2d4-7cd9-46fb-acc5-a030a3298505",
   "metadata": {},
   "source": [
    "# Generate references for HYCOM using kerchunk\n",
    "HYCOM data on AWS Open Data are stored in 63,341 NetCDF 64-bit offset files. The data in these files are stored as short integers with scale_factor and add_offset, but because these are not NetCDF4 files, there is no compression and no chunking. Each file contains one time step of data. \n",
    "\n",
    "We generate references for each file, and use kerchunk.utils.subchunk to create virtual chunks so that each vertical layer is treated as a chunk.  For this uncompressed data, the byte-ranges are the same for each file, so we only need to create references for one file and then clone that for all the files, changing only the URL and the time value. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59d48df1-290d-4cbd-98e4-6ad40fc1ce7e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from kerchunk.netCDF3 import NetCDF3ToZarr\n",
    "from kerchunk.combine import MultiZarrToZarr, auto_dask, JustLoad\n",
    "from kerchunk.utils import subchunk, inline_array\n",
    "from fsspec.implementations.reference import LazyReferenceMapper\n",
    "import fsspec\n",
    "import xarray as xr\n",
    "import datetime as dt\n",
    "import copy\n",
    "import kerchunk\n",
    "import base64\n",
    "import struct\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75d17cb9-e9ee-41a3-b4dc-ef9c580444c5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "fs = fsspec.filesystem('s3', anon=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46059963-8917-4989-aefc-2057ba8de4b6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "flist = fs.glob('hycom-gofs-3pt1-reanalysis/*/*.nc')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd3dff84-e61a-45f4-b1ad-e591fec46cdb",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(flist)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a41842d-4434-4a0f-bb3d-74820d76146f",
   "metadata": {},
   "outputs": [],
   "source": [
    "flist[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff666ef5-6b82-4a46-b41c-4258a50aa14d",
   "metadata": {},
   "source": [
    "Method to generate references.   Need only to use for the first file!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03f30c58-6933-42e8-bae9-4edf9063350d",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "d0 = NetCDF3ToZarr(\"s3://\" + flist[0], storage_options={\"anon\": True},\n",
    "                  inline_threshold=400, version=2).translate()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54ee1005-7786-457b-aa73-960dbdca067a",
   "metadata": {},
   "source": [
    "Subchunk the 4D data vars"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b482764-53c1-4f80-9575-6bb70a015838",
   "metadata": {},
   "outputs": [],
   "source": [
    "for v in ['salinity', 'water_temp', 'water_u', 'water_v']:\n",
    "    d0 = subchunk(store=d0, variable=v, factor=40)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "722ad1f4-f980-4071-9a50-4bcf0b9e8786",
   "metadata": {},
   "source": [
    "#### Open the references for the first file in the dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ec87f26-1778-4aa7-a21d-95913d8c58b6",
   "metadata": {},
   "source": [
    "Storage options (for accessing the NetCDF files from AWS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f373501-5e6b-470e-ac8a-c11f54f77308",
   "metadata": {},
   "outputs": [],
   "source": [
    "so = dict(anon=True, skip_instance_cache=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcc0786b-5e44-42be-8a84-51a24b22e557",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = xr.open_dataset(d0, engine='kerchunk', chunks={}, drop_variables='tau', \n",
    "                     backend_kwargs=dict(storage_options=dict(\n",
    "                    remote_protocol='s3', lazy=False, remote_options=so)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da85d8b0-270b-491c-999c-0d17b2370632",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "ds"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d1d320d-194e-42c5-8bd4-4a84f1ceea0b",
   "metadata": {},
   "source": [
    "Define some functions to replace all the URLs in the reference dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a498a97b-3218-4cc5-be3e-a9e70752181d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def float_to_base64(number):\n",
    "    # Pack the float into bytes\n",
    "    packed = struct.pack('>d', number)\n",
    "    \n",
    "    # Encode the bytes to base64\n",
    "    encoded = base64.b64encode(packed)\n",
    "    # Convert bytes to string and return\n",
    "    return encoded.decode('utf-8')\n",
    "\n",
    "# Example usage\n",
    "float_num = 122748.\n",
    "encoded_str = float_to_base64(float_num)\n",
    "print(f\"Original number: {float_num}\")\n",
    "print(f\"Base64 encoded: {encoded_str}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3aad6f84-2c26-49c1-b2da-1a6be800a780",
   "metadata": {},
   "outputs": [],
   "source": [
    "def replace_first_item(d, target_string, replacement_string):\n",
    "    for key, value in d.items():\n",
    "        if isinstance(value, dict):\n",
    "            # Recursively process nested dictionaries\n",
    "            replace_first_item(value, target_string, replacement_string)\n",
    "        elif isinstance(value, list) and value and isinstance(value[0], str):\n",
    "            # Check if the value is a non-empty list and the first item is a string\n",
    "            if value[0] == target_string:\n",
    "                value[0] = replacement_string\n",
    "    return d\n",
    "#replace_first_item(d, f's3://{flist[0]}', f's3://{flist[1]}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c35fb24-18b4-4db8-8324-ce4c8ea43a1e",
   "metadata": {},
   "source": [
    "Function to generate the time from the filename"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2529a74-833d-4e9a-949f-069f33536dfa",
   "metadata": {},
   "outputs": [],
   "source": [
    "def name2date(f):\n",
    "    year = f[51:55]\n",
    "    month = f[55:57]\n",
    "    day = f[57:59]\n",
    "#    hour = f[59:61]  #always 12 for this dataset\n",
    "    tau = f[64:66]\n",
    "    return dt.datetime(int(year), int(month), int(day), int(tau))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a07ecde-1e7e-4fcf-b33e-b4ac6faa80e0",
   "metadata": {},
   "source": [
    "Loop through all the files, generating the references for each file by replacing the URL and date in the reference dict template"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91b1d506-28af-4521-878e-a60794e28619",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "dlist = []\n",
    "time0 = dt.datetime(2000,1,1,0)\n",
    "for i,v in enumerate(flist):\n",
    "    dmod = copy.deepcopy(d0)\n",
    "    time1 = name2date(v) + dt.timedelta(hours=12)\n",
    "    time_val = (time1 - time0).total_seconds()/3600 \n",
    "    encoded_str = float_to_base64(time_val)\n",
    "    dmod['time/0'] = f'base64:{encoded_str}'\n",
    "    dmod = replace_first_item(dmod, f's3://{flist[0]}',f's3://{v}')\n",
    "    dlist.append(dmod)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da769023-e95d-417b-bc8c-bc7b8cd2d115",
   "metadata": {},
   "source": [
    "Generate the combined references and save to Parquet storage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da4c841b-b57e-4c42-8f64-e0d173eed832",
   "metadata": {},
   "outputs": [],
   "source": [
    "combined_parquet = 'hycom.parq'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1367ad0a-7af4-4d1d-a8a1-9697f43cc823",
   "metadata": {},
   "outputs": [],
   "source": [
    "out = LazyReferenceMapper.create(combined_parquet, fs=None, record_size=100000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e76d48fd-4431-4656-8843-43e936e48bc2",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "_ = MultiZarrToZarr(\n",
    "        dlist,\n",
    "        remote_protocol=\"s3\",\n",
    "        concat_dims=\"time\",\n",
    "        identical_dims=['lon', 'lat', 'depth'],\n",
    "        preprocess=kerchunk.combine.drop(\"tau\"),\n",
    "        out=out).translate()\n",
    "out.flush()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43555956-e7f1-4ebb-adcd-60b652d812fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "fs_write = fsspec.filesystem('s3', profile='osn-esip', skip_instance_cache=True, use_listings_cache=False,\n",
    "                             client_kwargs={'endpoint_url': 'https://ncsa.osn.xsede.org'})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35203549-36cd-4d6d-abd6-0ff5a8333626",
   "metadata": {},
   "outputs": [],
   "source": [
    "combined_parquet_aws = 's3://esip/rsignell/hycom.parq'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6768ad6-0c88-4503-a1e5-bea330072fc8",
   "metadata": {},
   "outputs": [],
   "source": [
    "fs_write.rm(combined_parquet_aws, recursive=True)    # delete any existing refs on OSN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "acf015de-1147-4382-9102-08971a016751",
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = fs_write.upload(combined_parquet, combined_parquet_aws, recursive=True)  # upload refs to OSN"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33fd6eae-eab4-40b8-a07d-810770c37f95",
   "metadata": {},
   "source": [
    "Check to make sure the references got updated"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01f19181-7377-49c8-83f9-fec652a7a8ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "fs_write.info(f'{combined_parquet_aws}/.zmetadata')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d9a3de6-d497-41e4-8cae-ddf0dc3ccc7e",
   "metadata": {},
   "source": [
    "#### Open the references for the entire dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f6c85e5-b781-48e2-bf8f-86dbf1d5334b",
   "metadata": {},
   "source": [
    "Target options (for accessing the reference files from OSN)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0846cab-0be9-4bbb-bcb3-bc1fde91ebf2",
   "metadata": {},
   "outputs": [],
   "source": [
    "to = dict(anon=True, skip_instance_cache=True, \n",
    "          client_kwargs={'endpoint_url': 'https://ncsa.osn.xsede.org'})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e92a92e8-390f-443c-9a1d-717c02e837da",
   "metadata": {},
   "outputs": [],
   "source": [
    "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": "code",
   "execution_count": null,
   "id": "78d7005e-9428-46e3-a3bf-6ffd22030beb",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "global-global-pangeo-dev",
   "language": "python",
   "name": "conda-env-global-global-pangeo-dev-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
}