{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "fe40a09b",
   "metadata": {},
   "source": [
    "# Volume scan from  multiple sweeps using xradar\n",
    "This example shows how to create a volume scan from multiple sweep files stored on AWS. The volume scan structure is based on [tree-like](https://xarray-datatree.readthedocs.io/en/latest/generated/datatree.DataTree.html) hierarchical collection of xarray objects "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba037495",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9d2cc01",
   "metadata": {},
   "outputs": [],
   "source": [
    "import fsspec\n",
    "import xradar as xd\n",
    "import numpy as np\n",
    "import cartopy.crs as ccrs\n",
    "import cartopy.feature as cfeature\n",
    "import matplotlib.pyplot as plt\n",
    "import cmweather\n",
    "import warnings\n",
    "from pandas import to_datetime\n",
    "from datetime import datetime\n",
    "from matplotlib import pyplot\n",
    "from datatree import DataTree, open_datatree\n",
    "from matplotlib.animation import FuncAnimation\n",
    "from IPython.display import HTML\n",
    "\n",
    "warnings.simplefilter(\"ignore\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59dbfdb5",
   "metadata": {},
   "source": [
    "## Access radar data from the Colombian radar network on AWS\n",
    "Access data from the IDEAM bucket on AWS. Detailed information can be found [here](https://openradar-docs--102.org.readthedocs.build/projects/xradar/en/102/notebooks/Read-plot-Sigmet-data-from-AWS.html) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37b38b59",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_query(date, radar_site):\n",
    "    \"\"\"\n",
    "    Creates a string for quering the IDEAM radar files stored in AWS bucket\n",
    "    :param date: date to be queried. e.g datetime(2021, 10, 3, 12). Datetime python object\n",
    "    :param radar_site: radar site e.g. Guaviare\n",
    "    :return: string with a IDEAM radar bucket format\n",
    "    \"\"\"\n",
    "    if (date.hour != 0) and (date.hour != 0):\n",
    "        return f\"l2_data/{date:%Y}/{date:%m}/{date:%d}/{radar_site}/{radar_site[:3].upper()}{date:%y%m%d%H}\"\n",
    "    elif (date.hour != 0) and (date.hour == 0):\n",
    "        return f\"l2_data/{date:%Y}/{date:%m}/{date:%d}/{radar_site}/{radar_site[:3].upper()}{date:%y%m%d}\"\n",
    "    else:\n",
    "        return f\"l2_data/{date:%Y}/{date:%m}/{date:%d}/{radar_site}/{radar_site[:3].upper()}{date:%y%m%d}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee045227",
   "metadata": {},
   "outputs": [],
   "source": [
    "date_query = datetime(2023, 4, 7, 3)\n",
    "radar_name = \"Barrancabermeja\"\n",
    "query = create_query(date=date_query, radar_site=radar_name)\n",
    "str_bucket = \"s3://s3-radaresideam/\"\n",
    "fs = fsspec.filesystem(\"s3\", anon=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4f2b780",
   "metadata": {},
   "outputs": [],
   "source": [
    "query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2dcf1f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "radar_files = sorted(fs.glob(f\"{str_bucket}{query}*\"))\n",
    "radar_files[:4]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3993934-42d3-499d-b044-21a8cb8a608f",
   "metadata": {},
   "source": [
    "## Let's check the elevation at each file using `xradar.datatree` module\n",
    "\n",
    "IDEAM radar network operates with a volume scan every five minutes. Each volume scan has four different tasks \n",
    "* *SURVP* \"super-resolution\" sweep at the lowest elevation angle, usually 0.5 deg,  with 720 degrees in azimuth (every 0.5 deg)\n",
    "* *PRECA* task with 1.5, 2.4, 3.0, and 5.0 elevation angles and shorter range than *SURVP*\n",
    "* *PRECB* task with 6.4 and 8.0 elevation angles and a shorter range than the previous task\n",
    "* *PRECC* task with 10.0, 12.5, and 15.0 with a shorter range than all the previous tasks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3cb9fb1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# List of first four task files\n",
    "task_files = [\n",
    "    fsspec.open_local(\n",
    "        f\"simplecache::s3://{i}\", s3={\"anon\": True}, filecache={\"cache_storage\": \".\"}\n",
    "    )\n",
    "    for i in radar_files[:4]\n",
    "]\n",
    "# list of xradar datatrees\n",
    "ls_dt = [xd.io.open_iris_datatree(i).xradar.georeference() for i in task_files]\n",
    "\n",
    "# sweeps and elevations within each task\n",
    "for i in ls_dt:\n",
    "    sweeps = list(i.children.keys())\n",
    "    print(f\"task sweeps: {sweeps}\")\n",
    "    for j in sweeps:\n",
    "        if j.startswith(\"sweep\"):\n",
    "            print(\n",
    "                f\"{j}: {i[j].sweep_fixed_angle.values: .1f} [deg], {i[j].range.values[-1] / 1e3:.1f} [km]\"\n",
    "            )\n",
    "    print(\"----------------------------------------------------------------\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a5e5481",
   "metadata": {},
   "source": [
    "## Create a single-volume scan\n",
    "Let's use the first four files, tasks *SURVP*, *PRECA*, *PRECB*, *PRECC*, to create a single volume scan using each task as a datatree. The new volume scan is a tree-like hierarchical object with all four tasks as children."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c906964",
   "metadata": {},
   "outputs": [],
   "source": [
    "vcp_dt = DataTree(\n",
    "    name=\"root\",\n",
    "    children=dict(SURVP=ls_dt[0], PRECA=ls_dt[1], PRECB=ls_dt[2], PRECC=ls_dt[3]),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f78bb9c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "vcp_dt.groups"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d026f7ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Size of data in tree = {vcp_dt.nbytes / 1e6 :.2f} MB\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "391c40cf",
   "metadata": {},
   "source": [
    "## PPI plot from the Datatree object\n",
    "\n",
    "Now that we have a tree-like hierarchical volume scan object. We can access data at each scan/sweep using dot method `vcp_dt.SURVP` or dictionary-key method `vcp_dt['PRECB']`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41260d70",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, (ax, ax1) = plt.subplots(1, 2, figsize=(10, 4))\n",
    "# dot method\n",
    "vcp_dt.SURVP.sweep_0.DBZH.plot(\n",
    "    x=\"x\", y=\"y\", cmap=\"ChaseSpectral\", vmin=-10, vmax=50, ax=ax\n",
    ")\n",
    "\n",
    "ax.set_title(\n",
    "    f\"SURPV sweep_0 ({vcp_dt.SURVP.sweep_0.sweep_fixed_angle.values: .1f} [deg])\"\n",
    ")\n",
    "m2km = lambda x, _: f\"{x/1000:g}\"\n",
    "ax.xaxis.set_major_formatter(m2km)\n",
    "ax.yaxis.set_major_formatter(m2km)\n",
    "ax.set_ylabel(\"$North - South \\ distance \\ [km]$\")\n",
    "ax.set_xlabel(\"$East - West \\ distance \\ [km]$\")\n",
    "\n",
    "# Dictionary-key method\n",
    "vcp_dt[\"PRECB\"][\"sweep_0\"].DBZH.plot(\n",
    "    x=\"x\", y=\"y\", cmap=\"ChaseSpectral\", vmin=-10, vmax=50, ax=ax1\n",
    ")\n",
    "\n",
    "ax1.set_title(\n",
    "    f\"PRECB sweep_0 ({vcp_dt.PRECB.sweep_0.sweep_fixed_angle.values: .1f} [deg])\"\n",
    ")\n",
    "m2km = lambda x, _: f\"{x/1000:g}\"\n",
    "ax1.xaxis.set_major_formatter(m2km)\n",
    "ax1.yaxis.set_major_formatter(m2km)\n",
    "ax1.set_xlim(ax.get_xlim())\n",
    "ax1.set_ylim(ax.get_ylim())\n",
    "ax1.set_ylabel(\"$North - South \\ distance \\ [km]$\")\n",
    "ax1.set_xlabel(\"$East - West \\ distance \\ [km]$\")\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67d9d6d0",
   "metadata": {},
   "source": [
    "## Multiple volumes scan into one `datatree` object\n",
    "\n",
    "Similarly, we can create a tree-like hierarchical object with multiple volume scans."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ed5a84b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def data_accessor(file):\n",
    "    \"\"\"\n",
    "    Open AWS S3 file(s), which can be resolved locally by file caching\n",
    "    \"\"\"\n",
    "    return fsspec.open_local(\n",
    "        f\"simplecache::s3://{file}\",\n",
    "        s3={\"anon\": True},\n",
    "        filecache={\"cache_storage\": \"./tmp/\"},\n",
    "    )\n",
    "\n",
    "\n",
    "def create_vcp(ls_dt):\n",
    "    \"\"\"\n",
    "    Creates a tree-like object for each volume scan\n",
    "    \"\"\"\n",
    "    return DataTree(\n",
    "        name=\"root\",\n",
    "        children=dict(SURVP=ls_dt[0], PRECA=ls_dt[1], PRECB=ls_dt[2], PRECC=ls_dt[3]),\n",
    "    )\n",
    "\n",
    "\n",
    "def mult_vcp(radar_files):\n",
    "    \"\"\"\n",
    "    Creates a tree-like object for multiple volumes scan every 4th file in the bucket\n",
    "    \"\"\"\n",
    "    ls_files = [radar_files[i : i + 4] for i in range(len(radar_files)) if i % 4 == 0]\n",
    "    ls_sigmet = [\n",
    "        [xd.io.open_iris_datatree(data_accessor(i)).xradar.georeference() for i in j]\n",
    "        for j in ls_files\n",
    "    ]\n",
    "    ls_dt = [create_vcp(i) for i in ls_sigmet]\n",
    "    return DataTree.from_dict({f\"vcp_{idx}\": i for idx, i in enumerate(ls_dt)})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40c830a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# let's test it using the first 24 files in the bucket. We can include more files for visualization. e.g. radar_files[:96]\n",
    "vcps_dt = mult_vcp(radar_files[:24])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60c82dcc-46f5-4be8-b075-0079f9fc33bc",
   "metadata": {},
   "source": [
    "Now we have 6 vcps in one tree-like hierarchical object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21b82847",
   "metadata": {},
   "outputs": [],
   "source": [
    "list(vcps_dt.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f28f154",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Size of data in tree = {vcps_dt.nbytes / 1e9 :.2f} GB\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e604675",
   "metadata": {},
   "source": [
    "### PPI animation using the lowest elevation angle\n",
    "\n",
    "We can create an animation using the `FuncAnimation` module from `matplotlib` package"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8646efdb",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(subplot_kw={\"projection\": ccrs.PlateCarree()})\n",
    "proj_crs = xd.georeference.get_crs(vcps_dt.vcp_1.SURVP)\n",
    "cart_crs = ccrs.Projection(proj_crs)\n",
    "sc = vcps_dt.vcp_1.SURVP.sweep_0.DBZH.plot.pcolormesh(\n",
    "    x=\"x\",\n",
    "    y=\"y\",\n",
    "    vmin=-10,\n",
    "    vmax=50,\n",
    "    cmap=\"ChaseSpectral\",\n",
    "    edgecolors=\"face\",\n",
    "    transform=cart_crs,\n",
    "    ax=ax,\n",
    ")\n",
    "\n",
    "title = f\"SURVP - {vcps_dt.vcp_1.SURVP.sweep_0.sweep_fixed_angle.values: .1f} [deg]\"\n",
    "ax.set_title(title)\n",
    "gl = ax.gridlines(\n",
    "    crs=ccrs.PlateCarree(),\n",
    "    draw_labels=True,\n",
    "    linewidth=1,\n",
    "    color=\"gray\",\n",
    "    alpha=0.3,\n",
    "    linestyle=\"--\",\n",
    ")\n",
    "plt.gca().xaxis.set_major_locator(plt.NullLocator())\n",
    "gl.top_labels = False\n",
    "gl.right_labels = False\n",
    "ax.coastlines()\n",
    "\n",
    "\n",
    "def update_plot(vcp):\n",
    "    sc.set_array(vcps_dt[vcp].SURVP.sweep_0.DBZH.values.ravel())\n",
    "\n",
    "\n",
    "ani = FuncAnimation(fig, update_plot, frames=list(vcps_dt.keys()), interval=150)\n",
    "plt.close()\n",
    "HTML(ani.to_html5_video())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bddf03ce",
   "metadata": {},
   "source": [
    "## Bonus!!\n",
    "### Analysis-ready data, cloud-optimized (ARCO) format \n",
    "\n",
    "Tree-like hierarchical data can be stored using ARCO format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f44f72a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "zarr_store = \"./multiple_vcp_test.zarr\"\n",
    "_ = vcps_dt.to_zarr(zarr_store)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65183026",
   "metadata": {},
   "source": [
    "ARCO format can be read by using `open_datatree` module"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76f182c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "vcps_back = open_datatree(zarr_store, engine=\"zarr\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "249c1079",
   "metadata": {},
   "outputs": [],
   "source": [
    "display(vcps_back)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96b5c030-3087-402b-b33b-6daf6d6e6214",
   "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.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}