bench_lsdb_xmatch_vs_astropy.ipynb

{
 "cells": [
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "# Performance of LSDB's cross-matching algorithm\n",
    "\n",
    "We compare the performance of the default LSDB's cross-matching algorithm, `KdTreeCrossmatch`, with `astropy`'s functions: `match_coordinates_sky` and `search_around_sky`.\n",
    "While `KdTreeCrossmatch` allows to specify both search radius and number of neighbors to get, `match_coordinates_sky` may provide n-th neighbor only, and `search_around_sky` may provide all neighbors in a given search radius.\n",
    "\n",
    "This notebook shows, that for the simplest case of a small radius and a single cross-match pair, LSDB's algorithm works a bit slower than `match_coordinates_sky`, if the least one receives `SkyCoord` as an input.\n",
    "My hypothesis that it happens because `SkyCoord` coverts spherical coordinates to cartesian at the construction time, while `KdTreeCrossmatch` does this conversion itself. This makes the result of the comparison unfair.\n",
    "\n",
    "From the other side, `KdTreeCrossmatch` cross-matches coordinate catalogs two orders of magnitude faster than `SkyCoord` construction, so if we would include initialisation of `SkyCoord`s into the astropy benchmark, it would run much slower than LSDB's algorithm.\n",
    "    "
   ],
   "id": "b000e75ad9b820c6"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "### Import packages",
   "id": "91b0e8444c9b9da"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-06T20:48:24.365177Z",
     "start_time": "2024-05-06T20:48:22.356275Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import pandas as pd\n",
    "\n",
    "import lsdb\n",
    "from astropy.coordinates import Angle, SkyCoord, match_coordinates_sky, search_around_sky\n",
    "from lsdb.core.crossmatch.kdtree_match import KdTreeCrossmatch"
   ],
   "id": "initial_id",
   "outputs": [],
   "execution_count": 1
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "### Set cross-match parameters",
   "id": "44e94149e90a4b42"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-06T20:48:24.367520Z",
     "start_time": "2024-05-06T20:48:24.366103Z"
    }
   },
   "cell_type": "code",
   "source": [
    "N_NEIGHBORS = 10\n",
    "RADIUS_ARCSEC = 60.0"
   ],
   "id": "a3b1d79edcf34562",
   "outputs": [],
   "execution_count": 2
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "### Obtain LINCC Frameworks' \"half-degree\" catalogs: Gaia DR3 and ZTF DR14",
   "id": "b7b5e6ecb0cfd072"
  },
  {
   "metadata": {
    "jupyter": {
     "is_executing": true
    },
    "ExecuteTime": {
     "start_time": "2024-05-06T20:48:24.368076Z"
    }
   },
   "cell_type": "code",
   "source": [
    "%%time\n",
    "\n",
    "gaia_path = 'https://epyc.astro.washington.edu/~lincc-frameworks/half_degree_surveys/gaia/'\n",
    "ztf_path = 'https://epyc.astro.washington.edu/~lincc-frameworks/half_degree_surveys/ztf/ztf_object/'\n",
    "\n",
    "gaia_catalog = lsdb.read_hipscat(gaia_path, columns=[\"ra\", \"dec\"])\n",
    "ztf_catalog = lsdb.read_hipscat(ztf_path, columns=[\"ra\", \"dec\"])\n",
    "\n",
    "gaia_df = gaia_catalog.compute().reset_index(drop=True)\n",
    "ztf_df = ztf_catalog.compute().reset_index(drop=True)"
   ],
   "id": "738fca4aa2658312",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "### Run LSDB's cross-matching algorithm ",
   "id": "ddcc5a4037a86ed"
  },
  {
   "metadata": {
    "jupyter": {
     "is_executing": true
    }
   },
   "cell_type": "code",
   "source": [
    "def lsdb_crossmatch():\n",
    "    # LSDB's cross-matching algorithms are not designed to be used out of \n",
    "    # lsdb.catalog.Catalog.crossmatch(), so we construct the object with some\n",
    "    # mock properties.\n",
    "    kd_tree_crossmatch = KdTreeCrossmatch(\n",
    "        left=gaia_df,\n",
    "        right=ztf_df,\n",
    "        left_order=-1,\n",
    "        left_pixel=-1,\n",
    "        right_order=-1,\n",
    "        right_pixel=-1,\n",
    "        left_metadata=gaia_catalog.hc_structure,\n",
    "        right_metadata=ztf_catalog.hc_structure,\n",
    "        right_margin_hc_structure=None,\n",
    "        suffixes=(\"\", \"\"),\n",
    "    )\n",
    "\n",
    "    return kd_tree_crossmatch.crossmatch(n_neighbors=N_NEIGHBORS, radius_arcsec=RADIUS_ARCSEC)\n",
    "\n",
    "%timeit lsdb_crossmatch()\n",
    "lsdb_result = lsdb_crossmatch().reset_index(drop=True)"
   ],
   "id": "8092ad539cb05670",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "### Create `SkyCoord`s for astropy - note that it takes a while",
   "id": "41eb686a61cd7c28"
  },
  {
   "metadata": {
    "jupyter": {
     "is_executing": true
    }
   },
   "cell_type": "code",
   "source": [
    "%%time\n",
    "\n",
    "gaia_skycoord = SkyCoord(ra=gaia_df[\"ra\"], dec=gaia_df[\"dec\"], unit=\"deg\")\n",
    "ztf_skycoord = SkyCoord(ra=ztf_df[\"ra\"], dec=ztf_df[\"dec\"], unit=\"deg\")"
   ],
   "id": "8c36347df579fb7",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "### Run `astropy`'s `match_coordinates_sky`",
   "id": "2b05200c288902d7"
  },
  {
   "metadata": {
    "jupyter": {
     "is_executing": true
    }
   },
   "cell_type": "code",
   "source": [
    "def astropy_match_nthneighbor(nthneighbor: int) -> pd.DataFrame:\n",
    "    idx_ztf, d2, _d3 = match_coordinates_sky(gaia_skycoord, ztf_skycoord, nthneighbor=nthneighbor)\n",
    "    result = pd.concat(\n",
    "        [\n",
    "            gaia_df,\n",
    "            ztf_df.iloc[idx_ztf].reset_index(drop=True),\n",
    "            pd.DataFrame({\"_dist_arcsec\": d2.to_value('arcsec')}),\n",
    "        ],\n",
    "        axis=1,\n",
    "    )\n",
    "    return result[d2 < Angle(RADIUS_ARCSEC, 'arcsec')]\n",
    "\n",
    "def astropy_match():\n",
    "    result = pd.concat(\n",
    "        [astropy_match_nthneighbor(nth) for nth in range(1, N_NEIGHBORS + 1)],\n",
    "        axis=0,\n",
    "    )\n",
    "    return result\n",
    "\n",
    "%timeit astropy_match()\n",
    "astropy_match_result = astropy_match().sort_index(kind='stable').reset_index(drop=True)"
   ],
   "id": "9b1fd50b853f6262",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "### Run `astropy`'s `search_around_sky`",
   "id": "239a894405540c94"
  },
  {
   "metadata": {
    "jupyter": {
     "is_executing": true
    }
   },
   "cell_type": "code",
   "source": [
    "def astropy_search():\n",
    "    idx_gaia, idx_ztf, d2, _d3 = search_around_sky(\n",
    "        gaia_skycoord,\n",
    "        ztf_skycoord,\n",
    "        Angle(RADIUS_ARCSEC, 'arcsec'),\n",
    "    )\n",
    "    idx_df = pd.DataFrame(\n",
    "        {\"idx_gaia\": idx_gaia, \"idx_ztf\": idx_ztf, \"_dist_arcsec\": d2.to_value(\"arcsec\")},\n",
    "    )\n",
    "    idx_df = idx_df.groupby(\"idx_gaia\").apply(lambda x: x.nsmallest(N_NEIGHBORS, \"_dist_arcsec\"))\n",
    "    idx_df = idx_df.reset_index(drop=True)\n",
    "    return pd.concat(\n",
    "        [\n",
    "            gaia_df.iloc[idx_df[\"idx_gaia\"]].reset_index(drop=True),\n",
    "            ztf_df.iloc[idx_df[\"idx_ztf\"]].reset_index(drop=True),\n",
    "            idx_df[\"_dist_arcsec\"].to_frame(),\n",
    "        ],\n",
    "        axis=1,\n",
    "    )\n",
    "\n",
    "%timeit astropy_search()\n",
    "astropy_search_result = astropy_search().reset_index(drop=True)"
   ],
   "id": "584c257806cdd8a6",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "### Compare the results",
   "id": "db66c79e8fdc911a"
  },
  {
   "metadata": {
    "jupyter": {
     "is_executing": true
    }
   },
   "cell_type": "code",
   "source": [
    "pd.testing.assert_frame_equal(\n",
    "    lsdb_result,\n",
    "    astropy_match_result,\n",
    ")\n",
    "pd.testing.assert_frame_equal(\n",
    "    lsdb_result,\n",
    "    astropy_search_result,\n",
    ")\n",
    "print('All dataframes are equal')"
   ],
   "id": "8bb47110ba085888",
   "outputs": [],
   "execution_count": null
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}