{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "bbe11a9c-4251-4685-b339-666197ee2aa6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initializing console.\n",
      "[0 / 1] [\u001b[40;0m\u001b[33;1mINFO\u001b[39;0m\u001b[49;0m   ] libLSS version 6b1404bfd8011f03ee1b8635275f68d14d16bef3 built-in modules \n",
      "[0 / 1] [STD    ] Using OneAPI/TBB (default scheduler).\n",
      "[0 / 1] [\u001b[40;0m\u001b[33;1mINFO S\u001b[39;0m\u001b[49;0m ] Registered likelihoods:\n",
      "[0 / 1] [\u001b[40;0m\u001b[33;1mINFO S\u001b[39;0m\u001b[49;0m ]    BORG_NULL, BORG_POISSON, BORG_REDUCED_SHEAR, BORG_SHEAR, BORG_STUDENT_T\n",
      "[0 / 1] [\u001b[40;0m\u001b[33;1mINFO S\u001b[39;0m\u001b[49;0m ]    EFT_BIAS_LIKE, EFT_BIAS_MARG_LIKE, GAUSSIAN_BROKEN_POWER_LAW_BIAS, GAUSSIAN_LINEAR_BIAS\n",
      "[0 / 1] [\u001b[40;0m\u001b[33;1mINFO S\u001b[39;0m\u001b[49;0m ]    GAUSSIAN_LINEAR_BIAS_DEGRADE2, GAUSSIAN_LINEAR_BIAS_DEGRADE4, GAUSSIAN_LINEAR_BIAS_DEGRADE8\n",
      "[0 / 1] [\u001b[40;0m\u001b[33;1mINFO S\u001b[39;0m\u001b[49;0m ]    GAUSSIAN_MO_WHITE_BIAS, GAUSSIAN_POWER_LAW_BIAS, GENERIC_GAUSSIAN_MANY_POWER_1^1\n",
      "[0 / 1] [\u001b[40;0m\u001b[33;1mINFO S\u001b[39;0m\u001b[49;0m ]    GENERIC_GAUSSIAN_MANY_POWER_1^2, GENERIC_POISSON_BROKEN_POWERLAW_BIAS\n",
      "[0 / 1] [\u001b[40;0m\u001b[33;1mINFO S\u001b[39;0m\u001b[49;0m ]    GENERIC_POISSON_BROKEN_POWERLAW_BIAS_DEGRADE2, GENERIC_POISSON_BROKEN_POWERLAW_SIGMOID_BIAS\n",
      "[0 / 1] [\u001b[40;0m\u001b[33;1mINFO S\u001b[39;0m\u001b[49;0m ]    GENERIC_POISSON_MANY_POWER_1^1, GENERIC_POISSON_MANY_POWER_1^2, GENERIC_POISSON_MANY_POWER_1^2_DEGRADE2\n",
      "[0 / 1] [\u001b[40;0m\u001b[33;1mINFO S\u001b[39;0m\u001b[49;0m ]    GENERIC_POISSON_MANY_POWER_2^2, GENERIC_POISSON_MANY_POWER_2^2_DEGRADE8\n",
      "[0 / 1] [\u001b[40;0m\u001b[33;1mINFO S\u001b[39;0m\u001b[49;0m ]    GENERIC_POISSON_MO_WHITE_BIAS, GENERIC_POISSON_POWERLAW_BIAS, GENERIC_POISSON_POWERLAW_BIAS_DEGRADE2\n",
      "[0 / 1] [\u001b[40;0m\u001b[33;1mINFO S\u001b[39;0m\u001b[49;0m ]    HADES_LINEAR, NOOP_POISSON, OCT_LIKELIHOOD, OCT_LIKELIHOOD_MARG, POISSON\n",
      "[0 / 1] [\u001b[40;0m\u001b[33;1mINFO S\u001b[39;0m\u001b[49;0m ]    POISSON_SIGMOID, ROBUST_POISSON_BROKEN_POWERLAW_BIAS, ROBUST_POISSON_BROKEN_POWERLAW_SIGMOID_BIAS\n",
      "[0 / 1] [\u001b[40;0m\u001b[33;1mINFO S\u001b[39;0m\u001b[49;0m ]    ROBUST_POISSON_MANY_POWER_1^1, ROBUST_POISSON_MANY_POWER_1^2, ROBUST_POISSON_POWERLAW_BIAS\n",
      "[0 / 1] [\u001b[40;0m\u001b[33;1mINFO\u001b[39;0m\u001b[49;0m   ] Registering data IO converters\n",
      "[0 / 1] [\u001b[40;0m\u001b[33;1mINFO S\u001b[39;0m\u001b[49;0m ] Registered forward models:\n",
      "[0 / 1] [\u001b[40;0m\u001b[33;1mINFO S\u001b[39;0m\u001b[49;0m ]    2LPT_CIC, 2LPT_CIC_OPENMP, 2LPT_DOUBLE, 2LPT_NGP, ALTAIR_AP, Downgrade\n",
      "[0 / 1] [\u001b[40;0m\u001b[33;1mINFO S\u001b[39;0m\u001b[49;0m ]    EnforceMass, HADES_LOG, HADES_PT, Haar, LPT_CIC, LPT_CIC_OPENMP, LPT_CIC_TBB\n",
      "[0 / 1] [\u001b[40;0m\u001b[33;1mINFO S\u001b[39;0m\u001b[49;0m ]    LPT_DOUBLE, LPT_NGP, LPT_NU_CIC, LPT_NU_CIC_OPENMP, PATCH_MODEL, PM_CIC\n",
      "[0 / 1] [\u001b[40;0m\u001b[33;1mINFO S\u001b[39;0m\u001b[49;0m ]    PM_CIC_OPENMP, PM_CIC_OPENMP2, PM_CIC_TBB, PM_TILE_CIC, PM_TILE_TBB_CIC\n",
      "[0 / 1] [\u001b[40;0m\u001b[33;1mINFO S\u001b[39;0m\u001b[49;0m ]    PRIMORDIAL, PRIMORDIALV2, PRIMORDIAL_AS, PRIMORDIAL_FNL, QLPT, QLPT_RSD\n",
      "[0 / 1] [\u001b[40;0m\u001b[33;1mINFO S\u001b[39;0m\u001b[49;0m ]    Softplus, TRANSFER_CLASS, TRANSFER_EHU, TRANSFER_EHUV2, Transfer, Upgrade\n",
      "[0 / 1] [\u001b[40;0m\u001b[33;1mINFO S\u001b[39;0m\u001b[49;0m ]    bias::BrokenPowerLaw, bias::BrokenPowerLawSigmoid, bias::DoubleBrokenPowerLaw\n",
      "[0 / 1] [\u001b[40;0m\u001b[33;1mINFO S\u001b[39;0m\u001b[49;0m ]    bias::EFT, bias::EFT_Thresh, bias::Gaussian_Linear_Bias_Degrade2, bias::Gaussian_Linear_Bias_Degrade4\n",
      "[0 / 1] [\u001b[40;0m\u001b[33;1mINFO S\u001b[39;0m\u001b[49;0m ]    bias::Gaussian_Linear_Bias_Degrade8, bias::Linear, bias::ManyPower_1^1\n",
      "[0 / 1] [\u001b[40;0m\u001b[33;1mINFO S\u001b[39;0m\u001b[49;0m ]    bias::ManyPower_1^2, bias::ManyPower_1^4, bias::ManyPower_2^2, bias::Noop\n",
      "[0 / 1] [\u001b[40;0m\u001b[33;1mINFO S\u001b[39;0m\u001b[49;0m ]    bias::PowerLaw, bias::SecondOrderBias, bias::Sigmoid\n",
      "[0 / 1] [\u001b[40;0m\u001b[33;1mINFO S\u001b[39;0m\u001b[49;0m ] Registered forward models (v3):\n",
      "[0 / 1] [\u001b[40;0m\u001b[33;1mINFO S\u001b[39;0m\u001b[49;0m ]    KappaSphere\n",
      "[0 / 1] [\u001b[40;0m\u001b[33;1mINFO S\u001b[39;0m\u001b[49;0m ] ARES3 base version 6b1404bfd8011f03ee1b8635275f68d14d16bef3\n",
      "[0 / 1] [\u001b[40;0m\u001b[33;1mINFO S\u001b[39;0m\u001b[49;0m ] GIT report\n",
      "[0 / 1] [\u001b[40;0m\u001b[33;1mINFO S\u001b[39;0m\u001b[49;0m ] - Module: root, Version: 6b1404bfd8011f03ee1b8635275f68d14d16bef3\n",
      "[0 / 1] [\u001b[40;0m\u001b[33;1mINFO S\u001b[39;0m\u001b[49;0m ] \u001b[0m\n",
      "[0 / 1] [\u001b[40;0m\u001b[33;1mINFO S\u001b[39;0m\u001b[49;0m ] \n",
      "[0 / 1] [\u001b[40;0m\u001b[33;1mINFO\u001b[39;0m\u001b[49;0m   ] CPU features: MMX [!AVX] [!AVX2] [!AVX512F] SSE SSE2 SSE3 [!SSE4.1] [!SSE4.2] [!AES] [!BMI] [!BMI2] \n",
      "[0 / 1] [STD    ] Attempting default data path\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-01-31 09:00:11.933354: W external/xla/xla/service/gpu/nvptx_compiler.cc:836] The NVIDIA driver's CUDA version is 12.0 which is older than the PTX compiler version (12.6.77). Because the driver is older than the PTX compiler version, XLA is disabling parallel compilation, which may slow down compilation. You should update your NVIDIA driver or use the NVIDIA-provided CUDA forward compatibility packages.\n"
     ]
    }
   ],
   "source": [
    "import analysis\n",
    "import importlib\n",
    "importlib.reload(analysis)\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import corner\n",
    "import symbolic_pofk.syrenhalofit as syrenhalofit\n",
    "import borg_velocity.utils\n",
    "import aquila_borg as borg\n",
    "import Pk_library as PKL"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05a5d6ae-3eeb-46b9-994d-ef32117e74a1",
   "metadata": {},
   "source": [
    "# Density Field Slice"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f4137a68-a2a5-4ee1-ad4a-aa01eec8a59d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Computing ensemble mean field\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 239/239 [00:15<00:00, 15.61it/s]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 235016 characters]",
      "text/plain": [
       "<Figure size 800x700 with 9 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "borg_dirname = '/data101/bartlett/fsigma8/borg_velocity/velmass_ics'\n",
    "iter_min = 1000\n",
    "iter_max = 5526\n",
    "\n",
    "borg_dirname = '/data101/bartlett/fsigma8/borg_velocity/velmass_ics_model'\n",
    "iter_min = 1000\n",
    "iter_max = 8162\n",
    "\n",
    "borg_dirname = '/data101/bartlett/fsigma8/borg_velocity/velmass_ics_model_cosmo'\n",
    "iter_min = 1000\n",
    "iter_max = 2626\n",
    "\n",
    "borg_dirname = '/data101/bartlett/fsigma8/borg_velocity/velmass_ics_model_cola_v3'\n",
    "iter_min = 0\n",
    "iter_max = -1\n",
    "\n",
    "ini_name = borg_dirname + '/ini_file.ini'\n",
    "field_type = 'delta'\n",
    "nframe = 200\n",
    "cut_field = True\n",
    "cut_field = False\n",
    "\n",
    "true_delta = analysis.crop_velmass_to_borg(ini_name, field_type)\n",
    "if cut_field:\n",
    "    true_delta, _ = analysis.crop_field(ini_name, true_delta)\n",
    "borg_mean, borg_std = analysis.compute_ensemble_mean_field(ini_name, borg_dirname, nframe, iter_max, iter_min, cut_field=cut_field)\n",
    "\n",
    "fig, axs = plt.subplots(3, 3, figsize=(8,7))\n",
    "\n",
    "# Parameter for uncut box\n",
    "Rmax = analysis.get_borg_Rmax(ini_name)\n",
    "L0 = analysis.get_borg_Lbox(ini_name)\n",
    "N0 = analysis.get_borg_N(ini_name)\n",
    "corner0 = analysis.get_borg_corner(ini_name)\n",
    "\n",
    "x = np.linspace(corner0, corner0+L0, N0)\n",
    "if cut_field:\n",
    "    m = np.abs(x) < Rmax\n",
    "    x = x[m]\n",
    "\n",
    "for i, field in enumerate([np.log10(2+true_delta), np.log10(2+borg_mean), borg_std]):\n",
    "    if i == 0:\n",
    "        vmin = field.min()\n",
    "        vmax = field.max()\n",
    "    N = field.shape[0]\n",
    "    j = N//2\n",
    "    axs[i,0].pcolor(x, x, field[j], vmin=vmin, vmax=vmax)\n",
    "    axs[i,1].pcolor(x, x, field[:,j,:], vmin=vmin, vmax=vmax)\n",
    "    axs[i,2].pcolor(x, x, field[:,:,j], vmin=vmin, vmax=vmax)\n",
    "    for ax in axs[i,:]:\n",
    "        circle = plt.Circle((0, 0), Rmax, color='red', fill=False, linestyle='dotted', lw=1.5)\n",
    "        xlim = ax.get_xlim()\n",
    "        ylim = ax.get_ylim()\n",
    "        ax.add_patch(circle)\n",
    "        ax.set_xlim(xlim)\n",
    "        ax.set_ylim(ylim)\n",
    "\n",
    "    for ax in axs[i,:]:\n",
    "        ax.set_aspect('equal')\n",
    "\n",
    "axs[0,1].set_title('Truth')\n",
    "axs[1,1].set_title('Ensemble Mean')\n",
    "axs[2,1].set_title('Ensemble Std')\n",
    "\n",
    "fig.tight_layout()\n",
    "\n",
    "savename = f'../figs/field_slice_velmass_{field_type}'\n",
    "if cut_field:\n",
    "    savename += '_cropped'\n",
    "fig.savefig(f'{savename}.png', bbox_inches='tight', facecolor='white')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f5f59bf-24f1-4b31-a0e7-cdb5a12cdb02",
   "metadata": {},
   "source": [
    "# Power and Cross Spectra"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ac6a9840-dcfc-4d16-98e4-39740656810b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "BOXSIZE 198.41269841269843\n",
      "\n",
      "Computing power spectrum of the field...\n",
      "Time to complete loop = 0.00\n",
      "Time taken = 0.03 seconds\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 110/110 [00:06<00:00, 15.96it/s]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 140960 characters]",
      "text/plain": [
       "<Figure size 1500x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "borg_dirname = '/data101/bartlett/fsigma8/borg_velocity/velmass_ics'\n",
    "iter_min = 0\n",
    "iter_max = 5526\n",
    "\n",
    "borg_dirname = '/data101/bartlett/fsigma8/borg_velocity/velmass_ics_model'\n",
    "iter_min = 0\n",
    "iter_max = 8162\n",
    "\n",
    "# borg_dirname = '/data101/bartlett/fsigma8/borg_velocity/velmass_ics_model_cosmo'\n",
    "# iter_min = 1000\n",
    "# iter_max = 2626\n",
    "\n",
    "borg_dirname = '/data101/bartlett/fsigma8/borg_velocity/velmass_ics_model_cola_v3'\n",
    "iter_min = 0\n",
    "iter_max = -1\n",
    "\n",
    "ini_name = borg_dirname + '/ini_file.ini'\n",
    "field_type = 'delta'\n",
    "# field_type = 'ics'\n",
    "nframe = 100\n",
    "cut_field = True\n",
    "# cut_field = False\n",
    "\n",
    "k, Pk_true, all_pk, all_r = analysis.get_spectra(ini_name, borg_dirname, nframe, iter_max, iter_min, which_field=field_type, cut_field=cut_field, mock_type='velmass', MAS='')\n",
    "all_mcmc = analysis.get_mcmc_steps(borg_dirname, nframe, iter_max, iter_min=iter_min)\n",
    "\n",
    "cpar = borg_velocity.utils.get_cosmopar(ini_name)\n",
    "pk_syrenhalofit = syrenhalofit.run_halofit(k, \n",
    "                       cpar.sigma8, cpar.omega_m, cpar.omega_b, cpar.h, cpar.n_s, 1.0,\n",
    "                       emulator='fiducial', extrapolate=True, which_params='Bartlett', add_correction=True)\n",
    "\n",
    "# Get axis range\n",
    "if field_type == 'delta':\n",
    "    ymin = 1e1\n",
    "    ymax = None\n",
    "else:\n",
    "    ymin = None\n",
    "    ymax = None\n",
    "\n",
    "# Normalize the index for coloring\n",
    "norm = plt.Normalize(all_mcmc[0], all_mcmc[-1])\n",
    "cmap = plt.cm.spring  # Choose a colormap\n",
    "\n",
    "fig, axs = plt.subplots(1, 2, figsize=(15,4))\n",
    "\n",
    "# Plot Pk and r\n",
    "for i in range(all_pk.shape[0]):\n",
    "    axs[0].loglog(k, all_pk[i, :], color=cmap(norm(all_mcmc[i])))\n",
    "    axs[1].semilogx(k, all_r[i, :], color=cmap(norm(all_mcmc[i])))\n",
    "axs[0].loglog(k, Pk_true, color='k', label='True')\n",
    "axs[0].loglog(k, pk_syrenhalofit, color='r', label='syren')\n",
    "for i in [-1, 0, 1]:\n",
    "    axs[1].axhline(i, color='k')\n",
    "axs[0].legend()\n",
    "\n",
    "sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)\n",
    "sm.set_array([])\n",
    "for ax in axs:\n",
    "    cbar = plt.colorbar(sm, ax=ax)\n",
    "    cbar.set_label('Step')\n",
    "    ax.set_title(field_type)\n",
    "    ax.set_xlabel(r'$k$')\n",
    "axs[0].set_ylim(ymin, ymax)\n",
    "axs[0].set_ylabel(r'$P(k)$')\n",
    "axs[1].set_ylabel(r'$r(k)$')\n",
    "fig.tight_layout()\n",
    "savename = f'../figs/pk_velmass_{field_type}'\n",
    "if cut_field:\n",
    "    savename += '_cropped'\n",
    "fig.savefig(f'{savename}.png', bbox_inches='tight', facecolor='white')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9bab5ce8-3171-41bf-bd97-30921e77664f",
   "metadata": {},
   "source": [
    "# Model Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "617d573a-03d2-4bea-839c-2b283ea39ac8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TO SAMPLE ['mua', 'alpha', 'lam', 'sig_v', 'bulk_flow']\n",
      "MY SAMPLER IS transformedblackjax\n",
      "['mua0', 'mua1', 'mua2', 'mua3', 'alpha0', 'alpha1', 'alpha2', 'alpha3', 'lam0', 'lam1', 'lam2', 'lam3', 'sig_v', 'bulk_flow_x', 'bulk_flow_y', 'bulk_flow_z']\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1190/1190 [00:01<00:00, 877.13it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['mua0', 'mua1', 'mua2', 'mua3', 'alpha0', 'alpha1', 'alpha2', 'alpha3', 'lam0', 'lam1', 'lam2', 'lam3', 'sig_v', 'bulk_flow_x', 'bulk_flow_y', 'bulk_flow_z']\n",
      "(16, 1190)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 519584 characters]",
      "text/plain": [
       "<Figure size 2000x1200 with 16 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 4205704 characters]",
      "text/plain": [
       "<Figure size 3490x3490 with 256 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "borg_dirname = '/data101/bartlett/fsigma8/borg_velocity/velmass_ics_model'\n",
    "iter_min = 0\n",
    "iter_max = 8162\n",
    "\n",
    "borg_dirname = '/data101/bartlett/fsigma8/borg_velocity/velmass_ics_model_cosmo'\n",
    "iter_min = 0\n",
    "iter_max = 2626\n",
    "\n",
    "borg_dirname = '/data101/bartlett/fsigma8/borg_velocity/velmass_ics_model_cola_v3'\n",
    "iter_min = 0\n",
    "iter_max = -1\n",
    "\n",
    "nframe = 2000\n",
    "# burn_in = 1000\n",
    "burn_in = 200\n",
    "\n",
    "ini_name = borg_dirname + '/ini_file.ini'\n",
    "param_names, all_mcmc, samples = analysis.load_param_samples(ini_name, borg_dirname, nframe, iter_max, iter_min)\n",
    "truths = analysis.get_truths(ini_name, param_names)\n",
    "\n",
    "print(param_names)\n",
    "print(samples.shape)\n",
    "\n",
    "# For Velmass some of these are not in fact truths\n",
    "for i, p in enumerate(param_names):\n",
    "    if p.startswith('alpha') or p.startswith('sig_v') or p.startswith('bulk_flow'):\n",
    "        truths[i] = None\n",
    "\n",
    "# print(param_names)\n",
    "# print(samples.shape)\n",
    "# for i, n in enumerate(param_names):\n",
    "#     if n.startswith('lam'):\n",
    "#         truths[i] = 1 / truths[i]\n",
    "#         samples[i,:] = 1 / samples[i,:]\n",
    "#         m = samples[i,:] < 1\n",
    "#         samples = samples[:,m]\n",
    "#         all_mcmc = all_mcmc[m]\n",
    "\n",
    "ncol = 4\n",
    "nrow = int(np.ceil(samples.shape[0] / ncol))\n",
    "fig, axs = plt.subplots(nrow, ncol, figsize=(ncol * 5, nrow * 3))\n",
    "axs = axs.flatten()\n",
    "for i in range(len(samples)):\n",
    "    ax = axs[i]\n",
    "    ax.plot(all_mcmc, samples[i])\n",
    "    ax.set_title(param_names[i])\n",
    "    if truths[i] is not None:\n",
    "        ax.axhline(truths[i], color='k')\n",
    "for i in range(len(samples), len(axs)):\n",
    "    axs[i].clear()\n",
    "    axs[i].remove()\n",
    "fig.tight_layout()\n",
    "fig.savefig('../figs/trace_velmass.png', bbox_inches='tight', facecolor='white')\n",
    "\n",
    "m = all_mcmc > burn_in\n",
    "\n",
    "# Full corner plot\n",
    "corner.corner(\n",
    "    samples[:,m].T,\n",
    "    labels=param_names,\n",
    "    truths=truths\n",
    ")\n",
    "fig = plt.gcf()\n",
    "fig.savefig('../figs/corner_velmass.png', bbox_inches='tight', facecolor='white')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e02c490-cddb-4ea3-817f-dfeebf209c73",
   "metadata": {},
   "source": [
    "# Check VELMASS Power Spectrum against Emulator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "7de0fa2b-de30-40a8-8cb2-891f06e7b598",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<CosmologicalParameters: omega_r=0, omega_k=0, omega_m=0.315, omega_b=0.049, omega_q=0.685, w=-1, n_s=0.97, fnl=0, wprime=0, sigma8=0.81, h=0.68, sum_mnu=0 eV, As=1.96592e-09>\n",
      "2000.0 500.0 500.0 96 160 256\n",
      "-0.9290850623568685 28.859140827923998 9.489479514679378e-17\n",
      "Lbox: 500.0\n",
      "\n",
      "Computing power spectrum of the field...\n",
      "Time to complete loop = 0.01\n",
      "Time taken = 0.01 seconds\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 78124 characters]",
      "text/plain": [
       "<Figure size 700x1000 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "borg_dirname = '/data101/bartlett/fsigma8/borg_velocity/velmass_ics_model'\n",
    "ini_name = borg_dirname + '/ini_file.ini'\n",
    "\n",
    "cpar = borg_velocity.utils.get_cosmopar(ini_name)\n",
    "print(cpar)\n",
    "\n",
    "# k, Pk_true, _, _ = analysis.get_spectra(ini_name, borg_dirname, 1., -1, 0, which_field='delta', cut_field=False, mock_type='velmass')\n",
    "delta1 = analysis.get_velmass_field(ini_name, 'delta')\n",
    "boxsize = analysis.get_velmass_Lbox(ini_name)\n",
    "\n",
    "Lbox_true = analysis.get_velmass_Lbox(ini_name)\n",
    "Lbox_borg = analysis.get_borg_Lbox(ini_name)\n",
    "N = delta1.shape[0]\n",
    "imin = int((1 - Lbox_borg / Lbox_true) * N / 2)\n",
    "imax = int((1 + Lbox_borg / Lbox_true) * N / 2)\n",
    "Lbox_new = (imax - imin) / N * Lbox_true\n",
    "print(Lbox_true, Lbox_borg, Lbox_new, imin, imax, N)\n",
    "boxsize = Lbox_new\n",
    "delta1 = delta1[imin:imax, imin:imax, imin:imax]\n",
    "# delta1 = delta1[:imax-imin, :imax-imin, :imax-imin]\n",
    "delta1 = (1 + delta1) / (1 + delta1).mean() - 1.\n",
    "print(delta1.min(), delta1.max(), delta1.mean())\n",
    "\n",
    "print(\"Lbox:\", boxsize)\n",
    "MAS = \"CIC\"\n",
    "Pk = PKL.Pk(delta1.astype(np.float32), boxsize, axis=0, MAS=MAS, threads=1, verbose=True)\n",
    "k      = Pk.k3D\n",
    "Pk_true     = Pk.Pk[:,0]\n",
    "\n",
    "\n",
    "\n",
    "k_max, k_per_decade = 10, 100\n",
    "extra_class = {}\n",
    "extra_class['YHe'] = '0.2458'\n",
    "cosmo = borg.cosmo.ClassCosmo(cpar, k_per_decade, k_max, extra=extra_class)\n",
    "\n",
    "# Compute integrand\n",
    "plin_class = cosmo.get_Pk_matter(k)\n",
    "\n",
    "pk_syrenhalofit = syrenhalofit.run_halofit(k, \n",
    "                       cpar.sigma8, cpar.omega_m, cpar.omega_b, cpar.h, cpar.n_s, 1.0,\n",
    "                       emulator='fiducial', extrapolate=True, which_params='Bartlett', add_correction=True)\n",
    "\n",
    "fig, axs = plt.subplots(2, 1, figsize=(7,10), sharex=True)\n",
    "axs[0].loglog(k, Pk_true, label='VELMASS')\n",
    "axs[0].loglog(k, pk_syrenhalofit, label='syren')\n",
    "axs[0].loglog(k, plin_class, label='linear (CLASS)')\n",
    "axs[1].semilogx(k, Pk_true / plin_class, label='VELMASS')\n",
    "axs[1].semilogx(k, pk_syrenhalofit / plin_class, label='syren')\n",
    "axs[1].semilogx(k, plin_class / plin_class, label='linear (CLASS)')\n",
    "axs[0].legend()\n",
    "axs[0].set_ylabel(r'$P(k)$')\n",
    "axs[1].set_ylabel(r'$P(k) / P_{\\rm lin}(k)$')\n",
    "axs[1].set_xlabel(r'$k$')\n",
    "axs[1].set_ylim(0.5, 2)\n",
    "fig.align_labels()\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "179592db-996d-4292-9350-3a0711048cb0",
   "metadata": {},
   "source": [
    "# TO DO\n",
    "\n",
    "* Looks like VELMASS ICs are after P(k) applied - extract white noise field"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c0c7b33-89a1-4926-8354-1e8f475d756d",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "borg_new",
   "language": "python",
   "name": "borg_new"
  },
  "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.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
