notebooks/h5py_testing_original_repacked_with_subsetting.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "10a088ec-3542-4f2b-87c1-020968f9d245",
   "metadata": {
    "user_expressions": []
   },
   "source": [
    "# H5PY with Subsetting Benchmarking\n",
    "\n",
    "Note, reading in lat/lon uses the full photon resolution data for simplicity right now, but could probably be improved by an order of magnitude by subsetting at the segment level using ref_photon_lat/lon or ph_index_beg parameters from ATL03."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9dbad514-c967-4411-8eed-02d3866eed9c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import cProfile\n",
    "import numpy as np\n",
    "import s3fs\n",
    "import xarray as xr\n",
    "import h5py\n",
    "import boto3\n",
    "import os\n",
    "import geopandas as gpd\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae4bc834-f0cf-4eb9-bdaa-23df6e77dbbf",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# !mamba install -c conda-forge awscli -y"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00c515de-9300-4ba0-9189-2538cbbf9ac3",
   "metadata": {
    "user_expressions": []
   },
   "source": [
    "## Setting up Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5a3fb09-d3ec-4e4e-b2e9-3b3e3ceedc48",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# take a look of the files in our s3 bucket\n",
    "\n",
    "s3_r = boto3.resource('s3')\n",
    "\n",
    "bucket_name = \"nasa-cryo-scratch\"\n",
    "\n",
    "bucket = s3_r.Bucket(bucket_name)\n",
    "\n",
    "objects = list(bucket.objects.all())\n",
    "\n",
    "for my_bucket_object in objects:\n",
    "    print(my_bucket_object)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df9a571b-7bdd-462e-af67-941903a5d01d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "s3 = boto3.resource('s3')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73b03bb7-e88b-4f09-9fd0-9ef5cf8234f5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Checkout the files\n",
    "!aws s3 ls s3://nasa-cryo-scratch/h5cloud/original/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "611e0a3d-2e70-4e3a-a6f2-45919df0d844",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Set the variables\n",
    "bucket = 'nasa-cryo-scratch'\n",
    "directory = 'h5cloud/original/'\n",
    "\n",
    "# Create a list of s3 objects\n",
    "s3 = s3fs.S3FileSystem(anon=False)\n",
    "\n",
    "# This generates a list of strings with filenames\n",
    "s3path_original = f's3://{bucket}/{directory}*'\n",
    "remote_files_original = [f's3://{path}' for path in s3.glob(s3path_original)]\n",
    "\n",
    "s3path_repacked = f's3://{bucket}/{directory}*'\n",
    "remote_files_repacked = [f's3://{path}' for path in s3.glob(s3path_repacked)]\n",
    "\n",
    "remote_files_original"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a39546be-fbdf-4bab-a7bf-4f535247985c",
   "metadata": {
    "tags": [],
    "user_expressions": []
   },
   "source": [
    "# Workflow 1 - Read Data H5PY"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1328823-4892-4922-9eac-cee2fa0a1be5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "group = '/gt2l/heights'\n",
    "\n",
    "# variables = [\n",
    "#     \"delta_time\", \"dist_ph_across\", \"dist_ph_along\", \"h_ph\", \"lat_ph\", \"lon_ph\", \"pce_mframe_cnt\", \"ph_id_channel\",\"ph_id_count\", \"ph_id_pulse\", \"quality_ph\",\"signal_conf_ph\",\n",
    "# ]\n",
    "\n",
    "# using the full list of variables takes a very long time to load for all files (~ 8 min on hackweek jupyterhub)\n",
    "# for now, using just lat, lon, heights for\n",
    "\n",
    "variables = ['lat_ph', 'lon_ph', 'h_ph']\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30293243-cae7-4a79-be61-d915e35cd0f7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def read_h5py(file, group, variables, verbose=False):\n",
    "    \n",
    "    with h5py.File(s3.open(file, 'rb')) as f:\n",
    "        \n",
    "        \n",
    "        if verbose: print(f'opening {file}...')\n",
    "        \n",
    "        group_data = []\n",
    "        \n",
    "        for variable in variables:\n",
    "            \n",
    "            if verbose: print(f'... reading {variable}')\n",
    "            \n",
    "            data = f[f'{group}/{variable}'][:]\n",
    "            \n",
    "            group_data.append(data)\n",
    "            \n",
    "    return group_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9aca8ca3-a415-4c27-8096-91974c39179b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def test_h5py(files, group='/gt2l/heights', \n",
    "              variables=[\"delta_time\", \"dist_ph_across\", \"dist_ph_along\",\n",
    "                         \"h_ph\", \"lat_ph\", \"lon_ph\", \n",
    "                         \"pce_mframe_cnt\", \"ph_id_channel\",\n",
    "                         \"ph_id_count\", \"ph_id_pulse\", \n",
    "                         \"quality_ph\",\"signal_conf_ph\"],\n",
    "             verbose=False):\n",
    "    \n",
    "    original_data = []\n",
    "\n",
    "    for file in files:\n",
    "        data = read_h5py(file, group, variables, verbose)\n",
    "        original_data.append(data)\n",
    "    \n",
    "    \n",
    "    return original_data\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07dc7109-af83-4d51-8da9-aea42780ba5b",
   "metadata": {
    "tags": [],
    "user_expressions": []
   },
   "source": [
    "# Workflow 2 - Spatially Subset H5PY"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f79b1b3f-11c7-46de-b5c8-46c247b15b1d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# read in the area of interest geojson\n",
    "aoi = gpd.read_file('/home/jovyan/h5cloud/notebooks/antarctic_aoi.geojson', crs='EPSG:4326')\n",
    "bounds = aoi.bounds.values[0] "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aab2258e-ded1-4c57-8ed8-8df52e19de98",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def read_spatial_subset_h5py(file, group, subset_variables, min_lon, max_lon, min_lat, max_lat, verbose=False):\n",
    "    \n",
    "    with h5py.File(s3.open(file, 'rb')) as f:\n",
    "        \n",
    "        if verbose: print(f'opening {file}...')\n",
    "        \n",
    "        group_data = []\n",
    "        \n",
    "        # read in the photon data to use for indexing\n",
    "        # may be sped up significantly by using a segment rate parameter instead\n",
    "        lat = f[f'{group}/lat_ph'][:]\n",
    "        lon = f[f'{group}/lon_ph'][:]\n",
    "        \n",
    "        ph_in_aoi = np.where((lat > min_lat) & (lat < max_lat) \\\n",
    "                             & (lon > min_lon) & (lon < max_lon))[0]\n",
    "\n",
    "        idx_start = ph_in_aoi[0]\n",
    "        idx_end = ph_in_aoi[-1]\n",
    "        \n",
    "        group_data.append(lat[idx_start:idx_end])\n",
    "        group_data.append(lon[idx_start:idx_end])\n",
    "        \n",
    "        for variable in subset_variables:\n",
    "            \n",
    "            if verbose: print(f'... reading {variable}')\n",
    "            \n",
    "            data = f[f'{group}/{variable}'][idx_start:idx_end]\n",
    "            \n",
    "            group_data.append(data)\n",
    "\n",
    "    return group_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7427a63b-35f3-4e7f-9bf7-db9a5bbd5dcf",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def test_h5py_subset(remote_files, bounds, group='/gt2l/heights', \n",
    "              variables=[\"delta_time\", \"dist_ph_across\", \"dist_ph_along\",\n",
    "                         \"h_ph\", \"pce_mframe_cnt\", \"ph_id_channel\",\n",
    "                         \"ph_id_count\", \"ph_id_pulse\", \n",
    "                         \"quality_ph\",\"signal_conf_ph\"], verbose=False):\n",
    "    \n",
    "    subset_data = []\n",
    "    \n",
    "    # specify the lat lon of the bounding box\n",
    "    min_lon = bounds[0]\n",
    "    min_lat = bounds[1]\n",
    "    max_lon = bounds[2]\n",
    "    max_lat = bounds[3]\n",
    "\n",
    "    # Loop through files / read\n",
    "    for file in remote_files:\n",
    "        data = read_spatial_subset_h5py(file, group, \n",
    "                                        variables, min_lon, max_lon, min_lat, max_lat, verbose)\n",
    "        subset_data.append(data)\n",
    "\n",
    "    return subset_data\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8740ce0-8a6f-4630-b860-4494618348b8",
   "metadata": {
    "user_expressions": []
   },
   "source": [
    "# Executing Workflows 1/2 for Original and Repacked H5 Files"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae2385cf-5b6e-4592-84c5-50183f683335",
   "metadata": {
    "user_expressions": []
   },
   "source": [
    "Running the test functions for the lists of original and repacked h5 files, and saving the time results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a358eadb-98b6-4ca0-b9bc-3462cc5293d1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "h5py_original_time = %timeit -n3 -r1 -o test_h5py(remote_files_original, verbose=True)\n",
    "h5py_subset_timeit = %timeit -n3 -r1 -o test_h5py_subset(remote_files_original, bounds, verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7a22bf8-688f-4cba-983a-2cf9207efa7e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "h5py_repacked_timeit = %timeit -n3 -r3 -o test_h5py(remote_files_repacked, verbose=True)\n",
    "h5py_repacked_subset_timeit = %timeit -n3 -r3 -o test_h5py_subset(remote_files_repacked, verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85531225-dba9-4e07-93b2-34caecc9e502",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "plt.plot(h5py_original_time.timings)\n",
    "plt.ylabel('seconds')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b74c974-92e4-4b20-85e9-3149fb11e5af",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}