{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 6,
   "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",
    "import csiborgtools\n",
    "\n",
    "from hmf import *\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "bin_edges = np.arange(11.8, 15.2, 0.25)\n",
    "xbin = (bin_edges[1:] + bin_edges[:-1]) / 2\n",
    "bin_edges = 10**bin_edges\n",
    "xbin = 10**xbin\n",
    "\n",
    "paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/20 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:05<00:00,  3.50it/s]\n",
      "100%|██████████| 20/20 [00:05<00:00,  3.58it/s]\n",
      "100%|██████████| 20/20 [00:05<00:00,  3.85it/s]\n"
     ]
    }
   ],
   "source": [
    "hmf_cb2_main = calculate_hmf(\"csiborg2_main\", bin_edges)\n",
    "hmf_cb2_random = calculate_hmf(\"csiborg2_random\", bin_edges)\n",
    "hmf_cb2_varysmall = calculate_hmf(\"csiborg2_varysmall\", bin_edges)\n",
    "hmf_mpdpl2 = MDPL2_HMF(bin_edges)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## CB2_varysmall HMF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 142108 characters]",
      "text/plain": [
       "<Figure size 1000x400 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "k = np.where(paths.get_ics(\"csiborg2_main\") == 16417)[0][0]\n",
    "\n",
    "fig, axs = plt.subplots(1, 2, figsize=(10, 4))\n",
    "lw = plt.rcParams['lines.linewidth']\n",
    "fig.subplots_adjust(hspace=0)\n",
    "axs[0].plot(xbin, hmf_cb2_main[k], zorder=1, color=\"red\", label=\"CB2 Main\")\n",
    "for i in range(len(hmf_cb2_varysmall)):\n",
    "    axs[0].plot(xbin, hmf_cb2_varysmall[i], lw=lw/2, color=\"black\", zorder=0,\n",
    "                label=\"CB2 Varysmall\" if i == 0 else None)\n",
    "\n",
    "    axs[1].plot(xbin, hmf_cb2_varysmall[i] / hmf_cb2_main[k], lw=lw/3, color=\"black\", zorder=0)\n",
    "\n",
    "\n",
    "xmin, xmax = xbin.min(), xbin.max()\n",
    "for i in range(2):\n",
    "    axs[i].set_xscale(\"log\")\n",
    "    axs[i].set_xlabel(r\"$M_{\\mathrm{FoF}}$ $[M_{\\odot}]$\")\n",
    "    axs[i].set_xlim(xmin, xmax)\n",
    "\n",
    "axs[0].legend()\n",
    "axs[0].set_ylabel(r\"HMF $[1 / (h^{-3} \\mathrm{Mpc}^3~\\mathrm{dex})]$\")\n",
    "axs[1].set_ylabel(r\"$\\mathrm{HMF}_{\\mathrm{Varysmall}}~/~\\mathrm{HMF}_{\\mathrm{Main}}$\")\n",
    "axs[0].set_yscale(\"log\")\n",
    "axs[1].set_ylim(0.9, 1.1)\n",
    "axs[1].axhline(1, color=\"red\", lw=lw, zorder=0)\n",
    "\n",
    "fig.tight_layout()\n",
    "fig.savefig(\"../../plots/varysmall_hmf.png\", dpi=450)\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 44304 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.figure()\n",
    "\n",
    "ylow, yhigh = np.percentile(hmf_cb2_main, [16, 84], axis=0)\n",
    "plt.fill_between(xbin, ylow, yhigh, alpha=0.75, label=\"CB2 Main\")\n",
    "\n",
    "ylow, yhigh = np.percentile(hmf_cb2_random, [16, 84], axis=0)\n",
    "plt.fill_between(xbin, ylow, yhigh, alpha=0.75, label=\"CB2 Random\")\n",
    "\n",
    "plt.plot(xbin, hmf_mpdpl2, label=\"MDPL2\", color=\"red\", ls=\"--\")\n",
    "\n",
    "plt.yscale(\"log\")\n",
    "plt.xscale(\"log\")\n",
    "plt.xlabel(r\"$M_{\\rm FoF} ~ [M_\\odot / h]$\")\n",
    "plt.ylabel(r\"HMF $[1 / (h^{-3} \\mathrm{Mpc}^3~\\mathrm{dex})]$\")\n",
    "plt.legend()\n",
    "plt.xlim(xbin.min(), xbin.max())\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.savefig(\"../../plots/HMF_comparison.png\", dpi=450)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}
