notebooks/Find_closest_neighbors_kdtree_balltree.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0332940b",
   "metadata": {},
   "source": [
    "# Find closest neighbors (KDTree and BallTree)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d5a7877",
   "metadata": {},
   "outputs": [],
   "source": [
    "# solve issue with autocomplete\n",
    "%config Completer.use_jedi = False\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e44900e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from geotree import gtree\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "524f47c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# instantiate gtree\n",
    "mytree = gtree()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc23ae73",
   "metadata": {},
   "source": [
    "## Define the first set of points or `base`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a011262",
   "metadata": {},
   "outputs": [],
   "source": [
    "npoints = 200\n",
    "lons = np.random.randint(-180, 180, npoints)\n",
    "lats = np.random.randint(-90, 90, npoints)\n",
    "depths = np.zeros(npoints)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55161276",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add lons/lats/depths of the first set of points\n",
    "mytree.add_lonlatdep(lons=lons, \n",
    "                     lats=lats, \n",
    "                     depths=depths)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bdba2d9b",
   "metadata": {},
   "source": [
    "## Define queries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47186edc",
   "metadata": {},
   "outputs": [],
   "source": [
    "q_npoints = 10\n",
    "q_lons = np.random.randint(-150, 150, q_npoints)\n",
    "q_lats = np.random.randint(-70, 70, q_npoints)\n",
    "q_depths = np.zeros(q_npoints)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cea3e198",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add lons/lats/depths of queries\n",
    "mytree.add_lonlatdep_query(lons=q_lons, \n",
    "                           lats=q_lats, \n",
    "                           depths=q_depths)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b298ce66",
   "metadata": {},
   "source": [
    "## Find neighbors, KDTree"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "edffdf2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create KDTree (kdt)\n",
    "mytree.create_kdt()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "733da042",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Choose the desired number of neighbors (and upper bound for distance, if needed):\n",
    "mytree.query_kdt(num_neighs=3, distance_upper=np.inf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27fe5beb",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Now, for each query, distances to the closest `base` neighbors and their indices are stored in (row-wise):\n",
    "\n",
    "# distances to the closest `base` neighbors \n",
    "mytree.dists2query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b158189d",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# indices of the closest `base` neighbors\n",
    "mytree.indxs2query"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d46c3e2",
   "metadata": {},
   "source": [
    "### Plot the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd354cfd",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(15, 7))\n",
    "\n",
    "colormap = plt.cm.Spectral\n",
    "list_colors = [colormap(i) for i in np.linspace(0, 1,len(q_lons))]\n",
    "\n",
    "# plot queries\n",
    "plt.scatter(q_lons, q_lats, \n",
    "            c=\"b\", \n",
    "            marker=\"x\", \n",
    "            label=\"queries\")\n",
    "\n",
    "# --- plot a line between one query and its neighbours\n",
    "# neighboring base lons/lats\n",
    "b_lons = lons[mytree.indxs2query]\n",
    "b_lats = lats[mytree.indxs2query]\n",
    "for i in range(len(q_lons)):\n",
    "    for j in range(len(mytree.indxs2query[i])):\n",
    "        plt.plot((q_lons[i], b_lons[i, j]), \n",
    "                 (q_lats[i], b_lats[i, j]), \n",
    "                 c=list_colors[i])\n",
    "\n",
    "# plot base points\n",
    "plt.scatter(lons, lats,\n",
    "            c=\"k\", \n",
    "            marker=\"o\",\n",
    "            label=\"base\",\n",
    "            zorder=100)\n",
    "\n",
    "plt.legend(bbox_to_anchor=(0., 1.01, 1., .05), \n",
    "           loc=\"right\", ncol=2, \n",
    "           fontsize=16,\n",
    "           borderaxespad=0.)\n",
    "\n",
    "plt.grid()\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "245b3d45",
   "metadata": {},
   "source": [
    "## Project results using `cartopy`\n",
    "\n",
    "`cartopy` needs to be installed, see: https://scitools.org.uk/cartopy/docs/latest/index.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6774c16",
   "metadata": {},
   "outputs": [],
   "source": [
    "from cartopy import crs\n",
    "\n",
    "fig = plt.figure(figsize=(10, 7))\n",
    "\n",
    "ax = fig.add_subplot(1, 1, 1, \n",
    "                     projection=crs.InterruptedGoodeHomolosine())\n",
    "ax.coastlines(color=\"black\")\n",
    "\n",
    "# plot queries\n",
    "plt.scatter(\n",
    "    mytree.lons_q, \n",
    "    mytree.lats_q, \n",
    "    transform=crs.PlateCarree(), \n",
    "    c=\"b\",\n",
    "    marker=\"x\")\n",
    "\n",
    "# --- plot a line between one query and its neighbours\n",
    "# neighboring base lons/lats\n",
    "b_lons = lons[mytree.indxs2query]\n",
    "b_lats = lats[mytree.indxs2query]\n",
    "for i in range(len(q_lons)):\n",
    "    for j in range(len(mytree.indxs2query[i])):\n",
    "        plt.plot((q_lons[i], b_lons[i, j]), \n",
    "                 (q_lats[i], b_lats[i, j]), \n",
    "                 transform=crs.PlateCarree(), \n",
    "                 c = list_colors[i])\n",
    "\n",
    "# plot base points\n",
    "plt.scatter(\n",
    "    mytree.lons, \n",
    "    mytree.lats, \n",
    "    transform=crs.PlateCarree(), \n",
    "    color=\"k\")\n",
    "\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08c0aede",
   "metadata": {},
   "source": [
    "## Find neighbors, Ball tree"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3bd48a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "mytree.create_balltree()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bae1b581",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Choose the desired number of neighbors:\n",
    "mytree.query_balltree(num_neighs=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe6d7772",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Now, for each query, distances to the closest `base` neighbors and their indices are stored in (row-wise):\n",
    "\n",
    "# distances to the closest `base` neighbors \n",
    "mytree.dists2query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db58b7b8",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# indices of the closest `base` neighbors\n",
    "mytree.indxs2query"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}