{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy.integrate\n",
    "import symbolic_pofk.linear\n",
    "import symbolic_pofk.syrenhalofit as syrenhalofit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_sigma8_from_pk(k, pk):\n",
    "    \"\"\"Given a power spectrum P(k), compute sigma8.\"\"\"\n",
    "    R = 8.0\n",
    "    x = k * R\n",
    "    W = np.zeros(x.shape)\n",
    "    m = x < 1.e-3\n",
    "    W[m] = 1.0\n",
    "    W[~m] =3.0 / x[~m]**3 * (np.sin(x[~m]) - x[~m] * np.cos(x[~m]))\n",
    "    y = pk * W**2 * k**3\n",
    "    sigma2 = scipy.integrate.simpson(y, x=np.log(x))\n",
    "    sigma = np.sqrt(sigma2 / (2.0 * np.pi**2))\n",
    "\n",
    "    return sigma\n",
    "\n",
    "# Cosmological parameters\n",
    "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",
    "tau = 0.0561\n",
    "\n",
    "# Define k integration range\n",
    "\n",
    "\n",
    "def linear_sigma8(As, Om, Ob, h, ns):\n",
    "    \"\"\"Calculated from Deaglan's emulator.\"\"\"\n",
    "    return symbolic_pofk.linear.As_to_sigma8(As, Om, Ob, h, ns)\n",
    "    # print('sigma8 from emulator', sigma8)\n",
    "\n",
    "# # Test linear sigma8\n",
    "# pk_lin = symbolic_pofk.linear.plin_emulated(k, sigma8, Om, Ob, h, ns,\n",
    "#     emulator='fiducial', extrapolate=True)\n",
    "# new_sigma8 = compute_sigma8(k, pk_lin)\n",
    "# print('sigma8 from integral:', new_sigma8)\n",
    "\n",
    "# Get non-linear sigma8\n",
    "def nonlinear_sigma8(As, Om0, Ob, h, ns, ks):\n",
    "    a = 1.\n",
    "    sigma8 = linear_sigma8(As, Om0, Ob, h, ns)  # Linear sigma8\n",
    "    pk_nl = syrenhalofit.run_halofit(\n",
    "        ks, sigma8, Om, Ob, h, ns, a, emulator='fiducial', extrapolate=True,\n",
    "        which_params='Bartlett', add_correction=True)\n",
    "    return compute_sigma8_from_pk(ks, pk_nl)\n",
    "# print('non-linear sigma8:', sigma8_nl)\n",
    "# print('sigma8 non-linear bigger by a factor', sigma8_nl / sigma8)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "kmin, kmax, nk = 1e-4, 1e1, 256\n",
    "ks = np.logspace(np.log10(kmin), np.log10(kmax), nk) # Wavenumber"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.825706107176727"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "linear_sigma8(As, Om, Ob, h, ns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.9203769215120938"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nonlinear_sigma8(As, Om, Ob, h, ns, ks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2.1005829616811546e-09"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    " 1e-10 * np.exp(3.0448)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv_csiborg",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}