{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Vectorfield"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import hvplot.xarray  # noqa"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`vectorfield` accepts 2d arrays of magnitude and angle on a grid and produces an array of vectors. x and y can be 2d or 1d coordinates. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import xarray as xr\n",
    "import holoviews as hv\n",
    "import cartopy.crs as ccrs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample_data(shape=(20, 30)):\n",
    "    \"\"\"\n",
    "    Return ``(x, y, u, v, crs)`` of some vector data\n",
    "    computed mathematically. The returned crs will be a rotated\n",
    "    pole CRS, meaning that the vectors will be unevenly spaced in\n",
    "    regular PlateCarree space.\n",
    "\n",
    "    \"\"\"\n",
    "    crs = ccrs.RotatedPole(pole_longitude=177.5, pole_latitude=37.5)\n",
    "\n",
    "    x = np.linspace(311.9, 391.1, shape[1])\n",
    "    y = np.linspace(-23.6, 24.8, shape[0])\n",
    "\n",
    "    x2d, y2d = np.meshgrid(x, y)\n",
    "    u = 10 * (2 * np.cos(2 * np.deg2rad(x2d) + 3 * np.deg2rad(y2d + 30)) ** 2)\n",
    "    v = 20 * np.cos(6 * np.deg2rad(x2d))\n",
    "\n",
    "    return x, y, u, v, crs\n",
    "\n",
    "xs, ys, U, V, crs = sample_data()\n",
    "\n",
    "mag = np.sqrt(U**2 + V**2)\n",
    "angle = (np.pi/2.) - np.arctan2(U/mag, V/mag)\n",
    "\n",
    "ds = xr.Dataset({'mag': xr.DataArray(mag, dims=('y', 'x'), coords={'y': ys, 'x': xs}),\n",
    "                 'angle': xr.DataArray(angle, dims=('y', 'x'), coords={'y': ys, 'x': xs})}, \n",
    "                attrs={'crs': crs})\n",
    "ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds.hvplot.vectorfield(x='x', y='y', angle='angle', mag='mag', hover=False).opts(magnitude='mag')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Geographic Data\n",
    "If a dataset has an attr called `crs` which is a cartopy object or a proj4 string, then just by setting the option `geo=True` will use the correct crs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds.hvplot.vectorfield(x='x', y='y', angle='angle', mag='mag',\n",
    "                      hover=False, geo=True, tiles=\"CartoLight\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you set `coastline` or `features` it will keep the original crs and transform the features to the data crs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds.hvplot.vectorfield(x='x', y='y', angle='angle', mag='mag',\n",
    "                      hover=False, geo=True, coastline=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Large Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The visualization of vector fields from large datasets often presents a challenge.  Direct plotting methods can quickly consume excessive memory, leading to crashes or unresponsive applications. To address this issue, we introduce a dynamic downsampling technique that enables interactive exploration of vector fields without sacrificing performance. We first create a `sample_data` that contains 4,000,000 points."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xs, ys, U, V, crs = sample_data(shape=(2000, 2000))\n",
    "\n",
    "mag = np.sqrt(U**2 + V**2)\n",
    "angle = (np.pi/2.) - np.arctan2(U/mag, V/mag)\n",
    "\n",
    "ds = xr.Dataset({'mag': xr.DataArray(mag, dims=('y', 'x'), coords={'y': ys, 'x': xs}),\n",
    "                 'angle': xr.DataArray(angle, dims=('y', 'x'), coords={'y': ys, 'x': xs})}, \n",
    "                attrs={'crs': crs})\n",
    "ds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we just try to call `ds.hvplot.vectorfield` this is probably returning `MemoryError`. The alternative is to dynamically downsample the data based on the visible range. This helps manage memory consumption when dealing with large datasets, especially when plotting vector fields. We are going to use HoloViews to create a view (`hv.DynamicMap`) whose content is dynamically updated based on the data range displayed (tracked by the `hv.streams.RangeXY` stream), [find out more about these concepts](https://holoviews.org/user_guide/Custom_Interactivity.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def downsample_quiver(x_range=None, y_range=None, nmax=10):\n",
    "    \"\"\"\n",
    "    Creates a HoloViews vector field plot from a dataset, dynamically downsampling \n",
    "    data based on the visible range to optimize memory usage.\n",
    "\n",
    "    Args:\n",
    "        x_range (tuple, optional): Range of x values to include. Defaults to None (full range).\n",
    "        y_range (tuple, optional): Range of y values to include. Defaults to None (full range).\n",
    "        nmax (int, optional): Maximum number of points along each axis after coarsening. \n",
    "                              Defaults to 10.\n",
    "\n",
    "    Returns:\n",
    "        HoloViews DynamicMap: A dynamic vector field plot that updates based on the visible range.\n",
    "    \"\"\"\n",
    "\n",
    "    if x_range is None or y_range is None:\n",
    "        # No range provided, downsample the entire dataset for initial display\n",
    "        xs, ys = ds.x.size, ds.y.size        # Get dataset dimensions\n",
    "        ix, iy = xs // nmax, ys // nmax      # Calculate downsampling intervals\n",
    "\n",
    "        ix = max(1, ix)                     # Ensure interval is at least 1\n",
    "        iy = max(1, iy)\n",
    "\n",
    "        sub = ds.coarsen(x=ix, y=iy, side=\"center\", boundary=\"trim\").mean()  # Downsample\n",
    "    else:\n",
    "        # Select data within the specified range\n",
    "        sub = ds.sel(x=slice(*x_range), y=slice(*y_range))\n",
    "\n",
    "        # Downsample the selected data\n",
    "        xs, ys = sub.x.size, sub.y.size\n",
    "        ix, iy = xs // nmax, ys // nmax\n",
    "        ix = max(1, ix)\n",
    "        iy = max(1, iy)\n",
    "        sub = sub.coarsen(x=ix, y=iy, side=\"center\", boundary=\"trim\").mean()\n",
    "\n",
    "    # Create the vector field plot\n",
    "    quiver = sub.hvplot.vectorfield(\n",
    "        x=\"x\",\n",
    "        y=\"y\",\n",
    "        mag=\"mag\",\n",
    "        angle=\"angle\",\n",
    "        hover=False,\n",
    "    ).opts(magnitude=\"mag\")\n",
    "\n",
    "    return quiver"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create interactive plot components\n",
    "range_xy = hv.streams.RangeXY()               # Stream to capture range changes\n",
    "filtered = hv.DynamicMap(downsample_quiver, streams=[range_xy])  # Dynamic plot\n",
    "filtered                                      # Display the plot"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}