{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "e87edff3-c6c0-434f-876b-28f12d376a14",
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy.integrate\n",
    "from classy import Class\n",
    "from classy import CosmoComputationError\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "445fc9a0-6bb9-45aa-8ea2-726bada97199",
   "metadata": {},
   "outputs": [],
   "source": [
    "def class_compute(class_params):\n",
    "    '''\n",
    "    A function to handle CLASS computation and deal with potential errors.\n",
    "\n",
    "    Args:\n",
    "        :class_params (dict): Dictionary of CLASS parameters\n",
    "\n",
    "    Returns:\n",
    "        :cosmo (CLASS): Instance of the CLASS code\n",
    "        :isNormal (bool): Whether error occurred in the computation\n",
    "    '''\n",
    "    cosmo = Class()\n",
    "    cosmo.set(class_params)\n",
    "    try:\n",
    "        cosmo.compute()\n",
    "        isNormal=True\n",
    "    except CosmoComputationError as e:\n",
    "        if \"DeltaNeff < deltaN[0]\" in str(e):\n",
    "            # set YHe to 0.25. Value is from https://arxiv.org/abs/1503.08146 and Plank 2018(Section 7.6.1) https://arxiv.org/abs/1807.06209\n",
    "            warnings.warn(f\"Adjust YHe to 0.25 due to CLASS CosmoComputationError for cosmology {class_params}.\")\n",
    "            class_params['YHe'] = 0.25\n",
    "            cosmo.set(class_params)\n",
    "            cosmo.compute()\n",
    "            isNormal=False\n",
    "        else:\n",
    "            raise e\n",
    "    return cosmo, isNormal\n",
    "\n",
    "\n",
    "def get_class_linear(k, As, Om, Ob, h, ns):\n",
    "    \"\"\"\n",
    "    Compute linear P(k) using class for the cosmology of interest\n",
    "    \n",
    "    Args:\n",
    "        :k (np.ndarray): k values to evaluate P(k) at [h / Mpc]\n",
    "        :As (float): 10^9 times the amplitude of the primordial P(k)\n",
    "        :Om (float): The z=0 total matter density parameter, Omega_m\n",
    "        :Ob (float): The z=0 baryonic density parameter, Omega_b\n",
    "        :h (float): Hubble constant, H0, divided by 100 km/s/Mpc\n",
    "        :ns (float): Spectral tilt of primordial power spectrum\n",
    "        \n",
    "    Returns:\n",
    "        :plin_class (np.ndarray): Linear power spectrum at corresponding k values [(Mpc/h)^3]\n",
    "        :isNormal (bool): Whether error occurred in the computation\n",
    "    \"\"\"\n",
    "    \n",
    "    Oc =  Om - Ob\n",
    "    redshift = 0.0\n",
    "    \n",
    "    class_params = {\n",
    "        'h': h,\n",
    "        'omega_b': Ob * h**2,\n",
    "        'omega_cdm': Oc * h**2,\n",
    "        'A_s': As*1.e-9,\n",
    "        'n_s': ns,\n",
    "        'output': 'mPk',\n",
    "        'P_k_max_1/Mpc': k.max() * h,\n",
    "        'w0_fld': -1,\n",
    "        'wa_fld': 0,\n",
    "        'Omega_Lambda': 0,  # Set to 0 because we're using w0_fld and wa_fld instead\n",
    "        'tau_reio':0.0561,\n",
    "        'z_max_pk': 3.0,  # Max redshift for P(k) output\n",
    "    }\n",
    "    \n",
    "    cosmo, isNormal = class_compute(class_params)\n",
    "    plin_class = np.array([cosmo.pk_lin(kk*h, redshift) for kk in k]) * h ** 3\n",
    "    \n",
    "    # Memory cleanup for class\n",
    "    cosmo.struct_cleanup()\n",
    "    cosmo.empty()\n",
    "\n",
    "    return plin_class, isNormal\n",
    "\n",
    "\n",
    "def get_sig(R, As, Om, Ob, h, ns):\n",
    "    \"\"\"\n",
    "    Compute the bulk flow variance when the field is smoothed with a top\n",
    "    hat filter of side length R.\n",
    "    \n",
    "    sigma^2 = H_0^2 f^2 / (2 \\pi^2) \\int_0^\\infty dk W^2(k,R) P(k)\n",
    "    \n",
    "    where W(k,R) is the Fourier transform of the top hat filter\n",
    "    \n",
    "    Args:\n",
    "        :k (np.ndarray): k values to evaluate P(k) at [h / Mpc]\n",
    "        :As (float): 10^9 times the amplitude of the primordial P(k)\n",
    "        :Om (float): The z=0 total matter density parameter, Omega_m\n",
    "        :Ob (float): The z=0 baryonic density parameter, Omega_b\n",
    "        :h (float): Hubble constant, H0, divided by 100 km/s/Mpc\n",
    "        :ns (float): Spectral tilt of primordial power spectrum\n",
    "        \n",
    "    Returns:\n",
    "        :sigma (float): Bulk flow RMS [km/s]\n",
    "        \n",
    "    \"\"\"\n",
    "    \n",
    "    k = np.logspace(-4, 2, 300) # (h / Mpc)\n",
    "    plin_class, isNormal = get_class_linear(k, As, Om, Ob, h, ns) # (Mpc/h)^3\n",
    "    x = k * R\n",
    "    integrand = plin_class * (3 * (np.sin(x) - x * np.cos(x)) / x ** 3) ** 2 # (Mpc/h)^3\n",
    "    \n",
    "    # Check units\n",
    "    sigma = scipy.integrate.simpson(integrand, x=k) # (Mpc/h)^2\n",
    "    H0 = 100 # h km/s/Mpc\n",
    "    f = Om ** 0.55\n",
    "    sigma = np.sqrt(sigma / 2) * H0 * f / np.pi # (Mpc/h) * h km/s/Mpc = km/s\n",
    "\n",
    "    return sigma"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "0e09f7ec-92e7-474c-b9a1-46094ab7385f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:11<00:00,  1.73it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "58.40857421174422\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 11432 characters]",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "As = 2.105  # 10^9 A_s\n",
    "h = 0.6766\n",
    "Om = 0.3111\n",
    "Ob = 0.02242 / h ** 2\n",
    "ns = 0.9665\n",
    "\n",
    "As = 1.7988  # 10^9 A_s\n",
    "h = 0.68\n",
    "Om = 0.335\n",
    "Ob = 0.049\n",
    "ns = 0.97\n",
    "\n",
    "all_R = np.linspace(50, 1000, 20)\n",
    "all_sigma = np.empty(len(all_R))\n",
    "\n",
    "for i in tqdm(range(len(all_sigma))):\n",
    "    all_sigma[i] = get_sig(all_R[i], As, Om, Ob, h, ns)\n",
    "plt.plot(all_R, all_sigma)\n",
    "\n",
    "print(get_sig(500, As, Om, Ob, h, ns))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9581f953-0bb4-4c71-a1ba-32ade3a95e27",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "symreg",
   "language": "python",
   "name": "symreg"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
