{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "bb3745ee",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Quick-start\n",
    "In this tutorial we will show \n",
    "  * How to initialise the emulator.\n",
    "  * How to obtain multipoles for the standard $\\Lambda$CDM cosmology."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10fb98ee",
   "metadata": {},
   "source": [
    "Let’s first import `comet` as well as other required libraries:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "cd57a731",
   "metadata": {},
   "outputs": [],
   "source": [
    "from comet import comet\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6746c3d2",
   "metadata": {},
   "source": [
    "At initialisation we only need to specify the perturbation theory model that we want to use (valid specifiers are currently either `\"EFT\"` or `\"RS\"`; for an overview of the models implemented in COMET, see here) and we can configure COMET either in $\\mathrm{Mpc}$ units (`use_Mpc = True`, which is the default option) or in $h^{-1}\\mathrm{Mpc}$ units (`use_Mpc = False`). All quantities that are not dimensionless are then returned or assumed to be given in the respective unit system. Let’s define an emulator object for the EFT model using the standard $h^{-1}\\mathrm{Mpc}$ units:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a045c0b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "EFT=comet(model=\"EFT\", use_Mpc=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91a1d037",
   "metadata": {},
   "source": [
    "In order to make predictions for a given cosmological model we first need to specify the fiducial background cosmology, from which the Alcock-Paczynski distortions will be computed. This is done by calling the function `define_fiducial_cosmology` with a dictionary specifying the cosmological parameters and the redshift:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "3d73efd3",
   "metadata": {},
   "outputs": [],
   "source": [
    "params_fid = {'h':0.695, 'wc':0.11544, 'wb':0.0222191, 'z':0.57}\n",
    "\n",
    "# This assumes by default a \"lambda\" cosmology with w0 = -1, for other\n",
    "# options, see the in-depth examples below.\n",
    "EFT.define_fiducial_cosmology(params_fid=params_fid)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32c1514d",
   "metadata": {},
   "source": [
    "The function `Pell`, which returns the power spectrum multipoles takes generally three parameters:\n",
    "\n",
    "1. The scales for which to compute the multipoles (in the corresponding units)\n",
    "2. A parameter dictionary, specifying cosmological, bias, and (if applicable) additional redshift-space distortions parameters\n",
    "3. The multipole number, i.e. ell = 0, 2, 4, or a list of multipole numbers\n",
    "\n",
    "The parameter dictionary must include all shape parameters: the physical cold dark matter and baryon densities (`wc` and `wb`) and the scalar spectral index (`ns`). In case of a flat $\\Lambda$CDM model we also need to specify values for $h$ (`h`), the amplitude of scalar fluctuations (`As`) and redshift (`z`). For other cosmologies, see In-depth options for obtaining multipoles."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d82510bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's create a parameter dictionary\n",
    "params = {}\n",
    "\n",
    "# We always need to specify the shape parameter values, e.g.\n",
    "params['wc'] = 0.11544\n",
    "params['wb'] = 0.0222191\n",
    "params['ns'] = 0.9632\n",
    "\n",
    "# For a LCDM cosmology, we also need:\n",
    "params['h']  = 0.8\n",
    "params['As'] = 2.3\n",
    "params['z']  = 0.6"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3e19a8c",
   "metadata": {},
   "source": [
    "Finally, we define the values of the bias parameters. The complete list of parameters along with a brief explanation and their dictionary keywords can be found here. In the following we only specify values for the linear and quadratic bias, all other parameters are automatically set to zero:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "69d24be9",
   "metadata": {},
   "outputs": [],
   "source": [
    "params['b1'] = 2.\n",
    "params['b2'] = -0.5"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "416a12f5",
   "metadata": {},
   "source": [
    "Now, let’s compute the monopole (`ell=0`), quadrupole (`ell=2`) and hexadecapole (`ell=4`) for a range of scales from $0.001 h\\,\\mathrm{Mpc}^{−1}$ to $0.3h\\,\\mathrm{Mpc}^{−1}$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b45988ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "k_hMpc = np.logspace(-3,np.log10(0.3),100)\n",
    "Pell_LCDM = EFT.Pell(k_hMpc, params, ell=[0,2,4], de_model='lambda')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4fc85a38",
   "metadata": {},
   "source": [
    "The output of `Pell` is given in a dictionary format:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "e6cdf925",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dict_keys(['ell0', 'ell2', 'ell4'])\n"
     ]
    }
   ],
   "source": [
    "print(Pell_LCDM.keys())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1a9a6d2",
   "metadata": {},
   "source": [
    "So we can access our results and plot them as follow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "85576c9e-5bda-4d92-abda-a0e514928731",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 59456 characters]\n",
      "text/plain": [
       "<Figure size 1000x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "f = plt.figure(figsize=(10,5))\n",
    "ax = f.add_subplot(111)\n",
    "ax.semilogx(k_hMpc, k_hMpc**0.5*Pell_LCDM[\"ell0\"],c='C0',ls='-',label='$P_0$')\n",
    "ax.semilogx(k_hMpc, k_hMpc**0.5*Pell_LCDM[\"ell2\"],c='C1',ls='-',label='$P_2$')\n",
    "ax.semilogx(k_hMpc, k_hMpc**0.5*Pell_LCDM[\"ell4\"],c='C2',ls='-',label='$P_4$')\n",
    "ax.set_xlabel('$k$ [h/Mpc]',fontsize=15)\n",
    "ax.set_ylabel(r'$k^{1/2}\\,P_{\\ell}(k)$ [$(\\mathrm{Mpc}/h)^{5/2}$]',fontsize=15)\n",
    "ax.legend(fontsize=15)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0a7299d",
   "metadata": {},
   "source": [
    "# Exploring a few in-depth options"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13d93468",
   "metadata": {},
   "source": [
    "Let us now consider some of the more detailed options in 'COMET':\n",
    "   * Specifying fiducial background cosmologies\n",
    "   * Specifying Alcock-Paczynski parameters\n",
    "   * Specifying the shot noise normalisation\n",
    "   * Non-flat and non-$\\Lambda$ cosmologies\n",
    "   * Using the $f$-$\\sigma_{12}$ parameter space\n",
    "   * Options for providing different $k$-scales, float vs np.array vs list and the corresponding outputs\n",
    "   * Description of the `fixed_cosmo_boost` function, i.e., speedup when just changing bias parameters\n",
    "   * Using different bases for galaxy bias"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0189f509-fe4b-4141-a1fb-3bf9b1c96ae2",
   "metadata": {},
   "source": [
    "## Fiducial background cosmologies"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca41cc6c-6608-4e92-8eac-e241c3643a52",
   "metadata": {},
   "source": [
    "Above, we specified the fiducial background cosmology by setting the values of $h$, $\\omega_b$, $\\omega_c$ and redshift $z$. Alternatively, we can directly provide the values of the Hubble rate $H_{\\rm fid}(z)$ and comoving transverse distance $D_{m,\\rm fid}(z)$ as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "3ac5069f-63ef-4b50-a850-0998e5d00de9",
   "metadata": {},
   "outputs": [],
   "source": [
    "H_fid = 135    # in units of km/s/(Mpc/h)\n",
    "Dm_fid = 1490  # in units of Mpc/h\n",
    "\n",
    "EFT.define_fiducial_cosmology(HDm_fid=[H_fid, Dm_fid])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "005bf690-527e-472d-8fe2-802c266a36ff",
   "metadata": {},
   "source": [
    "Note that the units of $H_{\\rm fid}(z)$ and $D_{m,\\rm fid}(z)$ need to be either in $\\mathrm{km}\\,\\mathrm{s}^{-1}\\,\\mathrm{Mpc}^{-1}$ and $\\mathrm{Mpc}$ (if `use_Mpc=True`), or $\\mathrm{km}\\,\\mathrm{s}^{-1}\\,(h^{-1}\\mathrm{Mpc})^{-1}$ and $h^{-1}\\mathrm{Mpc}$ (if `use_Mpc=False`)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d36a70d-4c03-4262-8d0c-c31029f9c451",
   "metadata": {},
   "source": [
    "Moreover, we stress that `define_fiducial_cosmology` is only used to set the fiducial cosmological parameter values. It cannot be used to set default parameter values for the evaluation of the model."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c935f304-27cf-4604-ae5c-d0f31b50904a",
   "metadata": {},
   "source": [
    "## Alcock-Paczynski parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06f96453-a6a7-4088-be4b-07aa1c3cdb9b",
   "metadata": {},
   "source": [
    "By default the values of the Alcock-Paczynski parameters, $q_{\\parallel}$ and $q_{\\perp}$, are computed based on the given cosmological parameters and the fiducial background values for the Hubble rate and comoving transverse distance. These values can be overwritten by explicitly providing the Alcock-Paczynski parameters as an argument to the `Pell` function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "54fe838a-e672-4649-ab19-abad2a4a2d94",
   "metadata": {},
   "outputs": [],
   "source": [
    "q_para = 1.0\n",
    "q_perp = 1.0\n",
    "\n",
    "Pell_LCDM_noAP = EFT.Pell(k_hMpc, params, ell=[0,2,4], de_model='lambda', q_tr_lo=[q_perp,q_para])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3fba07c0-1d67-4967-82ce-15c55d831787",
   "metadata": {},
   "source": [
    "This can be useful when one would like to ignore Alcock-Paczynski distortions."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "499f5882-288b-4581-baee-907c42c6152a",
   "metadata": {},
   "source": [
    "## Shot noise normalisation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b604308-a739-4aea-9de8-39f2eeea2908",
   "metadata": {},
   "source": [
    "By default the shot noise parameters in the power spectrum model are assumed to be given in units of $L^3$ for `NP0` and $L^5$ for `NP20` and `NP22`, where $L = (\\mathrm{Mpc})^3$ (`use_Mpc=True`) or $L = (h^{-1}\\mathrm{Mpc})^3$ (`use_Mpc=False`). It is possible to define a fixed normalisation scale (i.e., corresponding to the Poisson shot noise $1/\\bar{n}$) as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "e83492e2-483b-4f20-8488-c725a720e631",
   "metadata": {},
   "outputs": [],
   "source": [
    "nbar = 1e-3  # in the respective units\n",
    "EFT.define_nbar(nbar)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31a43bb5-f150-4371-bbe1-8d246d3e0dce",
   "metadata": {},
   "source": [
    "In this case `NP0` is dimensionless, while `NP20` and `NP22` have dimension $L^2$."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b45b0612",
   "metadata": {},
   "source": [
    "## Non-flat and non-$\\Lambda$ cosmologies"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80a7d393-af33-4cc2-b653-e786d360e35d",
   "metadata": {},
   "source": [
    "Predictions for non-flat cosmologies can be obtained by simply specifying the curvature density parameter $\\Omega_k$ in the parameter dictionary:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "56e51af9-ec74-4a42-a948-aaa691c12493",
   "metadata": {},
   "outputs": [],
   "source": [
    "params['Ok'] = 0.05"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95f98c3e",
   "metadata": {},
   "source": [
    "For different dark energy models we need to provide a different `de_model` argument for the `Pell` function. For a non-time varying dark energy equation of state, we set `de_model='w0'`, while for a time-varying equation of state in the $w_0$-$w_a$ parametrisation, we set `de_model='w0wa'`. In those cases we need to specify the corresponding values of $w_0$ and $w_a$ in the parameter dictionary. Let's consider the following example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "40915127",
   "metadata": {},
   "outputs": [],
   "source": [
    "params['w0'] = -1.1\n",
    "params['wa'] = 0.1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53c02449",
   "metadata": {},
   "source": [
    "Then let's recompute the model by updating the previously set parameter values and compare with the $\\Lambda$CDM prediction:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "3a9be672",
   "metadata": {},
   "outputs": [],
   "source": [
    "Pell_w0wa = EFT.Pell(k_hMpc, params, ell=[0,2,4], de_model='w0wa')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "6a1889fb-2ff1-4adc-9d65-ba2100e9d594",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 99788 characters]\n",
      "text/plain": [
       "<Figure size 1000x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "f = plt.figure(figsize=(10,5))\n",
    "ax = f.add_subplot(111)\n",
    "ax.semilogx(k_hMpc, k_hMpc**0.5*Pell_LCDM[\"ell0\"],c='C0',ls='-',label='$P_0$, $\\Lambda$CDM')\n",
    "ax.semilogx(k_hMpc, k_hMpc**0.5*Pell_LCDM[\"ell2\"],c='C1',ls='-',label='$P_2$, $\\Lambda$CDM')\n",
    "ax.semilogx(k_hMpc, k_hMpc**0.5*Pell_LCDM[\"ell4\"],c='C2',ls='-',label='$P_4$, $\\Lambda$CDM')\n",
    "ax.semilogx(k_hMpc, k_hMpc**0.5*Pell_w0wa[\"ell0\"],c='C0',ls='--',label='$P_0$, $w_0 w_a$CDM')\n",
    "ax.semilogx(k_hMpc, k_hMpc**0.5*Pell_w0wa[\"ell2\"],c='C1',ls='--',label='$P_2$, $w_0 w_a$CDM')\n",
    "ax.semilogx(k_hMpc, k_hMpc**0.5*Pell_w0wa[\"ell4\"],c='C2',ls='--',label='$P_4$, $w_0 w_a$CDM')\n",
    "ax.set_xlabel('$k$ [h/Mpc]',fontsize=15)\n",
    "ax.set_ylabel(r'$k^{1/2}\\,P_{\\ell}(k)$ [$(\\mathrm{Mpc}/h)^{5/2}$]',fontsize=15)\n",
    "ax.legend(fontsize=15)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea58fc4b",
   "metadata": {},
   "source": [
    "## The $f$-$\\sigma_{12}$ parameter space"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da869351",
   "metadata": {},
   "source": [
    "When calling the `Pell` function for a specific dark energy model, it ignores any potential values of `s12`, `q_tr`, `q_lo` and `f` in the parameter dictionary and instead converts the $\\Lambda$CDM parameters to the $\\sigma_{12}$ parameter space. The internal values of those parameters (which can be accessed via `EFT.params`) have therefore been updated:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "62af9831",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'wc': 0.11544,\n",
       " 'wb': 0.0222191,\n",
       " 'ns': 0.9632,\n",
       " 's12': array([0.56448159]),\n",
       " 'f': 0.7025465611424653,\n",
       " 'b1': 2.0,\n",
       " 'b2': -0.5,\n",
       " 'g2': 0.0,\n",
       " 'g21': 0.0,\n",
       " 'c0': 0.0,\n",
       " 'c2': 0.0,\n",
       " 'c4': 0.0,\n",
       " 'cnlo': 0.0,\n",
       " 'cnloB': 0.0,\n",
       " 'NP0': 0.0,\n",
       " 'NP20': 0.0,\n",
       " 'NP22': 0.0,\n",
       " 'NB0': 0.0,\n",
       " 'MB0': 0.0,\n",
       " 'cB1': 0.0,\n",
       " 'cB2': 0.0,\n",
       " 'h': 0.8,\n",
       " 'As': 2.3,\n",
       " 'Ok': 0.05,\n",
       " 'w0': -1.1,\n",
       " 'wa': 0.1,\n",
       " 'z': 0.6,\n",
       " 'q_tr': 1.081799699202137,\n",
       " 'q_lo': 1.045999542223697}"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# s12, q_tr, q_lo and f are computed internally!\n",
    "EFT.params"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e6d7e64",
   "metadata": {},
   "source": [
    "If we want to use the $f$-$\\sigma_{12}$ parameter space directly, we need to provide explicit values for `s12`, `f`, `q_lo` ($q_{\\parallel}$) and `q_tr` ($q_{\\perp}$). As an example, let's redefine our parameter values:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "7ea6c3e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# For predictions using the RSD parameter space we also need to specify values for the following four parameters, e.g.\n",
    "params['s12']  = 0.6\n",
    "params['q_lo'] = 1.1\n",
    "params['q_tr'] = 0.9\n",
    "params['f']    = 0.7"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9586cc2f",
   "metadata": {},
   "source": [
    "_Note: When computing the multipoles using the $\\sigma_{12}$ parameter space and in $h^{-1}\\mathrm{Mpc}$ units, we need to specify a fiducial value for the Hubble rate (provided in the parameter dictionary). This is required to convert the native emulator output from Mpc to Mpc/h units._"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "b6bdab99",
   "metadata": {},
   "outputs": [],
   "source": [
    "Pell_s12 = EFT.Pell(k_hMpc, params, ell=[0,2,4])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "d365a99f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 85208 characters]\n",
      "text/plain": [
       "<Figure size 1000x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "f = plt.figure(figsize=(10,5))\n",
    "ax = f.add_subplot(111)\n",
    "ax.semilogx(k_hMpc, k_hMpc**0.5*Pell_LCDM[\"ell0\"],c='C0',ls='-',label='P0')\n",
    "ax.semilogx(k_hMpc, k_hMpc**0.5*Pell_LCDM[\"ell2\"],c='C1',ls='-',label='P2')\n",
    "ax.semilogx(k_hMpc, k_hMpc**0.5*Pell_LCDM[\"ell4\"],c='C2',ls='-',label='P4')\n",
    "ax.semilogx(k_hMpc, k_hMpc**0.5*Pell_s12[\"ell0\"],c='C0',ls='--')\n",
    "ax.semilogx(k_hMpc, k_hMpc**0.5*Pell_s12[\"ell2\"],c='C1',ls='--')\n",
    "ax.semilogx(k_hMpc, k_hMpc**0.5*Pell_s12[\"ell4\"],c='C2',ls='--')\n",
    "ax.set_xlabel('$k$ [h/Mpc]',fontsize=15)\n",
    "ax.set_ylabel(r'$k^{1/2}\\,P_{\\ell}(k)$ [$(\\mathrm{Mpc}/h)^{5/2}$]',fontsize=15)\n",
    "ax.legend(fontsize=15)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd0adcd0",
   "metadata": {},
   "source": [
    "## Providing different $k$-scales"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dda3c415",
   "metadata": {},
   "source": [
    "There are multiple options for specifying the scales for which to compute the multipoles: if given as a number or numpy array all specified multipoles will be computed for those scales, if given as a list, however, then the first entry of the list is evaluated for the first multipole, the second for the second multipole, etc.\n",
    "\n",
    "We can output at a single scale and single multipole number, e.g. for the quadrupole at $k = 0.1\\,h\\,\\mathrm{Mpc}^{-1}$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "bb6b72dc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'ell2': array([12734.89063923])}"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "EFT.Pell(0.1, params, ell=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c1ca966",
   "metadata": {},
   "source": [
    "Or for various multipoles and multiple scales:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "97b5881a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'ell0': array([21992.53270038,  8421.94731931,  5056.16492358]),\n",
       " 'ell2': array([12734.89063923,  7164.45814763,  5358.95648909]),\n",
       " 'ell4': array([3027.96053502, 2241.88093571, 1868.24333838])}"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "EFT.Pell(np.array([0.1,0.2,0.3]), params, ell=[0,2,4])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96760b59",
   "metadata": {},
   "source": [
    "Or at different scales for different multipoles (providing a list of numbers or numpy arrays):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "7eee3ba8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'ell0': array([21992.53270038,  8421.94731931]),\n",
       " 'ell4': array([1868.24333838])}"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "EFT.Pell([np.array([0.1,0.2]),0.3], params, ell=[0,4])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35fc682d-cd91-4290-836e-2d025e10d06e",
   "metadata": {},
   "source": [
    "_Note: In this case, the length of the list must match the length of the specified multipoles (`ell`)._"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a4ac0af-77ec-409e-8411-3c110f7d0112",
   "metadata": {},
   "source": [
    "Performance-wise it is advisable to compute all required multipoles and scales via the same function call (i.e., avoid calling `Pell` for individual wavemodes)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a8e77e5",
   "metadata": {},
   "source": [
    "## Speedup when changing just bias parameters."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1fdb6727",
   "metadata": {},
   "source": [
    "It is a common task to test the models at fixed cosmological parameters, and in that case COMET provides the function `Pell_fixed_cosmo_boost`, which accelerates the model computation. It computes all individual model contributions, which are kept fixed as long as the cosmological parameters are not changed, such that changing the bias parameters only is sped up drastically. In the following cells the diferences on time can be seen, which reflects a speed up of around 3 orders of magnitude."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "32d70c8b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "5.17 ms ± 33 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit EFT.Pell(k_hMpc, params, ell=[0,2,4], de_model=\"lambda\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "be5a631f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "9.45 µs ± 23.2 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit EFT.Pell_fixed_cosmo_boost(k_hMpc, params, ell=[0,2,4], de_model=\"lambda\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c9ad866-7345-400f-baca-f048bb17bcb0",
   "metadata": {},
   "source": [
    "_Note: Since the computation of all the individual contributions takes more time than the direct evaluation of the multipoles, this is really only useful at fixed cosmological parameters._"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d183037-a8cc-407e-b7d8-dccc32a9fe38",
   "metadata": {},
   "source": [
    "## Using different bases for galaxy bias"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4631aa3-a0f5-4385-814c-1e0940497360",
   "metadata": {},
   "source": [
    "By default COMET uses the galaxy bias expansion proposed in Eggemeier et al. (2019), but it is also possible to specify bias parameters of two other bases from:\n",
    "  * Assassi et al. (2014), used e.g. in the analysis by Ivanov et al. (2019)\n",
    "  * d'Amico et al. (2019)\n",
    "  \n",
    "The bias basis is defined at initialisation using the argument `bias_basis`, which can take the strings `\"EggScoSmi\"` (for the Eggemeier et al. basis), `\"AssBauGre\"` (for the Assassi et al. basis), or `\"AmiGleKok\"` (for the D'Amico et al. basis). It is also possible to change the bias basis later via the function `change_bias_basis`, e.g.:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "f67e9f17-2a8b-4779-8bcd-8ccbcdc0823e",
   "metadata": {},
   "outputs": [],
   "source": [
    "EFT.change_bias_basis(\"AssBauGre\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "091f255b-73e0-4b03-927c-7f5fe5067fcd",
   "metadata": {},
   "source": [
    "Changing the bias basis changes the parameter dictionary keys that need to be provided. The full list of available bias keys can be printed as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "51863649-6cdb-4b63-a040-acc6691836e2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['b1', 'b2', 'bG2', 'bGam3', 'c0', 'c2', 'c4', 'cnlo', 'NP0', 'NP20', 'NP22', 'NB0', 'MB0']\n"
     ]
    }
   ],
   "source": [
    "print(EFT.bias_params_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eadbd011-74ef-4c3a-a801-9e3e245b8d04",
   "metadata": {},
   "source": [
    "In this case we now need to provide values for `'bG2'` and `'bGam3'`, i.e., parameters for `'g2'` and `'g21'` are now ignored. In case of the d'Amico et al. basis we have:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "3c5150a4-55dc-4f20-945b-fda940448dbc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['b1t', 'b2t', 'b3t', 'b4t', 'c0', 'c2', 'c4', 'cnlo', 'NP0', 'NP20', 'NP22', 'NB0', 'MB0']\n"
     ]
    }
   ],
   "source": [
    "EFT.change_bias_basis(\"AmiGleKok\")\n",
    "print(EFT.bias_params_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b1a8aa4-8400-4804-b1be-6fc1ecd8a095",
   "metadata": {},
   "source": [
    "Let's change back to the default for the remainder of the tutorial:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "0b02a359-3c14-42dc-a156-e3f2746a2d04",
   "metadata": {},
   "outputs": [],
   "source": [
    "EFT.change_bias_basis(\"EggScoSmi\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5152156c-aca3-4b75-a2de-803863acf28d",
   "metadata": {},
   "source": [
    "# Binning and discreteness effects"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3bed43bd-ba90-4f33-b40a-4ddf83048260",
   "metadata": {},
   "source": [
    "## Power spectrum"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e069cd6-6df5-4edb-ac07-580ab5c940b2",
   "metadata": {},
   "source": [
    "Power spectrum multipoles are estimated in Fourier space from discrete grids of wave vectors, which means that a given multipole at scale $k$ is an average over the discrete set of wave vectors $\\mathbf{q}$ whose magnitude falls into the spherical shell defined by $k - \\Delta k/2 \\leq |\\mathbf{q}| \\leq k + \\Delta k/2$. This leads to differences from the theory predictions, which (per default) assume continuous wave vectors and infinitesimally thin shells ($\\Delta k \\to 0$). However, the discreteness and finite bin width effects can be accounted for by averaging the anisotropic theory power spectrum over the same set of modes as those that are averaged over when performing the measurements."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84074562-fa8c-494a-aa31-739e91f6e4cd",
   "metadata": {},
   "source": [
    "In COMET this can be done by specifying a binning dictionary, when calling `Pell` or `Pell_fixed_cosmo_boost`. In order to compute the set of discrete modes, it is necessary to know the size (i.e., the fundamental frequency) of the Fourier grid used for the measurements, as well as the bin width. These can be specified via the keys `'kfun'` and `'dk'` in the binning dictionary. For example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "74b90171-f6a6-4579-ac59-bcb09d50ee0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "binning = {'kfun':0.005, 'dk':0.005}\n",
    "\n",
    "k = 0.005 + np.arange(80)*0.005\n",
    "Pell_discrete = EFT.Pell(k, params, [0,2,4], 'lambda', binning=binning)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f5a48fc-adfa-4cd1-be0d-78c528c0c0d7",
   "metadata": {},
   "source": [
    "Note: when calling `Pell` with the binning dictionary, the wavemodes specified via the argument `k` are assumed to be the bin centres."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb058230-d492-4d0f-aa6e-e3aaad74e059",
   "metadata": {},
   "source": [
    "Hint: calling `Pell` for the first time with the binning dictionary takes a while longer as COMET has to find the set of discrete modes first. Subsequent calls (provided that the binning options or the maximum bin centre have not been changed) are much faster."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "999e10df-382b-4501-80c6-45648e3e7050",
   "metadata": {},
   "source": [
    "A common approximation to account for the finite bin width is to evaluate the power spectrum multipoles at the so-called effective wave modes, which are weighted averages over the discrete modes in a given bin. If one wants to evaluate the power spectrum multipoles at those effective modes, one can specify the additional key `'effective':True` (`False` by default) in the binning dictionary; the wave modes specified via `k` are still supposed to correspond to the bin centres also in this case."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "51f99e6a-5ebc-4c25-a1cf-81a578c03f3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "Pell_discrete_eff = EFT.Pell(k, params, [0,2,4], 'lambda', \n",
    "                             binning={'kfun':0.005, 'dk':0.005, 'effective':True})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5273b48d-14c1-4ed5-a1f9-c2a0a5e40166",
   "metadata": {},
   "source": [
    "Let's compare the two sets of predictions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "6b9950c2-d06e-4e40-aad6-af2b12da79b0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0, 0.5, '$k\\\\,P_{\\\\ell}(k)$ [$(h/\\\\mathrm{Mpc})^2$]')"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 61908 characters]\n",
      "text/plain": [
       "<Figure size 1000x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "f = plt.figure(figsize=(10,5))\n",
    "ax = f.add_subplot(111)\n",
    "\n",
    "ax.plot(k, k*Pell_discrete['ell0'], 'o', c='C0', mfc='none', ms=3.5, label='discrete')\n",
    "ax.plot(k, k*Pell_discrete['ell2'], 'o', c='C1', mfc='none', ms=3.5)\n",
    "ax.plot(k, k*Pell_discrete['ell4'], 'o', c='C2', mfc='none', ms=3.5)\n",
    "\n",
    "ax.plot(k, k*Pell_discrete_eff['ell0'], c='C0', label='effective')\n",
    "ax.plot(k, k*Pell_discrete_eff['ell2'], c='C1')\n",
    "ax.plot(k, k*Pell_discrete_eff['ell4'], c='C2')\n",
    "\n",
    "ax.legend()\n",
    "ax.set_xlabel('$k$ [$h/\\mathrm{Mpc}$]',fontsize=15)\n",
    "ax.set_ylabel('$k\\,P_{\\ell}(k)$ [$(h/\\mathrm{Mpc})^2$]',fontsize=15)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6931bfc4-f1e8-4fb8-b68f-dd917a4a68d4",
   "metadata": {},
   "source": [
    "## Bispectrum"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2548f16c-7789-4a16-ab00-f0a5b4f68877",
   "metadata": {},
   "source": [
    "COMET also provides the possibility to correct for binning and discreteness effects in the bispectrum, using the approximation introduced in Eggemeier et al. 2025. Like for the power spectrum, the user can call the Bell function with a binning dictionary. However, there are a number of additional options available, which are summarised below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "e0b87571-d787-4e0f-b819-630a7fdf25b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "binning = {\n",
    "   'kfun':0.0025,                 # fundamental frequency of Fourier grid\n",
    "   'dk':0.005,                    # bin width\n",
    "   'first_bin_centre':0.0025,     # k-mode of first bin centre\n",
    "   'do_rounding':False,           # apply rounding to fundamental configurations: True(default)/False\n",
    "   'decimals':[3,3],              # defines rounding precision, default: [3,3]\n",
    "   'shape_limits':[0.999,2.001],  # defines for which triangle configurations the binning/discreteness corrections are computed, default: [0.999,1.15]\n",
    "   'fiducial_cosmology':{         # defines for which fiducial cosmology the corrections are computed, default: Planck2018 + redshift in parameter dictionary\n",
    "      'h': 0.7, 'wc': 0.12,\n",
    "      'wb': 0.022, 'ns': 0.96,\n",
    "      'As': 2.2, 'w0': -1.0,\n",
    "      'wa': 0.0, 'z': 0.5\n",
    "   },\n",
    "   'filename_root_kernels':'test' # filename root to store binned tables\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48a9ab73-eec4-4f1b-9973-d5ea6f427984",
   "metadata": {},
   "source": [
    "With the settings above, it is possible to define the triangle configurations for which the binning and discreteness corrections are being computed, as well as the efficiency (at the expense of accuracy). The `'shape_limits'` property allows the user to specify a tuple of numbers `[a,b]`, which select the following triangle configurations:\n",
    "$$\n",
    "\\frac{k_2+k_3}{k_1} < b \\quad \\land \\quad \\frac{k_2+k_3}{k_1} > a\n",
    "$$\n",
    "In the following example with `binning['shape_limits'] = [0.999,1.15]` this corresponds to all triangle configurations between the two orange lines, i.e., triangle configurations that are closer to being equilateral (top right corner) are not considered for the binning correction. If one intends to compute the binning and discreteness corrections for all triangle configurations instead, one should set `binning['shape_limits'] = [0.999,2.001]`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "adb3d74c-f950-466d-8f07-8062f8654c1a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x2a7663dc0>]"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 23448 characters]\n",
      "text/plain": [
       "<Figure size 500x300 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "f = plt.figure(figsize=(5,3))\n",
    "ax = f.add_subplot(111)\n",
    "\n",
    "x1 = np.linspace(0,0.5)\n",
    "x2 = np.linspace(0.5,1)\n",
    "\n",
    "ax.set_xticks(np.linspace(0,1,5))\n",
    "ax.set_xlabel(r'$k_3/k_1$')\n",
    "ax.set_xticklabels(['0.00','0.25','0.50','0.75','1.00'])\n",
    "ax.set_yticks(np.linspace(0.5,1,3))\n",
    "ax.set_ylabel(r'$k_2/k_1$')\n",
    "ax.plot(x1,1.-x1,c='k',lw=1)\n",
    "ax.plot(x2,x2,c='k',lw=1)\n",
    "ax.plot(np.concatenate((x1,x2)),np.ones(100),c='k',lw=1)\n",
    "ax.set_xlim(-0.05,1.05)\n",
    "ax.set_ylim(0.45,1.05)\n",
    "\n",
    "shape_limits = [0.999, 1.15]\n",
    "x3 = np.linspace(shape_limits[1]-1,shape_limits[1]/2)\n",
    "x4 = np.linspace(shape_limits[0]-1,shape_limits[0]/2)\n",
    "ax.plot(x3,shape_limits[1]-x3,c='C1',lw=3)\n",
    "ax.plot(x4,shape_limits[0]-x4,c='C1',lw=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8865a88-acb8-46ef-8c3f-133322ab6606",
   "metadata": {},
   "source": [
    "The properties `'do_rounding'` in combination with `'decimals'` can be used to reduce the number of fundamental triangles over which the theory predictions have to be averaged in order to improve efficiency. For `binning['decimals'] = [d1, d2]` $k_1,\\,k_2,\\,k_3$ and $\\mu_1,\\,\\mu_2,\\,\\mu_3$ values are approximated as follows:\n",
    "$$\n",
    "   k_i \\approx \\left\\lfloor 10^{d_1}\\,\\frac{k_i}{\\Delta k} \\right\\rceil \\, 10^{-d_1}\\,\\Delta k \\newline\n",
    "   \\mu_i \\approx \\left\\lfloor 10^{d_2}\\,\\mu_i \\right\\rceil \\, 10^{-d_2}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f80fdf33-8762-41a3-bf0f-17664cd9f3ab",
   "metadata": {},
   "source": [
    "The COMET binning module constructs the list of triangle configurations based on the first bin centre, the binwdith (both given in the binning dictionary), and the maximum k-mode given in the `tri` array when calling `Bell`. Currently, it assumes that the bin centres strictly form a closed triangle, i.e. $k_1 \\leq k_2 + k_3$ for $k_1 \\geq k_2 \\geq k_3$."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22383c84-80b6-4e8e-8dbb-9124532dee46",
   "metadata": {},
   "source": [
    "Depending on the number of triangle configurations, the identification of the fundamental triangles and the averaging of the bispectrum kernel functions can be computationally demanding. However, for a given fundamental frequency, bin width and maximum k-mode, this only has to be performed once, such that the subsequent evaluation of the bispectrum model is very fast. For that reason, COMET allows to store any required information, such that at any later time (e.g., after re-initialising COMET), the computationally demanding steps can be skipped. By specifying the property `filename_root_kernels` one can set the root for the files that are generated, and when calling `Bell` again with the same binning dictionary, COMET will try to look for any existing files."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73578c65-7c8e-4cdf-864f-60c94356d64e",
   "metadata": {},
   "source": [
    "This only works if *all* properties of the binning dictionary are **identical**."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97b17591-843c-4f10-a67a-72729ecbdd78",
   "metadata": {},
   "source": [
    "Let's compare the bispectrum with and without the binning and discreteness corrections:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "1de442b5-9780-4ca5-b326-3e8651ed8841",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# define triangle configurations\n",
    "k_hMpc_lin = np.arange(binning['first_bin_centre'], 0.05, binning['dk'])\n",
    "tri =[]\n",
    "for i1,k1 in enumerate(k_hMpc_lin):\n",
    "    for i2,k2 in enumerate(k_hMpc_lin[:i1+1]):\n",
    "        for i3,k3 in enumerate(k_hMpc_lin[:i2+1]):\n",
    "            if i2 + i3 >= i1 - binning['first_bin_centre']/binning['dk']:\n",
    "                tri.append([k1, k2, k3])\n",
    "tri=np.asarray(tri)\n",
    "\n",
    "# let's evalaute with the parameters used in the fiducial cosmology\n",
    "for p in binning['fiducial_cosmology']:\n",
    "    params[p] = binning['fiducial_cosmology'][p]\n",
    "\n",
    "# evaluate bispectrum at the bin centres\n",
    "Bell = EFT.Bell(tri, params, [0,2], 'lambda', kfun=binning['kfun'])\n",
    "\n",
    "# evaluate bispectrum at the bin centres including the binning and discreteness corrections (this may take a few minutes)\n",
    "Bell_discrete = EFT.Bell(tri, params, [0,2], 'lambda', kfun=binning['kfun'], binning=binning)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "a5d9d7f5-0077-4724-beac-032a26c02ae8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 0, 'Triangle Index')"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 123852 characters]\n",
      "text/plain": [
       "<Figure size 1000x500 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, axs = plt.subplots(2,1, figsize=(10,5), sharex=True,)\n",
    "for i in range(2):\n",
    "    axs[i].semilogy(np.arange(tri.shape[0]), Bell[\"ell\"+str(2*i)],'o',c='C'+str(2*i),ls='-',ms=3,label='bin centres')\n",
    "    axs[i].semilogy(np.arange(tri.shape[0]), Bell_discrete[\"ell\"+str(2*i)],'D',c=['C9','C8'][i],ls='--',ms=2,label='discrete')\n",
    "    axs[i].set_ylabel(f'$B_{i*2}(k)$',fontsize=15)\n",
    "    axs[i].legend(loc=1)\n",
    "axs[1].set_xlabel('Triangle Index',fontsize=15)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "bf325d95-32df-458d-9908-93127167e911",
   "metadata": {},
   "outputs": [],
   "source": [
    "VDG = comet(model='VDG_infty', use_Mpc=False)\n",
    "VDG.define_fiducial_cosmology(params_fid=params_fid)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "d4b01ae7-3ffc-4e6b-9791-643d316691c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.optimize import curve_fit\n",
    "from scipy.interpolate import RegularGridInterpolator\n",
    "\n",
    "# extend the range of triangle configurations to see an effect of the damping\n",
    "k_hMpc_lin = np.arange(binning['first_bin_centre'], 0.14, binning['dk'])\n",
    "tri =[]\n",
    "for i1,k1 in enumerate(k_hMpc_lin):\n",
    "    for i2,k2 in enumerate(k_hMpc_lin[:i1+1]):\n",
    "        for i3,k3 in enumerate(k_hMpc_lin[:i2+1]):\n",
    "            if i2 + i3 >= i1 - binning['first_bin_centre']/binning['dk']:\n",
    "                tri.append([k1, k2, k3])\n",
    "tri=np.asarray(tri)\n",
    "\n",
    "# generate some realistic bispectrum covariance matrix\n",
    "Bell_cov = EFT.Bell_covariance(tri, params, [0,2], dk=binning['dk'], de_model='lambda', \n",
    "                               kfun=binning['kfun'], volume=3e9)\n",
    "\n",
    "def compute_sv_avir_mapping(EFT, VDG, tri, params_fid, kf, cov_matrix, \n",
    "                            navirB, nsv, sv_min=2, sv_max=10):\n",
    "    \"\"\"\n",
    "       This function fits the bispectrum multipoles (monopole and quadrupole) from an \n",
    "       expansion of the damping function to predictions that originate from the exact \n",
    "       damping function for a range of 'avirB' and 'sv' values.\n",
    "       \n",
    "       Parameters\n",
    "       ----------\n",
    "       EFT: PTEmu object\n",
    "          Comet instance of the EFT model (with default bispectrum counterterm prescription)\n",
    "       VDG: PTEmu object\n",
    "          Comet instance of the VDG_infty model\n",
    "       tri: numpy.array\n",
    "          Array of triangle configurations\n",
    "       params_fid: dictionary\n",
    "          Fiducial cosmological parameters (and linear bias) to use for the calibration\n",
    "       kf: float\n",
    "          Fundamental frequency\n",
    "       cov_matrix: numpy.array\n",
    "          Covariance matrix for the bispectrum multipoles\n",
    "       navirB: integer\n",
    "          Number of bins in 'avirB'\n",
    "       nsv: integer\n",
    "          Number of bins in 'sv'\n",
    "       sv_min: float\n",
    "          Minimum 'sv' value\n",
    "       sv_max: float\n",
    "          Maximum 'sv' value\n",
    "          \n",
    "       Returns\n",
    "       -------\n",
    "       avirB_list: numpy.array\n",
    "          List of covered 'avirB' values\n",
    "       sv_list: numpy.array\n",
    "          List of covered 'sv' values\n",
    "       mapping: numpy.array\n",
    "          Corresponding coefficients for the mapping to 'cnloB'\n",
    "    \"\"\"\n",
    "    def Bapprox(tri, a):\n",
    "        params['cnloB'] = -a*VDG.params['avirB']**1.75 - 0.5*VDG.params['sv']**1.75\n",
    "        B = EFT.Bell(tri, params, ell=[0,2], de_model='lambda', kfun=kf)\n",
    "        return np.hstack([B[m] for m in B.keys()])\n",
    "    \n",
    "    params = {}\n",
    "    for p in ['wc','wb','ns','h','As','z']:\n",
    "        params[p] = params_fid[p]\n",
    "    params['b1'] = params_fid['b1']\n",
    "    \n",
    "    avirB_list = np.logspace(-2,np.log10(10),navirB)\n",
    "    sv_list = np.linspace(sv_min,sv_max,nsv)\n",
    "    mapping = np.zeros((navirB,nsv))\n",
    "    for i,avirB in enumerate(avirB_list):\n",
    "        for j,sv in enumerate(sv_list):\n",
    "            params['avirB'] = avirB\n",
    "            VDG.params['sv'] = sv\n",
    "            Bref = VDG.Bell(tri, params, [0,2], 'lambda', kfun=kf)\n",
    "            Bref = np.hstack([Bref[m] for m in Bref])\n",
    "            popt, pcov = curve_fit(Bapprox, tri, Bref, sigma=cov_matrix)\n",
    "            mapping[i,j] = popt\n",
    "    return avirB_list, sv_list, mapping\n",
    "\n",
    "# this may take a few minutes\n",
    "avirB_list, sv_list, mapping = compute_sv_avir_mapping(EFT, VDG, tri, params, binning['kfun'], Bell_cov, 10, 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "6ea76117-b39d-4351-b85f-cc200db081de",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 0, 'sv')"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 10664 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.imshow((np.log(np.abs(mapping))))\n",
    "plt.ylabel('avirB',fontsize=15)\n",
    "plt.xlabel('sv',fontsize=15)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "42b52ab8-ffe6-4d5b-8e55-f54ea062b4db",
   "metadata": {},
   "outputs": [],
   "source": [
    "cnloB_spline = RegularGridInterpolator((avirB_list,sv_list), mapping)\n",
    "\n",
    "# Going back to the smaller triangle configuration grid\n",
    "k_hMpc_lin = np.arange(binning['first_bin_centre'], 0.05, binning['dk'])\n",
    "tri =[]\n",
    "for i1,k1 in enumerate(k_hMpc_lin):\n",
    "    for i2,k2 in enumerate(k_hMpc_lin[:i1+1]):\n",
    "        for i3,k3 in enumerate(k_hMpc_lin[:i2+1]):\n",
    "            if i2 + i3 >= i1 - binning['first_bin_centre']/binning['dk']:\n",
    "                tri.append([k1, k2, k3])\n",
    "tri=np.asarray(tri)\n",
    "\n",
    "for p in binning['fiducial_cosmology']:\n",
    "    params[p] = binning['fiducial_cosmology'][p]\n",
    "params['avirB'] = 4\n",
    "\n",
    "VDG.params['wc'] = 0.1 # to trigger re-evaluation of the emulators in the call below (so that the 'sv' value is updated)\n",
    "Bell_VDG = VDG.Bell(tri, params, [0,2], 'lambda', kfun=binning['kfun'])\n",
    "\n",
    "binning['filename_root_kernels'] = 'test_VDG' # need to use a different filename root\n",
    "Bell_VDG_discrete = VDG.Bell(tri, params, [0,2], 'lambda', kfun=binning['kfun'], \n",
    "                             binning=binning, cnloB_mapping=cnloB_spline)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "693a42b6-35b7-425b-89c4-2903465e37a5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 0, 'Triangle Index')"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 129056 characters]\n",
      "text/plain": [
       "<Figure size 1000x500 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, axs = plt.subplots(2,1, figsize=(10,5), sharex=True,)\n",
    "for i in range(2):\n",
    "    axs[i].semilogy(np.arange(tri.shape[0]), Bell_VDG[\"ell\"+str(2*i)],'o',c='C'+str(2*i),ls='-',ms=3,label='bin centres')\n",
    "    axs[i].semilogy(np.arange(tri.shape[0]), Bell_VDG_discrete[\"ell\"+str(2*i)],'D',c=['C9','C8'][i],ls='--',ms=2,label='discrete')\n",
    "    axs[i].set_ylabel(f'$B_{i*2}(k)$',fontsize=15)\n",
    "    axs[i].legend(loc=1)\n",
    "axs[1].set_xlabel('Triangle Index',fontsize=15)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "id": "47ad9914-cb00-4b84-bba8-75332332e1fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "mu_true = 0.5\n",
    "mu_biased = 3.5\n",
    "sig_true = 1.0\n",
    "diff = mu_biased-mu_true\n",
    "sig_biased = np.sqrt(sig_true**2+diff**2)\n",
    "\n",
    "samples_true = np.random.normal(loc=mu_true, scale=sig_true, size=50000)\n",
    "samples_biased = np.random.normal(loc=mu_biased, scale=sig_biased, size=50000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "id": "01efe11f-5d57-493e-a94a-a11809f67c94",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.lines.Line2D at 0x2a9b41540>"
      ]
     },
     "execution_count": 80,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 20312 characters]\n",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.hist(samples_true,30,histtype='step')\n",
    "plt.hist(samples_biased,30,histtype='step')\n",
    "plt.axvline(mu_true,color='C0')\n",
    "plt.axvline(mu_biased,color='C1')\n",
    "plt.axvline(mu_biased-sig_biased,color='C1',ls='--')\n",
    "plt.axvline(mu_biased+sig_biased,color='C1',ls='--')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "5af95f1e-95cb-4253-b3c1-43cf82fbc33f",
   "metadata": {},
   "outputs": [],
   "source": [
    "diff = samples_biased-samples_true\n",
    "\n",
    "diff_avg = np.average(diff)\n",
    "diff_std = np.std(diff)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "694676c0-2ad6-4559-843c-6e64f2c8e484",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.20243925832785903"
      ]
     },
     "execution_count": 60,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "diff_avg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "id": "5f6b4c50-a4b3-41d9-8dc0-295f231706a7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.5102040816326532"
      ]
     },
     "execution_count": 64,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "1/(1+0.2/0.5)**2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "a92efa59-c1a1-4f66-8d37-6d51e138c107",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.9720079321704735"
      ]
     },
     "execution_count": 61,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "diff_std**2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afe37546",
   "metadata": {},
   "source": [
    "# Beyond $P_{\\ell}$ predictions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d7f49e5",
   "metadata": {},
   "source": [
    "In the following we demonstrate a number of additional outputs that COMET can provide. Specifically:\n",
    "  * The linear power spectrum, with and without infra-red resummation\n",
    "  * The Gaussian covariance matrix for the power spectrum multipoles\n",
    "  * The tree-level bispectrum multipoles"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8874ecbd-12ff-4f01-a4f3-c368b000eb59",
   "metadata": {},
   "source": [
    "## Linear power spectrum"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07cbd433-d7d0-47ae-8173-569652e48710",
   "metadata": {},
   "source": [
    "The linear power spectrum (no infra-red resummation; simply the emulated CAMB output) can be obtained from the function `PL`, while the linear power spectrum with damped BAO wiggles (infra-red resummation) can be obtained from the function `Pdw` (note: this is not the smooth, no-wiggle power spectrum). The arguments are identical to those of `Pell` with the exception that we no longer need to specify a multipole number. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "81fe547e-27fc-4ed6-b30d-7b2f5c367714",
   "metadata": {},
   "outputs": [],
   "source": [
    "k = np.logspace(-3,np.log10(0.4),300)\n",
    "Pdw = EFT.Pdw(params=params, k=k, de_model='lambda')\n",
    "PL = EFT.PL(params=params, k=k, de_model='lambda')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82c951fb-8839-4d63-b0ca-f9c0c1c13bcb",
   "metadata": {},
   "source": [
    "Let's plot the ratio of the de-wiggled linear power spectrum over the linear power spectrum:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "25f7e68d-2a8a-4da7-81a6-f8adf44ce31a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 49264 characters]\n",
      "text/plain": [
       "<Figure size 1000x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "f = plt.figure(figsize=(10,5))\n",
    "ax = f.add_subplot(111)\n",
    "ax.semilogx(k, Pdw/PL,c='C0',ls='-')\n",
    "ax.set_xlabel('$k$ [h/Mpc]',fontsize=15)\n",
    "ax.set_ylabel(r'$P_{\\rm dw}(k)/P_{L}(k)$',fontsize=15)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7009e077",
   "metadata": {},
   "source": [
    "## Tree-level bispectrum"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c894659-771e-40e2-b3c0-6dfd9c981ba8",
   "metadata": {},
   "source": [
    "COMET can also output the tree-level bispectrum (in real-space, for the `RS` model) and its multipoles (in redshift-space, for the `EFT` model). These predictions are not emulated, but computed from the emulated de-wiggled power spectrum directly. For that purpose we provide the function `Bell` and in order to demonstrate its usage let's first generate a set of triangle configurations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "4e32a9d5-cefc-478c-ba43-4d3bac0bc338",
   "metadata": {},
   "outputs": [],
   "source": [
    "k_hMpc_lin = np.arange(0.005, 0.3, 0.005)\n",
    "tri =[]\n",
    "for i1,k1 in enumerate(k_hMpc_lin):\n",
    "    for i2,k2 in enumerate(k_hMpc_lin[:i1+1]):\n",
    "        for i3,k3 in enumerate(k_hMpc_lin[:i2+1]):\n",
    "            if k2 + k3 >= k1:\n",
    "                tri.append([k1, k2, k3])\n",
    "tri=np.asarray(tri)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5cde96c-f2e2-4970-a82f-4b70347e158c",
   "metadata": {},
   "source": [
    "The `Bell` function has the same arguments and functionality as the analogous `Pell` function for the power spectrum. However, it expects the triangle configurations to be always specified as a numpy array containing $k_1$, $k_2$, $k_3$ (it is not possible to evaluate the multipoles for different triangles at the moment), and in addition it includes the argument `kfun`, which is used for compressing the number of unique k-modes and is ideally chosen as a value that corresponds closely to the spacing between configurations (e.g. the bin-width for measured data), but must not be much larger. If in doubt, use a value much smaller than the typical spacing. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "020c5e79-d41d-47e4-9a3a-5bcbdcfed8d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "params['h'] = 0.69\n",
    "params['z'] = 0.57\n",
    "Bell = EFT.Bell(tri, params=params, ell=[0,2,4], de_model='lambda', kfun=0.005)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ee33d90-3ff3-49b0-a948-8039a4e651f6",
   "metadata": {},
   "source": [
    "_Note: The very first call of `Bell` for a given set of configurations can take a little longer (depending on the total number of triangle configurations) as some lookup-tables are generated. All subsequent calls, even with changing cosmological parameters, are then much faster. That implicitly means that one should avoid calling `Bell` multiple times with different triangle configurations, but once for all triangle configurations._"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "fe912a1f-d328-4ddf-97a4-5f20351e67b2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 441608 characters]\n",
      "text/plain": [
       "<Figure size 1000x500 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, axs = plt.subplots(3,1, figsize=(10,5), sharex=True,)\n",
    "for i in range(3):\n",
    "    axs[i].semilogy(np.arange(tri.shape[0]), Bell[\"ell\"+str(2*i)],'.',c='C'+str(2*i),ls='-',ms=1,lw=0.2)\n",
    "    axs[i].set_ylabel(f'$B_{i*2}(k)$',fontsize=15)\n",
    "\n",
    "fig.tight_layout()\n",
    "plt.subplots_adjust(wspace=0, hspace=0)\n",
    "axs[-1].set_xlabel('Triangle index - $k$ [h/Mpc]',fontsize=15)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86cfb4eb",
   "metadata": {},
   "source": [
    "## Computing covariance matrices"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c05ce32",
   "metadata": {},
   "source": [
    "Apart from the multipoles we can also generate (Gaussian) covariance matrices, for which there are again different flags, that are either defined for the $\\sigma_{12}$ or the dark energy parameter spaces. For the power spectrum, the first three arguments, `k`, `params`, and `ell`, are identical to those for `Pell`. In addition, we need to specify a binwidth `dk` and volume (both of which need to be given in the respective units for which the emulator is configured in), for example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "929e48ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "dk_hMpc = 0.005\n",
    "k_hMpc_lin = np.arange(0.001, 0.3, dk_hMpc)\n",
    "vol_hMpc = 3e9"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "c99753d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "Cov_hMpc = EFT.Pell_covariance(k_hMpc_lin, params, ell=[0,2,4], dk=dk_hMpc, volume=vol_hMpc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "a66eeda4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 66536 characters]\n",
      "text/plain": [
       "<Figure size 900x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.figure(figsize=(9,6))\n",
    "plt.title(r\"\")\n",
    "plt.title(r\"Correlation Matrix\")\n",
    "var_inv = np.diag(1./np.sqrt(np.diag(Cov_hMpc)))\n",
    "R_hMpc = var_inv @ Cov_hMpc @ var_inv\n",
    "plt.imshow(R_hMpc,cmap='magma_r')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c1f12ae-1aaf-4724-a4fc-45cc11dd13c9",
   "metadata": {},
   "source": [
    "The argument specifying the scales provides the same functionality as for `Pell`, that is, it can either be given as a number or numpy array, in which case all specified multipoles are evaluated for the same scales, or a list of numbers/numpy arrays, in which case the first entry is evaluated for the first multipole in `ell` etc."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c77e2f7",
   "metadata": {},
   "source": [
    "For the version with specified dark energy model it is also possible (in addition to providing the volume via the `volume` argument) to provide minimum and maximum redshifts, `zmin` and `zmax`, a sky fraction `fsky`, and a volume scaling factor `volfac` (by default equal to 1), such that the volume is computed in accordance with the given cosmological model. For example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "7e616fea",
   "metadata": {},
   "outputs": [],
   "source": [
    "Cov_hMpc_LCDM = EFT.Pell_covariance(\n",
    "                        k_hMpc, \n",
    "                        params, \n",
    "                        ell=[0,2,4],\n",
    "                        dk=2*np.pi/3780, \n",
    "                        zmin=params['z']-0.1, \n",
    "                        zmax=params['z']+0.1, \n",
    "                        fsky=15000./(360**2/np.pi), \n",
    "                        volfac=1, \n",
    "                        de_model=\"lambda\", \n",
    "                )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "307d96f4",
   "metadata": {},
   "source": [
    "As a further extension, in the case when using measurements from a periodic box that have been averaged over different lines of sight, we have added the averaging corrections for the covariance matrix. We have created the flags `avg_cov` (set to `False` by default) and `avg_los` (set to 3 by default) for the `Pell_covariance` function, so that when `avg_cov=True` it by default will compute the average along the three perpendicular axes (x,y,z), but it is also possible to average over just 2 directions. Note that this computation is quite slow since it involves a different  integral for each k-bin, it may be optimised in the future."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a9b7211-ad1a-4723-9c1a-efb96a2e419c",
   "metadata": {},
   "source": [
    "Similarly, we can compute the Gaussian covariance matrix of the bispectrum using the function `Bell_covariance`. Apart from the first argument, which specifies the triangle configurations (or a list of configurations for different multipoles), the arguments are identical to those of `Pell_covariance`. In addition, one can also specify `kfun` as in case of `Bell` (see above), which by default is set to the bin width `dk`. Let us compute the bispectrum covariance matrix for a reduced set of triangle configurations with different scale cuts for the monopole, quadrupole, and hexadecapole:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "a28661c1-8b07-44f4-8b6d-194b7193615b",
   "metadata": {},
   "outputs": [],
   "source": [
    "id0p1 = np.where(tri[:,0] < 0.1)\n",
    "id0p06 = np.where(tri[:,0] < 0.06)\n",
    "id0p03 = np.where(tri[:,0] < 0.03)\n",
    "\n",
    "# using the same scale cut for all multipoles\n",
    "Cov_Bisp_hMpc = EFT.Bell_covariance(tri[id0p1], params, ell=[0,2,4], dk=0.005, de_model='lambda',\n",
    "                                    kfun=0.005, volume=3e9)\n",
    "\n",
    "# using different scale cuts\n",
    "Cov_Bisp_hMpc_diff_scale_cut = EFT.Bell_covariance([tri[id0p1],tri[id0p06],tri[id0p03]], params, ell=[0,2,4], \n",
    "                                                   dk=0.005, de_model='lambda', kfun=0.005, volume=3e9)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3df5b88f-04c9-40f8-bc4c-b6355bbba307",
   "metadata": {},
   "source": [
    "In the Gaussian approximation each block in the bispectrum covariance matrix is diagonal. Let's plot these diagonals as a function of the triangle configuration index:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "a1b774ee-db1f-4e09-a8b2-a3b4519c8945",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 128572 characters]\n",
      "text/plain": [
       "<Figure size 1000x500 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, axs = plt.subplots(2,3, figsize=(10,5), sharex=True, sharey=True)\n",
    "\n",
    "ntri = id0p1[0].shape[0]\n",
    "\n",
    "labels = ['$C_{00}$', '$C_{22}$', '$C_{44}$', '$C_{02}$', '$C_{04}$', '$C_{24}$']\n",
    "colors = ['C0','C1','C2','C3','C4','C5']\n",
    "for i in range(3):\n",
    "    axs[0,i].semilogy(np.arange(ntri), np.diag(Cov_Bisp_hMpc[i*ntri:(i+1)*ntri,i*ntri:(i+1)*ntri]), c=colors[i], label=labels[i])\n",
    "    axs[0,i].legend(fontsize=15)\n",
    "    \n",
    "n = 0\n",
    "for i in range(2):\n",
    "    for j in range(i,3):\n",
    "        if i != j:\n",
    "            axs[1,n].semilogy(np.arange(ntri), np.diag(Cov_Bisp_hMpc[i*ntri:(i+1)*ntri,j*ntri:(j+1)*ntri]), c=colors[n+3], label=labels[n+3])\n",
    "            axs[1,n].legend(fontsize=15)\n",
    "            axs[1,n].set_xlabel('Triangle Index',fontsize=15)\n",
    "            n += 1\n",
    "\n",
    "fig.tight_layout()\n",
    "plt.subplots_adjust(wspace=0, hspace=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ce360a0-4a1e-474c-a762-16b40fecd55d",
   "metadata": {},
   "source": [
    "Note that both, `Pell_covariance` and `Bell_covariance`, allow also to specify the number of fundamental modes and fundamental triangle per bin, respectively. This is possible by using the optional arguments `Nmodes` and `Ntri`, which should be an array of the same length as either `k` or `tri` (and if either of these is given as a list, it should match the length of the longest entry in the list of scales or triangle configurations). If not provided, the following approximations are assumed when computing the covariance matrix:\n",
    "$$\n",
    "N_{\\rm modes} \\approx \\frac{V}{6 \\pi^2}\\,\\left[\\left(k+\\frac{\\Delta k}{2}\\right)^3 - \\left(k-\\frac{\\Delta k}{2}\\right)^3\\right]\\,, \\\\[1.5em]\n",
    "N_{\\rm tri} \\approx \\frac{V^2}{8 \\pi^4}\\,k_1\\,k_2\\,k_3\\,\\Delta k^3\\,.\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29ca9aaf-4ec8-42b9-84bb-b14bc03c0f76",
   "metadata": {},
   "source": [
    "# Working with data sets"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe6a14c2-3840-40ef-b58a-58c184e744ba",
   "metadata": {},
   "source": [
    "## Loading data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31ef68c0-fe8f-4f2b-a844-9c79021882ba",
   "metadata": {},
   "source": [
    "We can load measurements of the power spectrum and bispectrum multipoles into COMET using the `define_data_set` function. This function takes first an identifier for the data set (`obs_id`; this can be anything, it will be used to reference the data) and any one of the following arguments:\n",
    "   * `stat`. Can either be `'powerspectrum'` or `'bispectrum'`; if not provided, `stat` is deduced from the number of columns in `bins` (see below).\n",
    "   * `bins`. In case of the power spectrum: 1d-array of k-modes corresponding to the measurements; in case of the bispectrum: 2d-array with three columns corresponding to the triangle configuration ($k_1$, $k_2$, $k_3$) of the measurements.\n",
    "   * `signal`. The measurements of the power spectrum or bispectrum; the size of the first dimension must match the size of `bins`, and it is assumed that the first column corresponds to the monopole, the second to the quadrupole, and the third to the hexadecapole (one does not need to provide all three multipoles, i.e., one can provide only the monopole, or monopole + quadrupole, but one cannot leave out preceding multipoles).\n",
    "   * `cov`. The covariance matrix of the measurements, which must match the combined size of all given multipoles. If the dimension of `cov` is one-dimensional, it is assumed to be the diagonal of the covariance matrix.\n",
    "   * `theory_cov`. A flag that specifies whether the given covariance matrix was derived analytically or from a set of simulation measurements. In the latter case an Anderson-Hartlap correction is applied to the inverse, based on `n_realizations`.\n",
    "   * `n_realizations`. Number of realizations from which the covariance matrix was estimated, only used (and required) in case `theory_cov=False`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c580d9a-b0e6-45e6-b88b-2c094cd0b890",
   "metadata": {},
   "source": [
    "Let us load some mock power spectrum measurements:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "dfb73c6b-2d6b-4160-9037-000dc6f623ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = np.loadtxt('mock_Pk_mean.dat')\n",
    "Cov = np.loadtxt('mock_Pk_cov.dat')\n",
    "\n",
    "k = data[:,0]\n",
    "P0 = data[:,1]\n",
    "P2 = data[:,3]\n",
    "P4 = data[:,5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "a75859e4-eef0-4b37-879e-9424a4e94b73",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's call this data set 'mock_Pk'\n",
    "EFT.define_data_set(obs_id='mock_Pk', bins=k, signal=np.array([P0,P2,P4]).T, cov=Cov, theory_cov=False, n_realizations=300)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e26d4479-7cf7-407e-98d0-fd2740bf41fd",
   "metadata": {},
   "source": [
    "We can access the data through `EFT.data['mock_Pk']` and check, for example, that the type of statistic was correctly identified (since it was provided above):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "ab0fcad3-a781-4381-bfcf-4014c4608b5c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'powerspectrum'"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "EFT.data['mock_Pk'].stat"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5921de13",
   "metadata": {},
   "source": [
    "## Computing the $\\chi^2$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "560a7c74",
   "metadata": {},
   "source": [
    "Finally, we can let COMET directly compute $\\chi^2$ values based on the provided data set, a given set of model parameters and range of scales. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80c8020d",
   "metadata": {},
   "source": [
    "To do so, we call the function `chi2`, which takes as arguments the identifier of the data set, the parameter dictionary, a maximum k-mode value `kmax`, a model argument `de_model`. `kmax` can either be a number, in which case the same cutoff is applied for all multipoles, or a list of numbers for each individual multipole, as for the multipoles case. If the cutoff is zero (or smaller than the minimum scale of the observations) for a particular multipole, then it is excluded from the computation of the chi-square. `kmax` is also assumed to be in the units of the emulator. `de_model` can be one of the options specified before. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "756446eb-23ed-49f6-9d13-e2387be7e63e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "6754.176546673202"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "EFT.chi2(obs_id='mock_Pk',params=params, kmax=[0.30, 0.30, 0.30], de_model='lambda', convolve_window=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12f776a6",
   "metadata": {},
   "source": [
    "Moreover, in order to speed up the computation of the $\\chi^2$, in the same way as `Pell_fixed_cosmo_boost` function, we can specify the flag `chi2_decomposition` in order to avoid recomputing the quantities depending on cosmological parameters. Let's see how it works"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 117,
   "id": "ee80c0e6-20e2-4c40-9ad2-8fbc9d264b31",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "6.37 ms ± 153 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit EFT.chi2(obs_id='mock_Pk',params=params, kmax=[0.30, 0.30, 0.30], de_model='lambda', chi2_decomposition=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 118,
   "id": "999268cb-f0a7-43e8-8646-e0e168ba6dac",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "9.11 µs ± 20.6 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit EFT.chi2(obs_id='mock_Pk',params=params, kmax=[0.30, 0.30, 0.30], de_model='lambda', chi2_decomposition=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f2109c3-0e63-464f-b8bc-28e39cb69a8d",
   "metadata": {},
   "source": [
    "It is also possible to compute the $\\chi^2$ for multiple data sets by giving `chi2` a list of data identifiers. While in principle this could be useful to simultaneously analyse multiple power spectrum measurements at different redshifts, COMET currently does not support multiple parameter sets with different bias parameters, or at various redshifts (this will be possible in a future release). However, we can use this functionality to compute the joint $\\chi^2$ of the power spectrum and bispectrum. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a60b07c-3543-40af-a3ba-062c01efb93f",
   "metadata": {},
   "source": [
    "As an example, let's load some mock bispectrum data and store it in a new data container:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "7dc46b23-ea1d-43e2-ab24-9815a59ca37f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# data format: k1, k2, k3, B0, B0_var, B2, B2_var, B4, B4_var\n",
    "data = np.loadtxt('mock_Bk_mean.dat') \n",
    "\n",
    "EFT.define_data_set(obs_id='mock_Bk', bins=data[:,:3], signal=data[:,[3,5,7]], cov=np.hstack(data[:,[4,6,8]]), kfun=0.00166)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b609274-8267-4c8c-a90c-1bedfe1ea431",
   "metadata": {},
   "source": [
    "When providing a list of data identifiers, the `kmax` argument passed to `chi2` can be a dictionary of $k_{\\rm max}$ values, where the keys must match the data identifiers. If not given as a dictionary, the same $k_{\\rm max}$ is used for each of the data sets. The following call of `chi2` evaluates the $\\chi^2$ for the power spectrum and bispectrum data sets, using the power spectrum monopole and quadrupole up to $k_{\\rm max} = 0.3$ and $0.25\\,h\\,\\mathrm{Mpc}^{-1}$, respectively, and the bispectrum monopole and hexadecapole up to $k_{\\rm max} = 0.12$ and $0.05\\,h\\mathrm{Mpc}^{-1}$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "ee8f15f2-74d8-4072-90a4-cf821a95fb5c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "65495175908.83485"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "EFT.chi2(['mock_Pk','mock_Bk'], params, {'mock_Pk':[0.3,0.25,0.], 'mock_Bk':[0.12,0.0,0.05]}, de_model='lambda')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbe5c93e-a7b0-474b-9290-5a253a9d96b6",
   "metadata": {},
   "source": [
    "Note that the option `chi2_decomposition` is currently not available for the bispectrum."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5629f78f-a1e9-446f-b5ee-86162bdc20fd",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Convolution with survey window function"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5549d946-86b5-4f8a-85fa-55d289670b51",
   "metadata": {},
   "source": [
    "In order to compare the power spectrum model predictions to some actual measurements, we need to convolve with the survey window function. This can be done within COMET by providing a window function mixing matrix $W_{\\ell\\ell'}(k,k')$ that connects the convolved and unconvolved power spectra via a simple matrix multiplication (see e.g. d'Amico et al. 2019):\n",
    "$$\n",
    "P_{W,\\ell}(k) = W_{\\ell\\ell'}(k,k') \\cdot P_{\\ell'}(k')\\,,\n",
    "$$\n",
    "where the summation over multipole numbers is implicit. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "703d1bca-0e93-4252-a07e-c0e3f9c83926",
   "metadata": {},
   "source": [
    "The mixing matrix and the associated scales for which it has been computed, $k$ and $k'$, can be specified via `define_data_set` using the arguments `bins_mixing_matrix` and `W_mixing_matrix`. The former is a list, containing the arrays for $k$ and $k'$. For example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "30b1a19c-4030-431c-b07a-75f6985f73b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's load some sample window function and k_prime values\n",
    "W = np.fromfile('mock_Pk_window_W.npy').reshape((216, 4854))\n",
    "k_prime = np.loadtxt('mock_Pk_window_kp.dat')\n",
    "\n",
    "# The mixing matrix was computed for the following k-scales\n",
    "k = np.arange(1,73)*2*np.pi/1500\n",
    "\n",
    "# Load everything into COMET using the same data identifier as before ('mock_Pk')\n",
    "EFT.define_data_set(obs_id='mock_Pk', bins_mixing_matrix=[k, k_prime], W_mixing_matrix=W)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "598df954-e9b5-4686-8230-07b9914217ab",
   "metadata": {},
   "source": [
    "We can now obtain the window-convolved power spectrum by passing the additional argument `obs_id` to `Pell` (the same functionality applies also to `Pell_fixed_cosmo_boost`) using the corresponding data identifier:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "fd226c0a-a698-4f48-b83d-e76c6dc725d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "P_unconv = EFT.Pell(k, params, ell=[0,2,4], de_model='lambda')                  # unconvolved, equivalent with obs_id=None\n",
    "P_conv = EFT.Pell(k, params, ell=[0,2,4], de_model='lambda', obs_id='mock_Pk')  # convolved with window function for data set 'mock_Pk'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "2e55fa98-d4e8-4658-86e3-5034d01b9139",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x16ada6160>"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 44536 characters]\n",
      "text/plain": [
       "<Figure size 720x360 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "f = plt.figure(figsize=(10,5))\n",
    "ax = f.add_subplot(111)\n",
    "ax.plot(k, k*P_unconv['ell0'],c='C0',ls='-',label='$P_{0}$')\n",
    "ax.plot(k, k*P_conv['ell0'],c='C0',ls='--',label='$P_{W,0}$')\n",
    "ax.plot(k, k*P_unconv['ell2'],c='C1',ls='-',label='$P_{2}$')\n",
    "ax.plot(k, k*P_conv['ell2'],c='C1',ls='--',label='$P_{W,2}$')\n",
    "ax.plot(k, k*P_unconv['ell4'],c='C2',ls='-',label='$P_{4}$')\n",
    "ax.plot(k, k*P_conv['ell4'],c='C2',ls='--',label='$P_{W,4}$')\n",
    "ax.set_xlabel('$k$ [h/Mpc]',fontsize=15)\n",
    "ax.set_ylabel(r'$k\\,P_{\\ell}(k)$ [$(\\mathrm{Mpc}/h)^{2}$]',fontsize=15)\n",
    "ax.legend(fontsize=15,ncol=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b35e75f9-da27-4896-9833-7f722558bf50",
   "metadata": {},
   "source": [
    "We can also take the window function convolution into account when computing the $\\chi^2$. In that case we set the flag `convolve_window=True` (by default it is set to `False`):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "21e58345-1e39-4444-a9b8-60833aabcbae",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "6745.515807012584"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "EFT.chi2(obs_id='mock_Pk',params=params, kmax=[0.30, 0.30, 0.30], de_model='lambda', convolve_window=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb5809c1-5683-45f3-99a1-e542ac189ba4",
   "metadata": {},
   "source": [
    "This also works in combination with the option `chi2_decomposition=True`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd15423b-adff-43a5-9866-f922b150f68a",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
