{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# NDPointIndex\n",
    "\n",
    "This uses one of the new indexes xarray is providing. See [xarray-indexes](https://xarray-indexes.readthedocs.io/) for more in-depth details on that.\n",
    "\n",
    "Needs latest xarray '2025.7.1'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cmweather  # noqa\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "from open_radar_data import DATASETS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "xr.__version__"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {},
   "source": [
    "# Step-by-step guide"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": [
    "## Load radar data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "fname = DATASETS.fetch(\"DWD-Vol-2_99999_20180601054047_00.h5\")\n",
    "ds = xr.open_dataset(fname, engine=\"gamic\", group=\"sweep_9\")\n",
    "display(ds)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "## Georeference\n",
    "\n",
    "Add x,y,z - 2D-coordinates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = ds.xradar.georeference()\n",
    "display(ds)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "metadata": {},
   "source": [
    "## Add NDPointIndex\n",
    "\n",
    "This uses scipy.KDTree under the hood. See also Indexes-Section in the html-repr below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = ds.set_xindex((\"x\", \"y\"), xr.indexes.NDPointIndex)\n",
    "display(ds)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10",
   "metadata": {},
   "source": [
    "## Plot\n",
    "\n",
    "This works as usual with the new NDPointIndex."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds.DBZH.plot(\n",
    "    x=\"x\", y=\"y\", xlim=(-10e3, 10e3), ylim=(-10e3, 10e3), cmap=\"HomeyerRainbow\", vmin=0\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12",
   "metadata": {},
   "source": [
    "## Nearest neighbour interpolation with NDPointIndex\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13",
   "metadata": {},
   "source": [
    "### Create 1D DataArrays for x and y selection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "y = xr.DataArray(np.arange(-100e3, 100e3, 500), dims=\"y\", name=\"y\", attrs=ds.y.attrs)\n",
    "x = xr.DataArray(np.arange(-100e3, 100e3, 500), dims=\"x\", name=\"x\", attrs=ds.x.attrs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15",
   "metadata": {},
   "source": [
    "### Select with above 1D DataArrays"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "actual = ds.sel(y=y, x=x, method=\"nearest\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17",
   "metadata": {},
   "source": [
    "### Assign the 1D DataArrays"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [],
   "source": [
    "actual = actual.assign(x=x, y=y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {},
   "outputs": [],
   "source": [
    "display(actual)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20",
   "metadata": {},
   "source": [
    "## Plot cartesian representation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21",
   "metadata": {},
   "outputs": [],
   "source": [
    "actual.DBZH.plot(xlim=(-10e3, 10e3), ylim=(-10e3, 10e3), cmap=\"HomeyerRainbow\", vmin=0)"
   ]
  }
 ],
 "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}