{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copyright (C) 2024 Richard Stiskalek\n",
    "# This program is free software; you can redistribute it and/or modify it\n",
    "# under the terms of the GNU General Public License as published by the\n",
    "# Free Software Foundation; either version 3 of the License, or (at your\n",
    "# option) any later version.\n",
    "#\n",
    "# This program is distributed in the hope that it will be useful, but\n",
    "# WITHOUT ANY WARRANTY; without even the implied warranty of\n",
    "# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General\n",
    "# Public License for more details.\n",
    "#\n",
    "# You should have received a copy of the GNU General Public License along\n",
    "# with this program; if not, write to the Free Software Foundation, Inc.,\n",
    "# 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.\n",
    "from os.path import join\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from astropy.io import fits\n",
    "from tqdm import tqdm\n",
    "from scipy.ndimage import gaussian_filter\n",
    "import scienceplots\n",
    "from warnings import warn\n",
    "from astropy.cosmology import FlatLambdaCDM\n",
    "from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
    "\n",
    "import csiborgtools\n",
    "\n",
    "from reconstruction_comparison import simname_to_pretty\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline\n",
    "\n",
    "paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)\n",
    "fdir = \"/mnt/extraspace/rstiskalek/csiborg_postprocessing/peculiar_velocity\"\n",
    "\n",
    "CLUSTERS = {\n",
    "    \"Shapley\": (-116.37827331, 69.68752732, -14.90318191),\n",
    "    \"Virgo\": (-3.48, 14.86, -2.21),\n",
    "    \"Norma\": (-50.26, -7.06, 6.44),\n",
    "    \"Coma\": (0.48, 72.79, 10.59),\n",
    "    \"Perseus\": (49.94, -10.73, -12.98),\n",
    "    \"Centaurus\": (-34.25, 14.93, -7.56)\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from astropy.coordinates import SkyCoord\n",
    "# from astropy.coordinates import Supergalactic\n",
    "# from astropy.units import deg, Mpc\n",
    "# from astropy.cosmology import FlatLambdaCDM\n",
    "\n",
    "# SPEED_OF_LIGHT = 299_792.458 \n",
    "\n",
    "\n",
    "# cosmo = FlatLambdaCDM(H0=100, Om0=0.3)\n",
    "\n",
    "# RA = 196.490821\n",
    "# dec = -33.067461\n",
    "# dist = cosmo.comoving_distance(0.048)\n",
    "\n",
    "# c = SkyCoord(ra= RA * deg, dec= dec * deg, distance=dist)\n",
    "# c = c.transform_to(Supergalactic)\n",
    "# c.cartesian"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_field(simname, nsim, kind, MAS, grid):\n",
    "    # Open the field reader.\n",
    "    if simname == \"csiborg1\":\n",
    "        field_reader = csiborgtools.read.CSiBORG1Field(nsim)\n",
    "    elif \"csiborg2_\" in simname:\n",
    "        simkind = simname.split(\"_\")[-1]\n",
    "        field_reader = csiborgtools.read.CSiBORG2Field(nsim, simkind)\n",
    "    elif simname == \"csiborg2X\":\n",
    "        field_reader = csiborgtools.read.CSiBORG2XField(nsim)\n",
    "    elif simname == \"Carrick2015\":\n",
    "        folder = \"/mnt/extraspace/rstiskalek/catalogs\"\n",
    "        warn(f\"Using local paths from `{folder}`.\", RuntimeWarning)\n",
    "        if kind == \"density\":\n",
    "            fpath = join(folder, \"twompp_density_carrick2015.npy\")\n",
    "            return np.load(fpath).astype(np.float32)\n",
    "        elif kind == \"velocity\":\n",
    "            fpath = join(folder, \"twompp_velocity_carrick2015.npy\")\n",
    "            field = np.load(fpath).astype(np.float32)\n",
    "\n",
    "            # Because the Carrick+2015 data is in the following form:\n",
    "            # \"The velocities are predicted peculiar velocities in the CMB\n",
    "            # frame in Galactic Cartesian coordinates, generated from the\n",
    "            # \\(\\delta_g^*\\) field with \\(\\beta^* = 0.43\\) and an external\n",
    "            # dipole \\(V_\\mathrm{ext} = [89,-131,17]\\) (Carrick et al Table 3)\n",
    "            # has already been added.\"\"\n",
    "            field[0] -= 89\n",
    "            field[1] -= -131\n",
    "            field[2] -= 17\n",
    "            field /= 0.43\n",
    "\n",
    "            return field\n",
    "        else:\n",
    "            raise ValueError(f\"Unknown field kind: `{kind}`.\")\n",
    "    elif simname == \"CLONES\":\n",
    "        field_reader = csiborgtools.read.CLONESField(nsim)\n",
    "    elif simname == \"CF4\":\n",
    "        folder = \"/mnt/extraspace/rstiskalek/catalogs/CF4\"\n",
    "        warn(f\"Using local paths from `{folder}`.\", RuntimeWarning)\n",
    "\n",
    "        if kind == \"density\":\n",
    "            fpath = join(folder, f\"CF4_new_128-z008_realization{nsim}_delta.fits\")     # noqa\n",
    "        elif kind == \"velocity\":\n",
    "            fpath = join(folder, f\"CF4_new_128-z008_realization{nsim}_velocity.fits\")  # noqa\n",
    "        else:\n",
    "            raise ValueError(f\"Unknown field kind: `{kind}`.\")\n",
    "\n",
    "        field = fits.open(fpath)[0].data\n",
    "\n",
    "        # https://projets.ip2i.in2p3.fr//cosmicflows/ says to multiply by 52\n",
    "        if kind == \"velocity\":\n",
    "            field *= 52\n",
    "\n",
    "        return field.astype(np.float32)\n",
    "    elif simname == \"Lilow2024\":\n",
    "        folder = \"/mnt/extraspace/rstiskalek/catalogs\"\n",
    "        warn(f\"Using local paths from `{folder}`.\", RuntimeWarning)\n",
    "\n",
    "        if kind == \"density\":\n",
    "            fpath = join(folder, \"Lilow2024_density.npy\")\n",
    "            field = np.load(fpath)\n",
    "        elif kind == \"velocity\":\n",
    "            field = []\n",
    "            for p in [\"x\", \"y\", \"z\"]:\n",
    "                fpath = join(folder, f\"Lilow2024_{p}Velocity.npy\")\n",
    "                field.append(np.load(fpath).astype(np.float32))\n",
    "            field = np.stack(field)\n",
    "\n",
    "        return field.astype(np.float32)\n",
    "    else:\n",
    "        raise ValueError(f\"Unknown simulation name: `{simname}`.\")\n",
    "\n",
    "    # Read in the field.\n",
    "    if kind == \"density\":\n",
    "        field = field_reader.density_field(MAS=MAS, grid=grid)\n",
    "    elif kind == \"velocity\":\n",
    "        field = field_reader.velocity_field(MAS=MAS, grid=grid)\n",
    "    else:\n",
    "        raise ValueError(f\"Unknown field kind: `{kind}`.\")\n",
    "    \n",
    "    # NOTE added here \n",
    "    if kind == \"density\" and (\"csiborg\" in simname or simname == \"CLONES\"):\n",
    "        Om0 = csiborgtools.simname2Omega_m(simname)\n",
    "        cosmo = FlatLambdaCDM(H0=100, Om0=Om0)\n",
    "        mean_rho_matter = cosmo.critical_density0.to(\"Msun/kpc^3\").value\n",
    "        mean_rho_matter *= Om0\n",
    "        field /= mean_rho_matter\n",
    "\n",
    "    return field"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_slice(simname, xmin, xmax, ngrid):\n",
    "    boxsize = csiborgtools.simname2boxsize(simname)\n",
    "\n",
    "    paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)\n",
    "    nsims = paths.get_ics(simname)\n",
    "\n",
    "    if simname in [\"Carrick2015\", \"Lilow2024\"]:\n",
    "        frame = \"galactic\"\n",
    "    elif simname in [\"CF4\", \"CLONES\"]:\n",
    "        frame = \"supergalactic\"\n",
    "    elif \"csiborg\" in simname:\n",
    "        frame = \"icrs\"\n",
    "    else:\n",
    "        raise ValueError(f\"Unknown frame for simulation: `{simname}`.\")\n",
    "\n",
    "    out = np.full((len(nsims), ngrid, ngrid), np.nan)\n",
    "\n",
    "    \n",
    "    for i, k in enumerate(tqdm(nsims, desc=simname)):\n",
    "        field = get_field(simname, k, \"density\", \"SPH\", 1024)\n",
    "        slice_values = csiborgtools.field.xy_supergalactic_slice(\n",
    "            field, boxsize, xmin, xmax, ngrid, frame)\n",
    "\n",
    "        if simname == \"Lilow2024\":\n",
    "            slice_values[np.isnan(slice_values)] = 1\n",
    "\n",
    "        out[i] = slice_values\n",
    "\n",
    "    return out\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "load_from_disk = True\n",
    "fname = \"/mnt/extraspace/rstiskalek/dump/XY_slices.npz\"\n",
    "\n",
    "xmin = -155\n",
    "xmax = 155\n",
    "ngrid = 1000\n",
    "\n",
    "if not load_from_disk:\n",
    "    xy_carrick = make_slice(\"Carrick2015\", xmin, xmax, ngrid) + 1\n",
    "    xy_lilow = make_slice(\"Lilow2024\", xmin, xmax, ngrid)\n",
    "    xy_CF4 = make_slice(\"CF4\", xmin, xmax, ngrid)\n",
    "    xy_CB2 = make_slice(\"csiborg2_main\", xmin, xmax, ngrid)\n",
    "    xy_CB2X = make_slice(\"csiborg2X\", xmin, xmax, ngrid)\n",
    "    xy_CLONES = make_slice(\"CLONES\", xmin, xmax, ngrid)\n",
    "\n",
    "    np.savez(fname, carrick=xy_carrick, lilow=xy_lilow,\n",
    "             CF4=xy_CF4, CB2=xy_CB2, CB2X=xy_CB2X, CLONES=xy_CLONES)\n",
    "else:\n",
    "    data = np.load(fname)\n",
    "    xy_carrick = data[\"carrick\"]\n",
    "    xy_lilow = data[\"lilow\"]\n",
    "    xy_CF4 = data[\"CF4\"]\n",
    "    xy_CB2 = data[\"CB2\"]\n",
    "    xy_CB2X = data[\"CB2X\"]\n",
    "    xy_CLONES = data[\"CLONES\"]\n",
    "    data.close()\n",
    "\n",
    "\n",
    "\n",
    "rsmooth = None\n",
    "if rsmooth is not None:\n",
    "    sigma = rsmooth / ((xmax - xmin) / ngrid)\n",
    "    print(f\"Smoothing with sigma={sigma}\")\n",
    "\n",
    "    for field in [xy_carrick, xy_CF4, xy_CB2, xy_CB2X]:\n",
    "        for i in range(len(field)):\n",
    "            field[i] = gaussian_filter(field[i], sigma=sigma)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Individual plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 235624 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.figure()\n",
    "\n",
    "img = np.mean(xy_CLONES, axis=0)\n",
    "img = np.clip(img, None, np.percentile(img.ravel(), 99.75))\n",
    "img = np.log10(img)\n",
    "\n",
    "plt.imshow(img, origin=\"lower\", cmap=\"gray_r\", extent=[xmin, xmax, xmin, xmax])\n",
    "\n",
    "kwargs = {\"marker\": \"x\", \"s\": 15}\n",
    "for name, coords in CLUSTERS.items():\n",
    "    plt.scatter(*coords[:2], label=name, **kwargs)\n",
    "\n",
    "plt.scatter(0, 0, label=\"Local Group\", **kwargs)\n",
    "\n",
    "plt.xlabel(r\"$\\mathrm{SGX} ~ [\\mathrm{Mpc} / h]$\")\n",
    "plt.ylabel(r\"$\\mathrm{SGY} ~ [\\mathrm{Mpc} / h]$\")\n",
    "\n",
    "plt.legend(loc='lower center', bbox_to_anchor=(0.5, 1.025),\n",
    "           ncol=4, fontsize=\"small\", frameon=False,\n",
    "           handletextpad=0.1,  # Reduce space between marker and text\n",
    "           borderpad=0.1,      # Reduce padding between legend and border\n",
    "           columnspacing=0.1)  # Reduce space between columns\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.savefig(\"../../plots/slice_test.png\", dpi=450, bbox_inches=\"tight\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### CF4 data check"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 172040 characters]",
      "text/plain": [
       "<Figure size 830x622.5 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "with plt.style.context(\"science\"):\n",
    "    fig, axs = plt.subplots(1, 2, figsize=(8.3, 8.3 * 0.75), sharex=True, sharey=True)\n",
    "\n",
    "    kwargs = {\n",
    "        \"origin\": \"lower\", \"cmap\": \"viridis\", \"extent\": [xmin, xmax, xmin, xmax]}\n",
    "    im0 = axs[0].imshow(xy_CF4[0], **kwargs)\n",
    "    im1 = axs[1].imshow(xy_CF4[0] - xy_CF4[15], **kwargs)\n",
    "\n",
    "    axs[0].set_title(\"Random sample\")\n",
    "    axs[1].set_title(\"Difference of two random samples\")\n",
    "\n",
    "    # Adjust colorbars to be the same height as the plots using make_axes_locatable\n",
    "    labels = [r\"$\\delta$\", r\"$\\Delta \\delta$\"]\n",
    "    for i, im in enumerate([im0, im1]):\n",
    "        divider = make_axes_locatable(axs[i])\n",
    "        cax = divider.append_axes(\"right\", size=\"5%\", pad=0.1)  # Create colorbar axis\n",
    "        fig.colorbar(im, label=labels[i], cax=cax, orientation=\"vertical\")\n",
    "\n",
    "    for i in range(2):\n",
    "        axs[i].set_ylabel(r\"$\\mathrm{SGY} ~ [\\mathrm{Mpc / h}]$\")\n",
    "        axs[i].set_xlabel(r\"$\\mathrm{SGX} ~ [\\mathrm{Mpc / h}]$\")\n",
    "\n",
    "    fig.tight_layout()\n",
    "    fig.savefig(\"../../plots/CF4_test.png\", dpi=450)\n",
    "    fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Paper plots"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. Comparison of fields"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 791728 characters]",
      "text/plain": [
       "<Figure size 830x556.1 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "with plt.style.context(\"science\"):\n",
    "    plt.rcParams.update({'font.size': 9})\n",
    "\n",
    "    imshow_kwargs = {\"origin\": \"lower\", \"cmap\": \"coolwarm\",\n",
    "                     \"extent\": [xmin, xmax, xmin, xmax]}\n",
    "\n",
    "    figwidth = 8.3\n",
    "    fig, axs = plt.subplots(2, 3, figsize=(figwidth, 0.67 * figwidth),\n",
    "                            sharex=\"col\", sharey=True)\n",
    "    fig.subplots_adjust(wspace=0, hspace=0)\n",
    "\n",
    "    # Carrick 2015\n",
    "    img = np.mean(xy_carrick, axis=0)\n",
    "    img = np.clip(img, None, np.percentile(img.ravel(), 99.75))\n",
    "    axs[0, 0].imshow(img, **imshow_kwargs)\n",
    "    axs[0, 0].text(\n",
    "        0.05, 0.05, simname_to_pretty(\"Carrick2015\"), transform=axs[0, 0].transAxes,\n",
    "        bbox=dict(facecolor='white', alpha=0.75))\n",
    "\n",
    "    # Lilow 2024\n",
    "    img = np.mean(xy_lilow, axis=0)\n",
    "    img = np.clip(img, None, np.percentile(img.ravel(), 99.75))\n",
    "    axs[0, 1].imshow(img, **imshow_kwargs)\n",
    "    axs[0, 1].text(\n",
    "        0.05, 0.05, simname_to_pretty(\"Lilow2024\"), transform=axs[0, 1].transAxes,\n",
    "        bbox=dict(facecolor='white', alpha=0.75))\n",
    "\n",
    "    # CF4\n",
    "    img = np.mean(xy_CF4, axis=0)\n",
    "    img = np.clip(img, None, np.percentile(img.ravel(), 99.9))\n",
    "    axs[0, 2].imshow(img, **imshow_kwargs)\n",
    "    axs[0, 2].text(\n",
    "        0.05, 0.05, simname_to_pretty(\"CF4\"), transform=axs[0, 2].transAxes,\n",
    "        bbox=dict(facecolor='white', alpha=0.75))\n",
    "\n",
    "    # csiborg2_main \n",
    "    img = np.mean(xy_CB2, axis=0)\n",
    "    img = np.clip(img, None, np.percentile(img.ravel(), 99))\n",
    "    img = np.log10(img)\n",
    "    axs[1, 0].imshow(img, **imshow_kwargs)\n",
    "    axs[1, 0].text(\n",
    "        0.05, 0.05, simname_to_pretty(\"csiborg2_main\"), transform=axs[1, 0].transAxes,\n",
    "        bbox=dict(facecolor='white', alpha=0.75))\n",
    "\n",
    "    # Manticore\n",
    "    img = np.mean(xy_CB2X, axis=0)\n",
    "    img = np.clip(img, None, np.percentile(img.ravel(), 99))\n",
    "    img = np.log10(img)\n",
    "    axs[1, 1].imshow(img, **imshow_kwargs)\n",
    "    axs[1, 1].text(\n",
    "        0.05, 0.05, simname_to_pretty(\"csiborg2X\"), transform=axs[1, 1].transAxes,\n",
    "        bbox=dict(facecolor='white', alpha=0.75))\n",
    "\n",
    "    # Jenny's field\n",
    "    img = np.mean(xy_CLONES, axis=0)\n",
    "    img = np.clip(img, None, np.percentile(img.ravel(), 99.))\n",
    "    img = np.log10(img)\n",
    "    axs[1, 2].imshow(img, **imshow_kwargs)\n",
    "    axs[1, 2].text(\n",
    "        0.05, 0.05, simname_to_pretty(\"CLONES\"), transform=axs[1, 2].transAxes,\n",
    "        bbox=dict(facecolor='white', alpha=0.75))\n",
    "\n",
    "    for i in range(2):\n",
    "        axs[i, 0].set_ylabel(r\"$\\mathrm{SGY} ~ [\\mathrm{Mpc} / h]$\")\n",
    "\n",
    "    for i in range(3):\n",
    "        axs[-1, i].set_xlabel(r\"$\\mathrm{SGX} ~ [\\mathrm{Mpc} / h]$\")\n",
    "\n",
    "    for i in range(2):\n",
    "        for j in range(3):\n",
    "            axs[i, j].scatter(0, 0, marker=\"x\", s=15, color=\"firebrick\")\n",
    "\n",
    "    fig.tight_layout()\n",
    "    fig.savefig(\"../../plots/XY_slices.pdf\", 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
}
