{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Map Over Sweeps\n",
    "\n",
    "## map_over_sweeps Accessor and Decorator"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1",
   "metadata": {},
   "source": [
    "Have you ever wondered how to efficiently apply operations to a full volume of radar sweeps, rather than processing each sweep individually? \n",
    "\n",
    "Xradar has the solution: the `map_over_sweeps` accessor. In this notebook, we’ll explore how you can leverage Xradar’s powerful `map_over_sweeps` functionality to perform volume-level operations on radar data.\n",
    "\n",
    "In radar data analysis, it's common to work with multiple sweeps in a radar volume. Xradar allows you to apply custom functions across the entire dataset with ease, making complex operations, such as filtering reflectivity or calculating rain rate, efficient and scalable.\n",
    "\n",
    "Here's what you'll learn in this notebook:\n",
    "- How to load and inspect radar data using Xradar.\n",
    "- How to apply functions to process all radar sweeps in one go using both conventional and decorator-based methods.\n",
    "- How to visualize radar variables like reflectivity and rain rate before and after processing.\n",
    "\n",
    "Let’s get into it!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import gridspec\n",
    "from open_radar_data import DATASETS\n",
    "\n",
    "import xradar as xd"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": [
    "## Load Read/Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fetch the sample radar file\n",
    "filename = DATASETS.fetch(\"sample_sgp_data.nc\")\n",
    "\n",
    "# Open the radar file into a DataTree object\n",
    "dtree = xd.io.open_cfradial1_datatree(filename)\n",
    "dtree = dtree.xradar.georeference()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "## Exploring Data Variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "display(dtree)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "for var in dtree[\"sweep_0\"].ds:\n",
    "    print(var)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9",
   "metadata": {},
   "source": [
    "## Map Over Sweeps"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10",
   "metadata": {},
   "source": [
    "We define custom functions like `filter_radar` that filters radar reflectivity based on certain conditions, and `calculate_rain_rate` which is self explanatory, and use `map_over_sweeps` accessor to implement these."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11",
   "metadata": {},
   "source": [
    "## Example #1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "# It is just a demonstration, you can ignore the logic\n",
    "def filter_radar(ds):\n",
    "    ds = ds.assign(\n",
    "        DBZH_Filtered=ds.where(\n",
    "            (ds[\"corrected_reflectivity_horizontal\"] > 10)\n",
    "            & (ds[\"corrected_reflectivity_horizontal\"] < 70)\n",
    "            & (ds[\"copol_coeff\"] > 0.85)\n",
    "        )[\"corrected_reflectivity_horizontal\"]\n",
    "    )\n",
    "    return ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Apply the function across all sweeps\n",
    "dtree = dtree.xradar.map_over_sweeps(filter_radar)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14",
   "metadata": {},
   "source": [
    "## Comparison"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15",
   "metadata": {},
   "source": [
    "Now, let's compare the unfiltered and filtered reflectivity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set a larger figure size with a wider aspect ratio for readability\n",
    "fig, ax = plt.subplots(figsize=(10, 4))\n",
    "\n",
    "# Plot the unfiltered and filtered reflectivity with clear distinctions\n",
    "ax.plot(\n",
    "    dtree[\"sweep_0\"][\"range\"],\n",
    "    dtree[\"sweep_0\"][\"corrected_reflectivity_horizontal\"].sel(\n",
    "        azimuth=100, method=\"nearest\"\n",
    "    ),\n",
    "    alpha=0.7,\n",
    "    lw=2,\n",
    "    linestyle=\"-\",\n",
    "    color=\"m\",\n",
    "    label=\"Unfiltered\",\n",
    ")\n",
    "\n",
    "ax.plot(\n",
    "    dtree[\"sweep_0\"][\"range\"],\n",
    "    dtree[\"sweep_0\"][\"DBZH_Filtered\"].sel(azimuth=100, method=\"nearest\"),\n",
    "    alpha=0.7,\n",
    "    lw=2,\n",
    "    linestyle=\"--\",\n",
    "    color=\"green\",\n",
    "    label=\"Filtered\",\n",
    ")\n",
    "\n",
    "ax.set_xlim(1000, 40_000)\n",
    "ax.set_ylim(0, 50)\n",
    "\n",
    "# Set title and labels with enhanced font sizes\n",
    "ax.set_title(\n",
    "    \"Compare Unfiltered and \\\n",
    "Filtered Reflectivity Variables (Azimuth = 100°)\",\n",
    "    fontsize=16,\n",
    "    pad=15,\n",
    ")\n",
    "ax.set_ylabel(\"Reflectivity [dBZ]\", fontsize=14)\n",
    "ax.set_xlabel(\"Range [m]\", fontsize=14)\n",
    "\n",
    "# Add minor ticks and a grid for both major and minor ticks\n",
    "ax.minorticks_on()\n",
    "ax.grid(True, which=\"both\", linestyle=\"--\", linewidth=0.5)\n",
    "\n",
    "# Adjust legend to avoid overlapping with the plot, and make it clear\n",
    "ax.legend(loc=\"upper right\", fontsize=12, frameon=True)\n",
    "\n",
    "# Apply a tight layout to avoid label/title overlap\n",
    "plt.tight_layout()\n",
    "\n",
    "# Show the plot\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17",
   "metadata": {},
   "source": [
    "## Example #2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a function to calculate rain rate from reflectivity\n",
    "def calculate_rain_rate(ds, ref_field=\"DBZH\"):\n",
    "    def _rain_rate(dbz, a=200.0, b=1.6):\n",
    "        Z = 10.0 ** (dbz / 10.0)\n",
    "        return (Z / a) ** (1.0 / b)\n",
    "\n",
    "    ds = ds.assign(RAIN_RATE=_rain_rate(ds[ref_field]))\n",
    "    ds[\"RAIN_RATE\"].attrs = {\"units\": \"mm/h\", \"long_name\": \"Rain Rate\"}\n",
    "    return ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Apply the function across all sweeps\n",
    "dtree = dtree.xradar.map_over_sweeps(calculate_rain_rate, ref_field=\"DBZH_Filtered\")\n",
    "dtree"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20",
   "metadata": {},
   "source": [
    "If you expand the `dtree` groups and sweeps, you'll notice that the `RAIN_RATE` variable has been successfully added to the DataTree. This confirms that the function was applied across all sweeps, incorporating the calculated rain rate into the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a figure with a GridSpec layout\n",
    "fig = plt.figure(figsize=(8, 8))\n",
    "gs = gridspec.GridSpec(2, 2, wspace=0.05, hspace=0.1)  # Reduced spacing between plots\n",
    "\n",
    "# Define common plotting settings\n",
    "vmin = 0\n",
    "vmax = 50\n",
    "cmap = \"viridis\"\n",
    "\n",
    "# Create the subplots without shrinking due to colorbar\n",
    "axes = [fig.add_subplot(gs[i // 2, i % 2]) for i in range(4)]\n",
    "\n",
    "# Plot each sweep on the respective subplot\n",
    "for i, ax in enumerate(axes):\n",
    "    sweep = f\"sweep_{i}\"\n",
    "    im = dtree[sweep][\"RAIN_RATE\"].plot(\n",
    "        x=\"x\",\n",
    "        y=\"y\",\n",
    "        vmin=vmin,\n",
    "        vmax=vmax,\n",
    "        cmap=cmap,\n",
    "        ax=ax,\n",
    "        add_colorbar=False,  # Disable individual colorbars\n",
    "    )\n",
    "    ax.set_title(f\"Rain Rate for {sweep.replace('_', ' ').capitalize()}\", fontsize=14)\n",
    "\n",
    "    # Turn off ticks and labels for inner subplots\n",
    "    if i in [0, 1]:  # Top row\n",
    "        ax.set_xticklabels([])\n",
    "        ax.set_xlabel(\"\")\n",
    "    if i in [1, 3]:  # Right column\n",
    "        ax.set_yticklabels([])\n",
    "        ax.set_ylabel(\"\")\n",
    "\n",
    "# Create a single shared colorbar outside the subplots\n",
    "# Adjust [left, bottom, width, height] for colorbar position\n",
    "cbar_ax = fig.add_axes([0.92, 0.25, 0.02, 0.5])\n",
    "fig.colorbar(im, cax=cbar_ax, label=\"Rain Rate (mm/hr)\")\n",
    "\n",
    "# Set a main title for the entire figure\n",
    "fig.suptitle(\"Rain Rate Across Different Sweeps\", fontsize=16, fontweight=\"bold\")\n",
    "\n",
    "# Show the plot\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22",
   "metadata": {},
   "source": [
    "## With decorator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23",
   "metadata": {},
   "outputs": [],
   "source": [
    "@xd.map_over_sweeps\n",
    "def calculate_rain_rate2(ds, ref_field=\"DBZH\"):\n",
    "    def _rain_rate(dbz, a=200.0, b=1.6):\n",
    "        Z = 10.0 ** (dbz / 10.0)\n",
    "        return (Z / a) ** (1.0 / b)\n",
    "\n",
    "    ds = ds.assign(RAIN_RATE2=_rain_rate(ds[ref_field]))\n",
    "    ds.RAIN_RATE2.attrs = {\"units\": \"mm/h\", \"long_name\": \"Rain Rate\"}\n",
    "    return ds\n",
    "\n",
    "\n",
    "# invocation via decorator + pipe\n",
    "dtree3 = dtree.pipe(calculate_rain_rate2, ref_field=\"DBZH_Filtered\")\n",
    "display(dtree3[\"sweep_0\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24",
   "metadata": {},
   "source": [
    "## Conclusion"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25",
   "metadata": {},
   "source": [
    "In this notebook, we demonstrated how to handle and process radar data efficiently using Xradar’s `map_over_sweeps` accessor. By applying operations across an entire volume of radar sweeps, you can streamline your workflow and avoid the need to manually process each sweep.\n",
    "\n",
    "We explored two key use cases:\n",
    "- **Filtering Reflectivity**: We applied a custom filtering function across all sweeps in the radar dataset, allowing us to isolate meaningful reflectivity values based on specific criteria.\n",
    "-  **Calculating Rain Rate**: Using the radar reflectivity data, we calculated the rain rate for each sweep, demonstrating how to perform scientific computations across multiple sweeps with minimal effort.\n",
    "\n",
    "The `map_over_sweeps` functionality in Xradar opens the door to performing various radar data processing tasks efficiently. Whether it's filtering, calculating derived quantities like rain rate, or applying more complex algorithms, Xradar simplifies working with radar volumes, making it easier to scale your analysis."
   ]
  }
 ],
 "metadata": {
  "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.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}