{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2057f600-1c0f-46d9-a906-c5b1d528f210",
   "metadata": {},
   "source": [
    "# Computing Beta drift with HuracanPy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30ee360e-751e-4b76-ad94-3903a02f07b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import huracanpy\n",
    "\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffff6a10-35ac-4ef0-9fed-b9c02a6d4f0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = huracanpy.load(huracanpy.example_TRACK_file, source=\"TRACK\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39da5c34-0f37-44e2-87c8-e0c3e0b8d73f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1. Compute Radius of Maximum Wind\n",
    "# Distance between the vortex centre and the location of the maximum wind speed\n",
    "# identified by TRACK (feature_9)\n",
    "RMW = huracanpy.calc.distance(\n",
    "    data.lon, data.lat, data.feature_9_lon, data.feature_9_lat\n",
    ")\n",
    "data = data.assign(RMW=RMW)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a2f9e07-ac38-45a4-bd65-8a5cf429c221",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Verify result\n",
    "plt.scatter(data.lat - data.feature_9_lat, data.lon - data.feature_9_lon, c=RMW)\n",
    "plt.colorbar(label=\"Radius of maximum wind in km\")\n",
    "plt.xlabel(\"Latitude difference\")\n",
    "plt.ylabel(\"Longitude difference\")\n",
    "plt.axhline(color=\"k\")\n",
    "plt.axvline(color=\"k\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2ad0597-2641-444e-9ae9-8ca6424164b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 2. Compute beta-drift\n",
    "V_drift, theta_drift = huracanpy.tc.beta_drift(data.lat, data.feature_9, data.RMW)\n",
    "data = data.assign(V_drift=V_drift, theta_drift=theta_drift)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7576b47-6486-4a67-bdc5-29cef096cfb3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Verify result\n",
    "plt.scatter(\n",
    "    data.feature_9,\n",
    "    data.RMW,\n",
    "    c=data.V_drift,\n",
    ")\n",
    "plt.colorbar(label=\"Beta-drift wind in m/s\")\n",
    "plt.xlabel(\"Maximum wind speed in m/s\")\n",
    "plt.ylabel(\"radius of maximum wind in km\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6dc30ab8-fdf2-45f0-bce7-beb27a1d55b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.scatter(\n",
    "    data.feature_9,\n",
    "    data.V_drift,\n",
    "    c=data.theta_drift,\n",
    "    s=data.RMW / 1e3,\n",
    "    edgecolors=\"grey\",\n",
    ")\n",
    "plt.colorbar(label=\"Drift angle in degrees\")\n",
    "plt.xlabel(\"Maximum wind speed in m/s\")\n",
    "plt.ylabel(\"Beta drift wind in m/s\")\n",
    "plt.title(\"Dot size represents the radius of maximum winds\")"
   ]
  }
 ],
 "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.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}