{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lagrangian patch example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import scienceplots\n",
    "\n",
    "import csiborgtools\n",
    "\n",
    "\n",
    "%matplotlib inline\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [],
   "source": [
    "paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)\n",
    "bounds = {\"dist\": (0, 155)}\n",
    "\n",
    "nsim0 = 7444 + 24 * 61\n",
    "nsnap_final = paths.get_snapshots(nsim0, \"csiborg1\")[-1]\n",
    "snap0_initial = csiborgtools.read.CSiBORG1Snapshot(nsim0, nsnap=1)\n",
    "cat0 = csiborgtools.read.CSiBORG1Catalogue(nsim0, paths, snapshot=snap0_initial, bounds=bounds)\n",
    "\n",
    "nsimx = 9844\n",
    "nsnap_final = paths.get_snapshots(nsimx, \"csiborg1\")[-1]\n",
    "snapx_initial = csiborgtools.read.CSiBORG1Snapshot(nsimx, nsnap=1)\n",
    "catx = csiborgtools.read.CSiBORG1Catalogue(nsimx, paths, snapshot=snapx_initial, bounds=bounds)\n",
    "\n",
    "pair_overlapper = csiborgtools.summary.PairOverlap(cat0, catx, 13.25, )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "15.319891\n",
      "74.15419337073868\n"
     ]
    }
   ],
   "source": [
    "k = np.argsort(cat0[\"totmass\"])[::-1][1]\n",
    "\n",
    "print(np.log10(cat0[\"totmass\"][k]))\n",
    "print(cat0[\"dist\"][k])\n",
    "\n",
    "hid0 = cat0[\"index\"][k]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overlaps are  [7.8512335e-01 9.1506150e-03 2.9659141e-03 3.0178351e-03 2.0852157e-04]\n"
     ]
    }
   ],
   "source": [
    "print(\"Overlaps are \", pair_overlapper.overlap(True)[k])\n",
    "hidx = catx[\"index\"][pair_overlapper[\"match_indxs\"][k][0]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "overlapper = csiborgtools.match.ParticleOverlap(**{\"box_size\": 2048, \"bckg_halfsize\": 512})\n",
    "\n",
    "pos0_initial = snap0_initial.halo_coordinates(hid0)\n",
    "pos0_initial = csiborgtools.match.pos2cell(pos0_initial / cat0.boxsize, 2048)\n",
    "mass0_initial = snap0_initial.halo_masses(hid0)\n",
    "\n",
    "\n",
    "posx_initial = snapx_initial.halo_coordinates(hidx)\n",
    "posx_initial = csiborgtools.match.pos2cell(posx_initial / catx.boxsize, 2048)\n",
    "massx_initial = snapx_initial.halo_masses(hidx)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "smooth_kwargs = {\"sigma\": 1, \"mode\": \"constant\", \"cval\": 0}\n",
    "# smooth_kwargs = None\n",
    "\n",
    "delta1_initial, delta2_initial, cellmins, __ = overlapper.make_deltas(\n",
    "    pos0_initial, posx_initial, mass0_initial, massx_initial, smooth_kwargs=smooth_kwargs, )\n",
    "\n",
    "x0 = cellmins[0] / 2048 * 677.7\n",
    "xf = (cellmins[0] + delta1_initial.shape[0]) / 2048 * 677.7\n",
    "\n",
    "y0 = cellmins[1] / 2048 * 677.7\n",
    "yf = (cellmins[1] + delta1_initial.shape[1]) / 2048 * 677.7"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_2722989/1176586521.py:10: RuntimeWarning: Mean of empty slice\n",
      "  axs[0].imshow(np.nanmean(delta1_initial, axis=-1), **kwargs)\n",
      "/tmp/ipykernel_2722989/1176586521.py:11: RuntimeWarning: Mean of empty slice\n",
      "  axs[1].imshow(np.nanmean(delta2_initial, axis=-1), **kwargs)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 244392 characters]",
      "text/plain": [
       "<Figure size 830x498 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "with plt.style.context(\"science\"):\n",
    "    # plt.rcParams.update({'font.size': 9})\n",
    "\n",
    "    fig, axs = plt.subplots(1, 2, figsize=(8.3, 0.6* 8.3), sharey=True, sharex=True)\n",
    "    fig.subplots_adjust(wspace=0.05, hspace=0.05)\n",
    "\n",
    "    kwargs = {\"origin\": \"lower\", \"cmap\": \"viridis\", \"vmin\": 0,\n",
    "              \"extent\": [x0, xf, y0, yf]}\n",
    "\n",
    "    axs[0].imshow(np.mean(delta1_initial, axis=-1), **kwargs)\n",
    "    axs[1].imshow(np.mean(delta2_initial, axis=-1), **kwargs)\n",
    "\n",
    "    for i in range(2):\n",
    "        axs[i].set_xlabel(r\"$x ~ [\\mathrm{cMpc} / h]$\")\n",
    "    axs[0].set_ylabel(r\"$y ~ [\\mathrm{cMpc} / h]$\")\n",
    "\n",
    "    axs[0].set_title(f\"MCMC step {nsim0}\")\n",
    "    axs[1].set_title(f\"MCMC step {nsimx}\")\n",
    "\n",
    "    fig.tight_layout()\n",
    "    fig.savefig(\"../../plots/lagrangian_patch_example.pdf\", bbox_inches=\"tight\", dpi=450)\n",
    "    fig.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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
