linear_Pk_emu_PCA.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import camb\n",
    "from camb import model, initialpower"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook, we will train our iterative emulator on a simple scenario of linear Power spectrum, i.e, we assume that our data vector is the _linear_ power spectrum, $P(k)$. \n",
    "\n",
    "We will use our neural network emulator to predict the power spectrum as a function of cosmology. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We use `camb` to predict the linear power spectrum, as defined below. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_Pk(cosmo_pars):\n",
    "    As, ns, H0, ombh2, omch2 = cosmo_pars\n",
    "    #Now get matter power spectra and sigma8 at redshift 0 and 0.8\n",
    "    pars = camb.CAMBparams()\n",
    "    pars.set_cosmology(H0=H0, ombh2=ombh2, omch2=omch2)\n",
    "    pars.InitPower.set_params(As=As, ns=ns)\n",
    "    #Note non-linear corrections couples to smaller scales than you want\n",
    "    pars.set_matter_power(redshifts=[0.], kmax=2.0)\n",
    "\n",
    "    #Linear spectra\n",
    "    pars.NonLinear = model.NonLinear_none\n",
    "    results = camb.get_results(pars)\n",
    "    kh, z, pk = results.get_matter_power_spectrum(minkh=1e-3, maxkh=1, npoints = 2000)\n",
    "    s8 = np.array(results.get_sigma8())\n",
    "    \n",
    "    return kh, pk[0], s8"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To use the emulator, we need a function to compute the data vector as a function of cosmological parameters. \n",
    "\n",
    "For your specific case, you can simply replace this function along with the log probability function to use the emulator as proposed here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_datavector(cosmo_pars):\n",
    "    kh, pk, _ = get_Pk(cosmo_pars)\n",
    "    # We use the logarithm of the power spectrum as it is easier to train\n",
    "    return np.log(pk)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We use the following fiducial parameters in our analysis\n",
    "\n",
    "cosmo_pars_fid = np.array([2e-9, 0.97, 70., 0.0228528, 0.1199772])\n",
    "\n",
    "kh_fid, pk_fid, _ = get_Pk(cosmo_pars_fid)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialize a Latin Hypercube"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We begin our emulator training, we need a Latin Hypercube sample with the prior range."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyDOE import lhs\n",
    "\n",
    "N_dim = 5\n",
    "\n",
    "cosmo_prior = np.array([[1.2e-9, 2.7e-9], # As\n",
    "                       [0.87, 1.07],      # ns\n",
    "                       [60, 85],          # H0\n",
    "                       [0.01, 0.04],      # omega_b\n",
    "                       [0.01, 0.3]])      # omega_c\n",
    "\n",
    "def get_cosmo_lhs_samples(N_samples, cosmo_prior):\n",
    "    lhs_samples = lhs(N_dim, N_samples)\n",
    "    cosmo_samples = cosmo_prior[:,0] + (cosmo_prior[:,1] - cosmo_prior[:,0]) * lhs_samples\n",
    "    return cosmo_samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "from multiprocessing import Pool\n",
    "\n",
    "def calculate_datavector_batch(cosmo_samples):\n",
    "    \"\"\"\n",
    "    Function to calculate the data vectors for a batch of training samples\n",
    "    \"\"\"\n",
    "    train_dv_list = []\n",
    "    with Pool() as p:\n",
    "        train_dv_list = list(tqdm(p.imap(compute_datavector, cosmo_samples), total=len(cosmo_samples)))\n",
    "    return np.array(train_dv_list)    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1000/1000 [06:31<00:00,  2.55it/s]\n"
     ]
    }
   ],
   "source": [
    "N_train_samples = 1000\n",
    "\n",
    "train_cosmo_samples = get_cosmo_lhs_samples(N_train_samples, cosmo_prior)\n",
    "\n",
    "train_dv_arr = calculate_datavector_batch(train_cosmo_samples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 400/400 [02:39<00:00,  2.52it/s]\n"
     ]
    }
   ],
   "source": [
    "N_test_samples = 400\n",
    "\n",
    "test_cosmo_samples = get_cosmo_lhs_samples(N_test_samples, cosmo_prior)\n",
    "test_dv_arr = calculate_datavector_batch(test_cosmo_samples)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### PCA"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Take the PCAs of the training samples and then train your emulator to predict the PCAs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.decomposition import PCA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_Pk(pca_coeffs, pca):\n",
    "    \"\"\"\n",
    "    Function to get the power spectrum from the PCA coefficients\n",
    "    :pca_coefficient: PCA coefficients\n",
    "    :pca: sklearn PCA object\n",
    "    \"\"\"\n",
    "    log_Pk_pred  = pca.inverse_transform(pca_coeffs)\n",
    "    return np.exp(log_Pk_pred)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_pca = 5\n",
    "\n",
    "pca = PCA(n_pca)\n",
    "pca.fit(train_dv_arr)\n",
    "\n",
    "train_pca_coeff = pca.transform(train_dv_arr)\n",
    "\n",
    "train_pca_mean = train_pca_coeff.mean(0)\n",
    "train_pca_std  = train_pca_coeff.std(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "from emulator import NNEmulator\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loss: 0.08990950137376785:  42%|████▏     | 42/100 [00:27<00:29,  1.95it/s] "
     ]
    }
   ],
   "source": [
    "emu = NNEmulator(N_dim, n_pca, train_pca_mean, train_pca_std)\n",
    "emu.train(torch.Tensor(train_cosmo_samples), torch.Tensor(train_pca_coeff))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_pca_pred = emu.predict(torch.Tensor(test_cosmo_samples))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_pk_pred = get_Pk(test_pca_pred, pca)\n",
    "test_pk_arr  = np.exp(test_dv_arr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DeltaP_frac = np.abs((test_pk_pred - test_pk_arr) / test_pk_arr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.ylim(0., 0.15)\n",
    "plt.semilogx(kh_fid, np.median(DeltaP_frac, axis=0))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With the emulator to predict 5 PCA components, we get median accuracy of $\\sim \\mathcal{O}(5\\%)$. Increasing the number of training samples will likely improve the accuracy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cocoa",
   "language": "python",
   "name": "cocoa"
  },
  "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.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}