{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0914f300-aa40-438f-8281-5cc5c79d6e0a",
   "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": "6c2d91cb-a7d2-44bd-a153-d4a28191738b",
   "metadata": {},
   "source": [
    "Let’s first import `comet` as well as other requried libraries:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "12747bb4-8b3e-4fe3-8438-03e1d340ce64",
   "metadata": {},
   "outputs": [],
   "source": [
    "from comet import comet\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb7a313d-38dc-4949-8cbd-057f5d6591ac",
   "metadata": {},
   "source": [
    "At initialisation we only need to specify the perturbation theory model that we want to use (for an overview of the models implemented in COMET, see here) and we can configure COMET either in $Mpc$ units (`use_Mpc = True`, which is the default option) or in $h^{-1}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}Mpc$ units:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "fb7ce865-a10b-4450-9696-519486083437",
   "metadata": {},
   "outputs": [],
   "source": [
    "EFT=comet(model=\"EFT\", use_Mpc=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "841cff96-bbf0-4b8d-b69e-bc2982fda861",
   "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": "0defabbc-9a94-4541-94f4-9efb54e2a36f",
   "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": "91f09602-421d-4fc2-ada7-ec8d31d9b3a2",
   "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": "f2c0af2f-ff5d-4cb4-8b60-83b2756be76c",
   "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": "8e564305-fb83-48ef-831b-e702924968e8",
   "metadata": {},
   "source": [
    "Finally, we define the values of the bias parameters. The complete list of parameters along with a brief explanation and their dioctionary 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": "793934e8-4863-4390-960f-3364db97d6ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "params['b1'] = 2.\n",
    "params['b2'] = -0.5"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b6bfd15-809e-4690-b70c-2ef3a8d47d22",
   "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 hMpc^{−1}$ to $0.3hMpc^{−1}$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "e889beed-3cd5-416c-86cb-ca9e5668cf72",
   "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": "069cb347-be0b-473a-8fb1-80cb6832aee3",
   "metadata": {},
   "source": [
    "The output of `Pell` is given in a dictionary format:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "0ec02672-23c6-4cb8-b514-ca0dafb7104a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dict_keys(['ell0', 'ell2', 'ell4'])\n"
     ]
    }
   ],
   "source": [
    "print(Pell_LCDM.keys())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2fa79809-04b1-44c6-93ae-a851bf2c80b8",
   "metadata": {},
   "source": [
    "So we can access our results and plot them as follow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f4de477d-51e1-453e-a95b-8b909d2ffba2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 35432 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",
    "\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.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": "61b2d9aa-3380-4ebd-ad16-a9cc575be1e2",
   "metadata": {},
   "source": [
    "# In-depth options for obtaining multipoles"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82f83abe-822f-42e0-9631-7b055c909128",
   "metadata": {},
   "source": [
    "* Now, let's see somo details:\n",
    "    * different cosmologies ($\\omega_0 + \\omega_0\\omega_a$)\n",
    "    * using the $f-\\sigma_{12}$ parameter space\n",
    "    * the options for providing different $k$-scales, float vs np.array vs list and the corresponding outputs\n",
    "    * describe the `fixed_cosmo_boost` function, i.e., speedup when just changing bias parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a729c090-ec36-4d7d-b13f-dc28852fef15",
   "metadata": {},
   "source": [
    "## An alternative Dark energy model."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "701e6fae-82e4-4e49-b5d5-baae2086fd2b",
   "metadata": {},
   "source": [
    "As mentioned before, we can make use of a different cosmological model. Let's try the `w0wa` model. We need to add such parameters to our params dictionary first."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "6acbb3a6-cb26-4596-b1ba-da0368108ff1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# a non-flat\n",
    "# cosmology is assumed if `params_fid` includes the key `Ok`.\n",
    "# For other dark energy models one can set `de_model` to `w0` or `w0wa`, in\n",
    "# which case one needs to provide the values for w0, wa in `params_fid`.\n",
    "params['w0'] = -1.1\n",
    "params['wa'] = 0.1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17b56228-f2b4-4d80-9aec-d55f8eb2b622",
   "metadata": {},
   "source": [
    "Then let's recompute the model updating such parameters and compare with the $\\Lambda CDM$ prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "1c1690d2-142b-40ff-a502-12f612b637d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "Pell_w0wa = EFT.Pell(k_hMpc, params, ell=[0,2,4], de_model='w0wa')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "b313cebc-e9ef-4bcf-93de-b3901926675a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 43876 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",
    "\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_w0wa[\"ell0\"],c='C0',ls='--')\n",
    "ax.semilogx(k_hMpc, k_hMpc**0.5*Pell_w0wa[\"ell2\"],c='C1',ls='--')\n",
    "ax.semilogx(k_hMpc, k_hMpc**0.5*Pell_w0wa[\"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": "39689f76-9ca6-44c0-a74a-52c80a57f7fd",
   "metadata": {},
   "source": [
    "## The $f-\\sigma_{12}$ parameter space."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed5d4849-5e7b-466a-b3bc-7d0bf2f116da",
   "metadata": {},
   "source": [
    "The function called before ignores the values of `s12`, `alpha_tr`, `alpha_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": 15,
   "id": "030cd9a0-e541-4a03-8005-c37d43459fcc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'wc': 0.11544,\n",
       " 'wb': 0.0222191,\n",
       " 'ns': 0.9632,\n",
       " 's12': array([0.5993593]),\n",
       " 'f': 0.7252890877752092,\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",
       " 'N0': 0.0,\n",
       " 'N20': 0.0,\n",
       " 'N22': 0.0,\n",
       " 'h': 0.8,\n",
       " 'As': 2.3,\n",
       " 'Ok': 0.0,\n",
       " 'w0': -1.1,\n",
       " 'wa': 0.1,\n",
       " 'z': 0.6,\n",
       " 'alpha_tr': 1.0960392096062852,\n",
       " 'alpha_lo': 1.0718295294749038}"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# s12, alpha_tr, alpha_lo and f are computed internally!\n",
    "EFT.params"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5fffaf91-6e1b-4e76-8aec-2c85ef78377b",
   "metadata": {},
   "source": [
    "First, we need to redefine our parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "9e61118a-75a4-4f89-b8c8-270d3e391545",
   "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['alpha_lo'] = 1.1\n",
    "params['alpha_tr'] = 0.9\n",
    "params['f']        = 0.7"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fec79fa8-db9e-4ac8-9da9-64848a310177",
   "metadata": {},
   "source": [
    "_Note: When computing the multipoles using the $\\sigma_{12}$ parameter space we need to specify a fiducial value for the Hubble rate. This is required to convert the native emulator output from Mpc to Mpc/h units._"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "58d427d8-2d5d-4d8a-8fb8-159ef56abf20",
   "metadata": {},
   "outputs": [],
   "source": [
    "Pell_s12 = EFT.Pell(k_hMpc, params, ell=[0,2,4])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "581f22ec-b7ae-4c66-b258-de3f2ebe7f91",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 53468 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",
    "\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": "7d6d7c8c-21e6-4e72-a781-4d8f3f6834af",
   "metadata": {},
   "source": [
    "## How to provide different k-scales."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "349b5b7b-55ff-4de8-86a8-6ea8039bcd0e",
   "metadata": {},
   "source": [
    "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, the length must match the number of specified multipoles (`ell`) and the first entry of the list is evaluated for the first multipole etc.\n",
    "\n",
    "We can output at a single scale and single multipole number, e.g. for the quadrupole at $k = 0.1 1/Mpc$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "11d44074-613e-4922-b598-e7960058b8e4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'ell2': array([12735.02642573])}"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# The provided scales are also assumed to be in 1/Mpc units\n",
    "EFT.Pell(0.1, params, ell=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d25aba8-1971-45f3-a5fa-d691a13313d7",
   "metadata": {},
   "source": [
    "Or for various multipoles and multiple scales, in which case the output is a list with the first entry corresponding to the first multipole specified in `ell`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "4ab21bae-111b-4b11-a88c-be5cc9f2768c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'ell0': array([21993.63231466,  8419.32061295,  5052.40597447]),\n",
       " 'ell2': array([12735.02642573,  7163.41410577,  5357.67022233]),\n",
       " 'ell4': array([3027.55903402, 2244.05437432, 1870.91419438])}"
      ]
     },
     "execution_count": 20,
     "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": "d3b177cf-c736-42dd-ac3b-0fb0139b6b49",
   "metadata": {},
   "source": [
    "Or at different scales for different multipoles (providing a list of numbers or numpy arrays):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "07c6a121-6646-4a62-bdc2-e05d6dfd97a4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'ell0': array([21993.63231466,  8419.32061295]),\n",
       " 'ell4': array([1870.91419438])}"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "EFT.Pell([np.array([0.1,0.2]),0.3], params, ell=[0,4])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9931c7b9-d809-4192-9c1b-04486a34a665",
   "metadata": {},
   "source": [
    "## Speedup when changing just bias parameters."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1bd0f8ad-10df-40a4-9881-d1afaa6bf674",
   "metadata": {},
   "source": [
    "It is a common task to test the models just by changing parameters that does not involve any cosmological computation. On that case we can call the function `Pell_fixed_cosmo_boost`, which looks into the parameters specified and if any cosmological parameter has changed, it uses the computation from previous calls. 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": 30,
   "id": "a8413952-89be-4b9b-82e2-b3b6bc11c193",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "22.3 ms ± 676 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit EFT.Pell(k_hMpc, params, ell=[0,2,4], de_model=\"lambda\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "a5cbfd73-5829-4db5-b351-c85b92f3d619",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "22.8 µs ± 3.19 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%timeit EFT.Pell_fixed_cosmo_boost(k_hMpc, params, ell=[0,2,4], de_model=\"lambda\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3f6f0cf-308f-4a5c-ba9c-5cd3ad392bbb",
   "metadata": {},
   "source": [
    "# Beyond $P_{\\ell}$ predictions."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd329fa2-73a5-4d44-bb01-598906a40e45",
   "metadata": {},
   "source": [
    "We have included some required tools needed for cosmological analysis in order to make easiert to integrate it into a complete pipeline"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "473d7d6f-5144-42ac-a33c-b0faaded2a24",
   "metadata": {},
   "source": [
    "## Computing covariance matrices"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d4c4df8-2b98-4ced-b172-59bec7fe0258",
   "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 $\\Lambda$CDM parameter spaces. 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": 32,
   "id": "c3353c08-7356-4060-ab9a-9d6776e7fed0",
   "metadata": {},
   "outputs": [],
   "source": [
    "dk_hMpc = 0.001\n",
    "k_hMpc_lin = np.arange(0.001, 0.3, dk_hMpc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "18872cb3-2072-446d-8547-aafd25c11258",
   "metadata": {},
   "outputs": [],
   "source": [
    "Cov_hMpc = EFT.Pell_covariance(k_hMpc, params, ell=[0,2,4], dk=dk_hMpc, volume=3e9)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "d1bf1550-31a2-4f2e-a0fd-407eb565e905",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 34348 characters]\n",
      "text/plain": [
       "<Figure size 648x432 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.figure(figsize=(9,6))\n",
    "plt.title(r\"\")\n",
    "plt.title(r\"Corr matrix\", y=1.05)\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', interpolation=\"blackman\", filterrad=10)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb264fee-3dcf-4cca-a2aa-572c62d1a1c2",
   "metadata": {},
   "source": [
    "For the $\\Lambda$CDM version, instead of providing a volume, we can provide minimum and maximum redshifts, `zmin` and `zmax`, a sky fraction `fsky`, and a volume scaling factor `volfac` (by default equal to 1), which computes then evalutes the corresponding volume, also assuming a flat $\\Lambda$CDM model. For example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "95789e3f-5647-45b8-8a4e-4937340cf7c5",
   "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",
    "                        volume=3780**3\n",
    "                )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "82b2f79a-3b5c-411c-9ab5-a3f514e9e242",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 31136 characters]\n",
      "text/plain": [
       "<Figure size 648x432 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.figure(figsize=(9,6))\n",
    "plt.title(r\"\")\n",
    "plt.title(r\"Corr matrix\", y=1.05)\n",
    "var_inv = np.diag(1./np.sqrt(np.diag(Cov_hMpc_LCDM)))\n",
    "R_hMpc = var_inv @ Cov_hMpc_LCDM @ var_inv\n",
    "plt.imshow(R_hMpc,cmap='magma_r', interpolation=\"blackman\", filterrad=10)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92520666-6832-46c3-b916-809a4c7916fc",
   "metadata": {},
   "source": [
    "## Computing the $\\chi^2$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3dcb51d7-047a-44f4-9e65-85ffbfd26643",
   "metadata": {},
   "source": [
    "Finally, the Comet class provides a function to directly compute the $\\chi^2$. In order to do so, we need to specify the data set, i.e., range of scales of the observations, the measured multipoles and their covariance matrix. As before, these must be given in the same unit system as that for which the emulator has been configured."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "24a56597-e0ac-4d82-9c78-1850b80187d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Loading a sample data set:\n",
    "data = np.loadtxt('../../../../data/cmass/Minerva_HOD_zs_z0.57_pkmulti_mean.dat')\n",
    "cov = np.loadtxt('../../../../data/cmass/Minerva_HOD_zs_z0.57_pkmulti_covar.dat')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cfc05a3a-9796-4930-93b2-01e8df31fe8c",
   "metadata": {},
   "source": [
    "As the data vector would have an assosiated shot noise, we can specify it by indicating the number density of the sample."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "4dd2cc1c-d9ae-4160-b7da-30ef7f8c2095",
   "metadata": {},
   "outputs": [],
   "source": [
    "EFT.define_nbar(nbar=3.95898e-4)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c412bb8b-d363-4e3d-96c2-4ec1bed06f34",
   "metadata": {},
   "source": [
    "We have the option of specifying whether the covariance matrix is a \"theory\" covariance matrix or not. If `theory_cov = False`, the Anderson-Hartlap factor is included in the inverse covariance matrix, which is why we need to also provide the number of n_realizations from which the covariance matrix was estimated (if `theory_cov = True`, `Nrealizations` can be ignored)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "dba08e32-10a5-4424-bba3-315885197882",
   "metadata": {},
   "outputs": [],
   "source": [
    "EFT.define_data_set(obs_id='Pk', bins=data[:,0], signal=data[:,(1,3,5)], cov=cov, theory_cov=False, n_realizations=300)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57aa7e2d-ed7f-4879-8717-faa1633390c9",
   "metadata": {},
   "source": [
    "Now we can call `chi2`, which takes as arguments 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 specified before. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "6fc37422-5f3c-406a-816a-6becfa0f6969",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "9453.159316045188"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "EFT.chi2(obs_id='Pk',params=params, kmax=[0.3,0.30, 0.30], de_model='lambda', chi2_decomposition=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ac9332b-b1bb-420c-9395-17ca7476e290",
   "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 recompute the cosmological dependent quantities. Let's see how it works"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "ee80c0e6-20e2-4c40-9ad2-8fbc9d264b31",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "23 ms ± 538 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit EFT.chi2(obs_id='Pk',params=params, kmax=[0.3,0.30, 0.30], de_model='lambda', chi2_decomposition=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "999268cb-f0a7-43e8-8646-e0e168ba6dac",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "20.7 µs ± 470 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%timeit EFT.chi2(obs_id='Pk',params=params, kmax=[0.3,0.30, 0.30], de_model='lambda', chi2_decomposition=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50a626d2-f938-446a-8616-4e00acf647eb",
   "metadata": {},
   "outputs": [],
   "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.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
