notebooks/benchmark-h5repack.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8a7ac7b4-40fe-4396-8531-76ecd67a040b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import cProfile\n",
    "import earthaccess\n",
    "import h5py\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import s3fs\n",
    "import tqdm\n",
    "import xarray as xr\n",
    "\n",
    "try:\n",
    "    import h5coro\n",
    "except:\n",
    "    !pip install -U git+https://github.com/ICESat2-SlideRule/h5coro.git@8035f01c4b793313496e246870a53f40056407c8\n",
    "    import h5coro\n",
    "\n",
    "try:\n",
    "    from gedi_subset.h5frame import H5DataFrame\n",
    "except ImportError:\n",
    "    !pip install git+https://github.com/MAAP-Project/gedi-subsetter.git@0.6.0\n",
    "    from gedi_subset.h5frame import H5DataFrame\n",
    "\n",
    "from h5coro import h5coro, s3driver, filedriver\n",
    "h5coro.config(errorChecking=True, verbose=False, enableAttributes=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66d2a3bb-5ff8-429b-abf4-1a4972cb6a20",
   "metadata": {
    "tags": [],
    "user_expressions": []
   },
   "source": [
    "# Hi!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd29c436-f9b0-489c-b9e3-26749bd07d68",
   "metadata": {
    "tags": [],
    "user_expressions": []
   },
   "source": [
    "# Benchmarking h5repack Format\n",
    "\n",
    "This is a HDF5 format that has been rechunked from the original format from NSIDC,\n",
    "done by Luis Lopez.\n",
    "\n",
    "We will benchmark the following workflows:\n",
    "\n",
    "1. Reading 1 full group (`h_ph`)\n",
    "2. Spatially subsetting that group (tbd)\n",
    "\n",
    "via the following libraries: h5py, h5coro and xarray"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7750e965-d0c3-4fca-97d5-3f81f0b558e6",
   "metadata": {
    "tags": [],
    "user_expressions": []
   },
   "source": [
    "## Workflow 1 - Read full group\n",
    "\n",
    "### Setup Steps\n",
    "\n",
    "NOTE, below we use the `h5repack` directory, this should be replaced if you are working with a different dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ef607f86-b21f-46fd-95a2-36203d82ccef",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2023-08-09 05:30:04 7760000000 ATL03_20181120182818_08110112_006_02_repacked.h5\n",
      "2023-08-09 05:30:04 7008000000 ATL03_20190219140808_08110212_006_02_repacked.h5\n",
      "2023-08-09 05:30:04 6936000000 ATL03_20200217204710_08110612_006_01_repacked.h5\n",
      "2023-08-09 05:30:04 8400000000 ATL03_20211114142614_08111312_006_01_repacked.h5\n",
      "2023-08-09 05:30:04 7960000000 ATL03_20230211164520_08111812_006_01_repacked.h5\n"
     ]
    }
   ],
   "source": [
    "# Checkout the files\n",
    "# !mamba install -c conda-forge awscli -y\n",
    "!aws s3 ls s3://nasa-cryo-scratch/h5cloud/h5repack/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "36f2c445-b576-4d22-968c-66ae166a3dcb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Set the variables\n",
    "bucket = 'nasa-cryo-scratch'\n",
    "directory = 'h5cloud/h5repack/'\n",
    "group = '/gt2l/heights'\n",
    "variable = 'h_ph'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8e42a3b7-ed22-411c-9863-c93569c18f20",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(['s3://nasa-cryo-scratch/h5cloud/h5repack/ATL03_20181120182818_08110112_006_02_repacked.h5',\n",
       "  's3://nasa-cryo-scratch/h5cloud/h5repack/ATL03_20190219140808_08110212_006_02_repacked.h5',\n",
       "  's3://nasa-cryo-scratch/h5cloud/h5repack/ATL03_20200217204710_08110612_006_01_repacked.h5',\n",
       "  's3://nasa-cryo-scratch/h5cloud/h5repack/ATL03_20211114142614_08111312_006_01_repacked.h5',\n",
       "  's3://nasa-cryo-scratch/h5cloud/h5repack/ATL03_20230211164520_08111812_006_01_repacked.h5'],\n",
       " ['nasa-cryo-scratch/h5cloud/h5repack/ATL03_20181120182818_08110112_006_02_repacked.h5',\n",
       "  'nasa-cryo-scratch/h5cloud/h5repack/ATL03_20190219140808_08110212_006_02_repacked.h5',\n",
       "  'nasa-cryo-scratch/h5cloud/h5repack/ATL03_20200217204710_08110612_006_01_repacked.h5',\n",
       "  'nasa-cryo-scratch/h5cloud/h5repack/ATL03_20211114142614_08111312_006_01_repacked.h5',\n",
       "  'nasa-cryo-scratch/h5cloud/h5repack/ATL03_20230211164520_08111812_006_01_repacked.h5'])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Create a list of s3 objects\n",
    "s3 = s3fs.S3FileSystem(anon=False)\n",
    "\n",
    "# This generates a list of strings with filenames\n",
    "s3path = f's3://{bucket}/{directory}*'\n",
    "remote_files_no_protocol = s3.glob(s3path)\n",
    "remote_files = [f's3://{path}' for path in remote_files_no_protocol]\n",
    "remote_files, remote_files_no_protocol"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e85dc339-f339-482c-8f29-70b728e67a11",
   "metadata": {
    "tags": [],
    "user_expressions": []
   },
   "source": [
    "### Option 1a: Read the group with `h5py`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "4c329e67-3553-434a-b8d0-acc29b7c58b9",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 5/5 [03:07<00:00, 37.52s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 30.1 s, sys: 7.16 s, total: 37.2 s\n",
      "Wall time: 3min 7s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "final_h5py_array = []\n",
    "for file in tqdm.tqdm(remote_files):\n",
    "    with h5py.File(s3.open(file, 'rb')) as f:\n",
    "        data = f[f'{group}/{variable}'][:]\n",
    "        # Need to test if using concatenate is faster\n",
    "        final_h5py_array = np.insert(final_h5py_array, len(final_h5py_array), data, axis=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "901b9476",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "298271231"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(final_h5py_array)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e353e811-35e6-4a51-9773-712c280f959c",
   "metadata": {
    "user_expressions": []
   },
   "source": [
    "### Option 1b: Read the group with gedi subsetter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "6625b8a9-0528-47f4-950b-b43901637f37",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 5/5 [02:15<00:00, 27.02s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 39.3 s, sys: 13.5 s, total: 52.8 s\n",
      "Wall time: 2min 16s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "dataframes = []\n",
    "for file in tqdm.tqdm(remote_files):\n",
    "    with h5py.File(name=s3.open(file, 'rb')) as h5:\n",
    "        df = H5DataFrame(h5[f\"{group[1:]}\"])\n",
    "        dataframes.append(df[variable])\n",
    "final_dataframe: pd.Series = pd.concat(objs=dataframes, axis=\"index\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "18fb35db-c540-4833-9c3c-168c153ef425",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "298271231"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(final_dataframe)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f74714b-5192-49c7-9a65-7f738de90c68",
   "metadata": {
    "tags": [],
    "user_expressions": []
   },
   "source": [
    "### Option 2: Read the group with `xarray`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "ea37b50c-500c-40dd-a982-0e0a878ab7eb",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 2min 22s, sys: 38.5 s, total: 3min 1s\n",
      "Wall time: 7min 59s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "s3_fileset = [s3.open(file) for file in remote_files]\n",
    "xrds = xr.open_mfdataset(s3_fileset, group=group, combine='by_coords', engine='h5netcdf')\n",
    "final_xr_array = xrds['h_ph']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "5df0e6d2-e9bc-4381-a7a5-dfbecc956fd1",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "298271231"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(final_xr_array)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e05d15a-1e58-4c13-baea-27702074ddce",
   "metadata": {
    "tags": [],
    "user_expressions": []
   },
   "source": [
    "### Option 3: Read the group with `h5coro`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "90c7790c",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 5/5 [01:12<00:00, 14.50s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 11.4 s, sys: 3.94 s, total: 15.3 s\n",
      "Wall time: 1min 12s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "final_h5coro_array = []\n",
    "for file in tqdm.tqdm(remote_files_no_protocol):\n",
    "    h5obj = h5coro.H5Coro(file, s3driver.S3Driver)\n",
    "    output = h5obj.readDatasets(datasets=[f'{group}/{variable}'], block=True)\n",
    "    data = h5obj[f'{group}/{variable}'].values\n",
    "    final_h5coro_array = np.insert(final_h5coro_array, len(final_h5coro_array), data, axis=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "45076b16-c481-4708-91e3-7710199d64ec",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "298271231"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(final_h5coro_array)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5943e67-625a-44fe-a4b4-d43dc467a86f",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "f93ecc88-ad98-443a-9d7c-3ad8b318231e",
   "metadata": {
    "tags": [],
    "user_expressions": []
   },
   "source": [
    "# Workflow 2 - Spatially Subset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0e9f3d6-d13d-45a9-92b2-b3de098a7169",
   "metadata": {
    "user_expressions": []
   },
   "source": [
    "## Option 1: Spatially subset with `h5py`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ac0d696-51ea-48bd-a082-60255f3a3e94",
   "metadata": {
    "user_expressions": []
   },
   "source": [
    "## Option 2: Spatially subset with `xarray`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9509f24-1d80-4798-8b86-d6ff94c857c6",
   "metadata": {
    "user_expressions": []
   },
   "source": [
    "## Option 3: Spatially subset with `h5coro`"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}