{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Matching of haloes to clusters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import cm\n",
    "from matplotlib.colors import Normalize\n",
    "from scipy import stats\n",
    "\n",
    "from match_clusters import *\n",
    "\n",
    "import scienceplots\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot $R - \\log M$ as a function of significance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {},
   "outputs": [],
   "source": [
    "cat = csiborgtools.read.CSiBORG2Catalogue(\n",
    "    17417, 99, \"main\", bounds={\"totmass\": (1e12, None)})\n",
    "pos = cat[\"cartesian_pos\"]\n",
    "logmass = np.log10(cat[\"totmass\"])\n",
    "\n",
    "model = csiborgtools.match.MatchingProbability(pos, logmass)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 95348 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "xrange = np.linspace(13, 15.2, 250)\n",
    "log_p_values = np.linspace(-3, -0.25, 1000)\n",
    "cmap = cm.viridis\n",
    "norm = Normalize(vmin=log_p_values.min(), vmax=log_p_values.max())\n",
    "sm = cm.ScalarMappable(cmap=cmap, norm=norm)\n",
    "sm.set_array([])  # Necessary for the colorbar to use the ScalarMappable\n",
    "\n",
    "\n",
    "fig, ax = plt.figure(), plt.gca()\n",
    "for p in log_p_values:\n",
    "    y = [model.inverse_cdf(10**p, x) for x in xrange]\n",
    "    ax.plot(xrange, y, color=cmap(norm(p)))\n",
    "\n",
    "ls = [\":\", \"--\", \"-.\"]\n",
    "for n in range(1, 4):\n",
    "    p = stats.norm.sf(n) * 2\n",
    "    y = [model.inverse_cdf(p, x) for x in xrange]\n",
    "    ax.plot(xrange, y, color=\"red\", label=fr\"${n}\\sigma$\", linestyle=ls[n-1])\n",
    "\n",
    "ax.set_xlabel(r\"$\\log M_{\\mathrm{FoF}} ~ [M_\\odot / h]$\")\n",
    "ax.set_ylabel(r\"$R ~ [\\mathrm{Mpc} / h]$\")\n",
    "ax.set_ylim(1, 55)\n",
    "ax.legend()\n",
    "cbar = fig.colorbar(sm, ax=ax, label=r\"$\\log p$\", pad=0)\n",
    "ax.set_xlim(xrange.min(), xrange.max())\n",
    "fig.tight_layout()\n",
    "fig.savefig(\"../../plots/matching_probability.png\", dpi=450)\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test matching"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Opening catalogues: 100%|██████████| 20/20 [00:07<00:00,  2.59it/s]\n"
     ]
    }
   ],
   "source": [
    "simname = \"csiborg2_main\"\n",
    "bounds = {\"totmass\": (1e12, None)}\n",
    "\n",
    "cats = open_cats(simname, bounds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:04<00:00,  4.32it/s]\n"
     ]
    }
   ],
   "source": [
    "model = csiborgtools.match.MatchCatalogues(cats)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "14.690196080028514\n"
     ]
    }
   ],
   "source": [
    "cluster = csiborgtools.clusters[\"Leo\"]\n",
    "x0 = cluster.cartesian_pos(cats[0].boxsize)\n",
    "logmass0 = np.log10(cluster.mass)\n",
    "print(logmass0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Matching catalogues:   0%|          | 0/20 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found 2 plausible matches in catalogue 0.\n",
      "    1: CDF = 1.640e-04, index = 69, logM = 14.423 Msun / h, dx = 1.103 Mpc / h.\n",
      "    2: CDF = 3.798e-02, index = 1, logM = 15.135 Msun / h, dx = 22.790 Mpc / h.\n",
      "\n",
      "Found 2 plausible matches in catalogue 1.\n",
      "    1: CDF = 2.669e-04, index = 58, logM = 14.474 Msun / h, dx = 1.381 Mpc / h.\n",
      "    2: CDF = 4.125e-02, index = 1, logM = 15.120 Msun / h, dx = 22.606 Mpc / h.\n",
      "\n",
      "Found 2 plausible matches in catalogue 2.\n",
      "    1: CDF = 1.993e-04, index = 68, logM = 14.460 Msun / h, dx = 1.232 Mpc / h.\n",
      "    2: CDF = 3.619e-02, index = 1, logM = 15.145 Msun / h, dx = 22.957 Mpc / h.\n",
      "\n",
      "Found 2 plausible matches in catalogue 3.\n",
      "    1: CDF = 1.290e-04, index = 45, logM = 14.545 Msun / h, dx = 1.187 Mpc / h.\n",
      "    2: CDF = 4.352e-02, index = 2, logM = 15.117 Msun / h, dx = 22.863 Mpc / h.\n",
      "\n",
      "Found 2 plausible matches in catalogue 4.\n",
      "    1: CDF = 9.658e-05, index = 47, logM = 14.541 Msun / h, dx = 1.072 Mpc / h.\n",
      "    2: CDF = 4.109e-02, index = 1, logM = 15.120 Msun / h, dx = 22.605 Mpc / h.\n",
      "\n",
      "Found 2 plausible matches in catalogue 5.\n",
      "    1: CDF = 1.171e-04, index = 53, logM = 14.493 Msun / h, dx = 1.074 Mpc / h.\n",
      "    2: CDF = 4.936e-02, index = 1, logM = 15.097 Msun / h, dx = 22.768 Mpc / h.\n",
      "\n",
      "Found 2 plausible matches in catalogue 6.\n",
      "    1: CDF = 2.107e-04, index = 41, logM = 14.537 Msun / h, dx = 1.383 Mpc / h.\n",
      "    2: CDF = 4.406e-02, index = 1, logM = 15.116 Msun / h, dx = 22.923 Mpc / h.\n",
      "\n",
      "Found 2 plausible matches in catalogue 8.\n",
      "    1: CDF = 3.236e-04, index = 48, logM = 14.487 Msun / h, dx = 1.497 Mpc / h.\n",
      "    2: CDF = 3.737e-02, index = 2, logM = 15.137 Msun / h, dx = 22.756 Mpc / h.\n",
      "\n",
      "Found 2 plausible matches in catalogue 9.\n",
      "    1: CDF = 2.925e-04, index = 51, logM = 14.505 Msun / h, dx = 1.482 Mpc / h.\n",
      "    2: CDF = 3.504e-02, index = 2, logM = 15.145 Msun / h, dx = 22.737 Mpc / h.\n",
      "\n",
      "Found 2 plausible matches in catalogue 10.\n",
      "    1: CDF = 1.577e-04, index = 58, logM = 14.451 Msun / h, dx = 1.127 Mpc / h.\n",
      "    2: CDF = 3.792e-02, index = 2, logM = 15.133 Msun / h, dx = 22.693 Mpc / h.\n",
      "\n",
      "Found 2 plausible matches in catalogue 11.\n",
      "    1: CDF = 2.399e-04, index = 64, logM = 14.438 Msun / h, dx = 1.276 Mpc / h.\n",
      "    2: CDF = 3.375e-02, index = 2, logM = 15.157 Msun / h, dx = 23.132 Mpc / h.\n",
      "\n",
      "Found 2 plausible matches in catalogue 12.\n",
      "    1: CDF = 2.907e-05, index = 60, logM = 14.460 Msun / h, dx = 0.648 Mpc / h.\n",
      "    2: CDF = 3.438e-02, index = 2, logM = 15.155 Msun / h, dx = 23.149 Mpc / h.\n",
      "\n",
      "Found 2 plausible matches in catalogue 13.\n",
      "    1: CDF = 6.630e-05, index = 51, logM = 14.508 Msun / h, dx = 0.906 Mpc / h.\n",
      "    2: CDF = 4.297e-02, index = 2, logM = 15.128 Msun / h, dx = 23.342 Mpc / h.\n",
      "\n",
      "Found 2 plausible matches in catalogue 14.\n",
      "    1: CDF = 3.599e-05, index = 53, logM = 14.501 Msun / h, dx = 0.733 Mpc / h.\n",
      "    2: CDF = 3.438e-02, index = 1, logM = 15.154 Msun / h, dx = 23.077 Mpc / h.\n",
      "\n",
      "Found 2 plausible matches in catalogue 15.\n",
      "    1: CDF = 2.172e-05, index = 65, logM = 14.459 Msun / h, dx = 0.588 Mpc / h.\n",
      "    2: CDF = 3.481e-02, index = 1, logM = 15.150 Msun / h, dx = 22.989 Mpc / h.\n",
      "\n",
      "Found 2 plausible matches in catalogue 16.\n",
      "    1: CDF = 1.061e-04, index = 77, logM = 14.420 Msun / h, dx = 0.952 Mpc / h.\n",
      "    2: CDF = 3.451e-02, index = 1, logM = 15.158 Msun / h, dx = 23.358 Mpc / h.\n",
      "\n",
      "Found 2 plausible matches in catalogue 17.\n",
      "    1: CDF = 6.105e-04, index = 73, logM = 14.446 Msun / h, dx = 1.758 Mpc / h.\n",
      "    2: CDF = 4.261e-02, index = 2, logM = 15.118 Msun / h, dx = 22.738 Mpc / h.\n",
      "\n",
      "Found 2 plausible matches in catalogue 18.\n",
      "    1: CDF = 1.621e-04, index = 69, logM = 14.444 Msun / h, dx = 1.128 Mpc / h.\n",
      "    2: CDF = 3.606e-02, index = 2, logM = 15.144 Msun / h, dx = 22.879 Mpc / h.\n",
      "\n",
      "Found 2 plausible matches in catalogue 19.\n",
      "    1: CDF = 2.105e-04, index = 44, logM = 14.532 Msun / h, dx = 1.374 Mpc / h.\n",
      "    2: CDF = 4.477e-02, index = 2, logM = 15.119 Msun / h, dx = 23.189 Mpc / h.\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Matching catalogues: 100%|██████████| 20/20 [00:00<00:00, 324.46it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{0: 0.0001639847755937618, 1: 0.0002668512052416139, 2: 0.00019926862221641972, 3: 0.00012900322742348536, 4: 9.657721769740224e-05, 5: 0.00011705910455406254, 6: 0.00021073714493358153, 7: 0.00015055517039153088, 8: 0.00032362529099960735, 9: 0.00029247443413660257, 10: 0.0001576842700352632, 11: 0.00023994027660023676, 12: 2.9066704152591605e-05, 13: 6.630406935614896e-05, 14: 3.598720831643476e-05, 15: 2.1723041493393325e-05, 16: 0.00010607153796671209, 17: 0.0006105268685839826, 18: 0.0001620825661784453, 19: 0.00021048883480756508}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "ps, indxs = model(x0, logmass0, pvalue_threshold=0.05, rmax=30)\n",
    "print(ps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ps_coma = np.array(list(ps.values()))\n",
    "# ps_centaurus = np.array(list(ps.values()))\n",
    "# ps_shapley = np.array(list(ps.values()))\n",
    "ps_virgo = np.array(list(ps.values()))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 9592 characters]",
      "text/plain": [
       "<Figure size 350x262.5 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "ps = {\"Coma\": ps_coma, \"Centaurus\": ps_centaurus, \"Shapley\": ps_shapley, \"Virgo\": ps_virgo}\n",
    "\n",
    "with plt.style.context('science'):\n",
    "    plt.figure()\n",
    "    for name, p in ps.items():\n",
    "        plt.hist(np.log10(p), bins=\"auto\", label=name, density=True, alpha=0.5)\n",
    "\n",
    "    plt.legend(loc=\"upper left\", ncol=2)\n",
    "    plt.xlabel(r\"$\\log p$\")\n",
    "    plt.ylabel(\"Counts\")\n",
    "    plt.tight_layout()\n",
    "    plt.savefig(\"../../plots/matching_probability_histogram.png\", dpi=450, bbox_inches=\"tight\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "dx = np.asarray([np.linalg.norm(cats[i][\"cartesian_pos\"][indx] - x0) for i, indx in indxs.items()])\n",
    "logM = np.asarray([np.log10(cats[i][\"totmass\"][indx]) for i, indx in indxs.items()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{0: 1.4193360633019125e-05,\n",
       " 1: 1.5814267730451803e-05,\n",
       " 2: 1.578598264362352e-05,\n",
       " 3: 2.9062728345241595e-05,\n",
       " 4: 2.329484601226728e-05,\n",
       " 5: 2.9712920921687314e-05,\n",
       " 6: 2.075305503379532e-05,\n",
       " 7: 1.074050678318006e-05,\n",
       " 8: 1.888041209263669e-05,\n",
       " 9: 1.9893007069504698e-05,\n",
       " 10: 1.7657580479335522e-05,\n",
       " 11: 2.8592305873931245e-05,\n",
       " 12: 1.878843367530969e-05,\n",
       " 13: 8.355070255428032e-06,\n",
       " 14: 1.413250003556854e-05,\n",
       " 15: 1.4124073620669364e-05,\n",
       " 16: 1.2492311356582242e-05,\n",
       " 17: 2.5997234836760974e-05,\n",
       " 18: 1.3209139431147321e-05,\n",
       " 19: 1.0199631684804089e-05}"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 114,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 15824 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.figure()\n",
    "plt.scatter(np.arange(len(dx)), dx)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "metadata": {},
   "outputs": [],
   "source": [
    "ps, dlog, indxs = model.cdf_per_halo(x0, logmass0)\n",
    "r = np.linalg.norm(pos - x0.reshape(1, -1), axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "None None [[15.430657 15.189978 15.118925 ... 12.098506 12.76248  12.336433]]\n"
     ]
    }
   ],
   "source": [
    "p, k = model.match_halo(x0, logmass0, pvalue_threshold=0.1, max_absdlogmass=1)\n",
    "print(p, k, logmass[k])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 111,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 34756 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.figure()\n",
    "plt.plot(xrange, y)\n",
    "plt.plot(xrange, y2)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [],
   "source": [
    "x, y = model.pdf_marginalized(xrange, 14.0, 42.5, 4)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv_csiborg",
   "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.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
