Open_access_of_2e5_emulator_NN_tt.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/keirkwame/axionEmu/blob/main/Open_access_of_2e5_emulator_NN_tt.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "t31B4nK6CyIs",
    "outputId": "174b258c-5d1c-4da0-d443-df38e477add6"
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "\n",
    "# setting the seed for reproducibility\n",
    "np.random.seed(2)\n",
    "tf.random.set_seed(4)\n",
    "\n",
    "#from google.colab import drive\n",
    "#drive.mount('/content/drive')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "DK3nZhT3C4g6"
   },
   "outputs": [],
   "source": [
    "model_parameters = ['h', \n",
    "                    'tau_reio', \n",
    "                    'omega_b', \n",
    "                    'n_s', \n",
    "                    'ln10A_s', \n",
    "                    'omega_cdm',\n",
    "                    'log10ma',\n",
    "                    'omega_ax' \n",
    "                    ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "8AU1QixoKwY8",
    "outputId": "0f4230e5-07c0-4507-d057-e822b248e281"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 0.,  2.,  4.,  6.,  8., 10., 12., 14., 16., 18., 20., 22., 24.,\n",
       "       26., 28.])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.linspace(0,28,15)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "RrjemD2nC6nu",
    "outputId": "3ce57769-fe57-4c61-f6ae-369e4c160309"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2000\n",
      "2000\n",
      "2000\n",
      "2000\n",
      "2000\n",
      "2000\n",
      "2000\n",
      "2000\n",
      "2000\n",
      "2000\n",
      "2000\n",
      "2000\n",
      "2000\n",
      "2000\n",
      "2000\n"
     ]
    }
   ],
   "source": [
    "import pickle\n",
    "collection_list = []\n",
    "for i in np.linspace(0,28,15):\n",
    "    f = open('/home/anran/axionCAMB/test_data_collect_9params_5e5_mp_test_'+str(int(i))+'.pkl', 'rb')\n",
    "    collection = pickle.load(f)\n",
    "    f.close()\n",
    "    print(len(collection['C_tt']))\n",
    "    collection_list.append(collection)\n",
    "    \n",
    "import random\n",
    "random.shuffle(collection_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "JhpiDxrRKbuw",
    "outputId": "144824b3-1187-4343-8b44-d6e59b428ef2"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2000\n",
      "2000\n",
      "2000\n",
      "2000\n",
      "2000\n",
      "2000\n",
      "2000\n",
      "2000\n",
      "2000\n",
      "2000\n",
      "2000\n",
      "2000\n",
      "2000\n",
      "2000\n",
      "2000\n"
     ]
    }
   ],
   "source": [
    "import pickle\n",
    "for i in np.linspace(0,28,15):\n",
    "    f = open('/home/anran/axionCAMB/test_data_collect_9params_5e5_mp_test_'+str(int(i)+30)+'.pkl', 'rb')\n",
    "    collection = pickle.load(f)\n",
    "    f.close()\n",
    "    print(len(collection['C_tt']))\n",
    "    collection_list.append(collection)\n",
    "    \n",
    "import random\n",
    "random.shuffle(collection_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "19BWUXvODCts",
    "outputId": "990bd519-0a55-4fc6-fe96-fc2eae7bda28"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "parameters_list = {}\n",
    "for key in collection['params']:\n",
    "    parameters_list[key] = np.array([])\n",
    "\n",
    "C_tt_list = []\n",
    "\n",
    "for i in range(30):\n",
    "    para = collection_list[i]['params']\n",
    "    C_tt_list.extend(collection_list[i]['C_tt'])\n",
    "    for key in para:\n",
    "      parameters_list[key] = np.concatenate((parameters_list[key], para[key]))\n",
    "cut_off = 54000\n",
    "import copy\n",
    "training_parameters_ = copy.deepcopy(parameters_list)\n",
    "h_0 = training_parameters_['H_0'][:cut_off]/100\n",
    "ln10_10A_s = np.log(training_parameters_['A_s'][:cut_off]*10**10)\n",
    "ma_mass = np.array(training_parameters_['ma'][:cut_off]*10**32,dtype='float64')\n",
    "ma_mass = np.log10(ma_mass)\n",
    "training_parameters = dict()\n",
    "training_parameters['omega_b'] = training_parameters_['omega_b'][:cut_off]\n",
    "training_parameters['omega_cdm'] = training_parameters_['omega_cdm'][:cut_off]\n",
    "training_parameters['h'] = h_0\n",
    "training_parameters['tau_reio'] = training_parameters_['tau_reio'][:cut_off]\n",
    "training_parameters['n_s'] = training_parameters_['n_s'][:cut_off]\n",
    "training_parameters['ln10A_s'] = ln10_10A_s\n",
    "training_parameters['log10ma'] = ma_mass\n",
    "training_parameters['omega_ax'] = training_parameters_['omega_ax'][:cut_off]\n",
    "ell_range = np.linspace(2,2508, 2507)\n",
    "\n",
    "test_parameters_ = copy.deepcopy(parameters_list)\n",
    "h_0 = test_parameters_['H_0'][cut_off:]/100\n",
    "ln10_10A_s = np.log(test_parameters_['A_s'][cut_off:]*10**10)\n",
    "ma_mass = np.array(test_parameters_['ma'][cut_off:]*10**32, dtype = 'float64')\n",
    "ma_mass = np.log10(ma_mass)\n",
    "test_parameters = dict()\n",
    "test_parameters['omega_b'] = test_parameters_['omega_b'][cut_off:]\n",
    "test_parameters['omega_cdm'] = test_parameters_['omega_cdm'][cut_off:]\n",
    "test_parameters['h'] = h_0\n",
    "test_parameters['tau_reio'] = test_parameters_['tau_reio'][cut_off:]\n",
    "test_parameters['n_s'] = test_parameters_['n_s'][cut_off:]\n",
    "test_parameters['ln10A_s'] = ln10_10A_s\n",
    "test_parameters['omega_ax'] = test_parameters_['omega_ax'][cut_off:]\n",
    "test_parameters['log10ma'] = ma_mass\n",
    "\n",
    "spectra_= np.array(C_tt_list)/(7.4311*10**(12))\n",
    "print(spectra_.shape)\n",
    "spectra_ = spectra_[:,:2507]/(ell_range*(ell_range+1)/(2.*np.pi))\n",
    "training_log_spectra = np.log10(spectra_)[:cut_off,:]\n",
    "testing_spectra = spectra_[cut_off:,:]\n",
    "print('number of training spectra: ', training_log_spectra.shape)\n",
    "print('number of testing spectra: ', testing_spectra.shape)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "y-SR8DFGQLPD"
   },
   "outputs": [],
   "source": [
    "test_collection = {}\n",
    "test_collection['training_parameters'] = training_parameters\n",
    "test_collection['test_parameters'] = test_parameters\n",
    "test_collection['training_log_spectra'] = training_log_spectra\n",
    "test_collection['testing_spectra'] = testing_spectra\n",
    "import pickle\n",
    "f = open('test_collection.pkl', 'wb')\n",
    "pickle.dump(test_collection, f)\n",
    "f.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "jZXssxsdnrJY",
    "outputId": "fad62178-9bfa-4fc9-a210-bbfa9e2e35a9"
   },
   "outputs": [],
   "source": [
    "##Only for Google Colab\n",
    "!git clone https://github.com/keirkwame/axionEmu.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "APuVRpykD9Hf",
    "outputId": "6c02fa8a-cc58-4f30-81fe-d2ca6381e17b"
   },
   "outputs": [],
   "source": [
    "from axionEmu.cosmopower import cosmopower_NN\n",
    "\n",
    "# instantiate NN class\n",
    "cp_nn = cosmopower_NN(parameters=model_parameters, \n",
    "                      modes=ell_range, \n",
    "                      n_hidden = [512, 512, 512,512], # 4 hidden layers, each with 512 nodes\n",
    "                      verbose=True, # useful to understand the different steps in initialisation and training\n",
    "                      )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "myZDPxFVD_gL",
    "outputId": "4eebcb70-2a4e-4639-f203-71da3290b0e1",
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "with tf.device('/device:GPU:0'): # ensures we are running on a GPU\n",
    "    # train\n",
    "    cp_nn.train(training_parameters=training_parameters,\n",
    "                training_features=training_log_spectra,\n",
    "                filename_saved_model='TT_cp_NN_2e5_test_30_t0',\n",
    "                # cooling schedule\n",
    "                validation_split=0.1,\n",
    "                learning_rates=[1e-2, 1e-3, 1e-4, 1e-5, 1e-6],\n",
    "                batch_sizes=[1024,1024,1024,1024,1024],\n",
    "                gradient_accumulation_steps = [1, 1, 1, 1, 1],\n",
    "                # early stopping set up\n",
    "                patience_values = [100,100,100,100,100],\n",
    "                max_epochs = [1000,1000,1000,1000,1000],\n",
    "                )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "A8LpGeVbD_mP"
   },
   "outputs": [],
   "source": [
    "cp_nn = cosmopower_NN(restore=True, \n",
    "                      restore_filename='TT_cp_NN_2e5_test_30_t0',\n",
    "                      )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 224
    },
    "id": "0Ktj-D7zFKLt",
    "outputId": "7b7a2cf5-9f7f-4cc5-ffca-998949a1595b"
   },
   "outputs": [],
   "source": [
    "predicted_testing_spectra = cp_nn.ten_to_predictions_np(test_parameters)\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import gridspec\n",
    "fig, ax = plt.subplots(nrows=1, ncols=3, figsize=(50,10))\n",
    "for i in range(3):\n",
    "    pred = predicted_testing_spectra[i]*ell_range*(ell_range+1)/(2.*np.pi)\n",
    "    true = testing_spectra[i]*ell_range*(ell_range+1)/(2.*np.pi)\n",
    "    ax[i].semilogx(ell_range, true, 'blue', label = 'Original')\n",
    "    ax[i].semilogx(ell_range, pred, 'red', label = 'NN reconstructed', linestyle='--')\n",
    "    ax[i].set_xlabel('$\\ell$', fontsize='x-large')\n",
    "    ax[i].set_ylabel('$\\\\frac{\\ell(\\ell+1)}{2 \\pi} C_\\ell$', fontsize='x-large')\n",
    "    ax[i].legend(fontsize=15)\n",
    "plt.savefig('examples_reconstruction_TT_2e5_30_0.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "BjvhpdYCF2TV",
    "outputId": "d3040cda-c742-433e-e476-1790169c00f7"
   },
   "outputs": [],
   "source": [
    "##Only for Google Colab\n",
    "!git clone https://github.com/simonsobs/so_noise_models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 869
    },
    "id": "7gYkkDJCFNmm",
    "outputId": "e398fdd3-b53b-4ffd-88d1-9f3ce5018e87"
   },
   "outputs": [],
   "source": [
    "# load noise models from the SO noise repo\n",
    "noise_levels_load = np.loadtxt('./so_noise_models/LAT_comp_sep_noise/v3.1.0/SO_LAT_Nell_T_atmv1_goal_fsky0p4_ILC_CMB.txt')\n",
    "conv_factor = (2.7255e6)**2\n",
    "\n",
    "ells = noise_levels_load[:, 0]\n",
    "SO_TT_noise = noise_levels_load[:, 1][:2509-40] / conv_factor ##6001-40\n",
    "new_ells = ells[:2509-40] ##6001-40\n",
    "\n",
    "f_sky = 0.4\n",
    "prefac = np.sqrt(2/(f_sky*(2*new_ells+1)))\n",
    "\n",
    "##Cut off some spectra\n",
    "print(new_ells)\n",
    "print(prefac.shape, testing_spectra.shape, SO_TT_noise.shape)\n",
    "#testing_spectra = testing_spectra[:6001-40]\n",
    "\n",
    "denominator = prefac*(testing_spectra[:, 38:]+SO_TT_noise)  # use all of them\n",
    "diff = np.abs((predicted_testing_spectra[:, 38:] - testing_spectra[:, 38:])/(denominator))\n",
    "\n",
    "# Compute percentiles\n",
    "percentiles = np.zeros((4, diff.shape[1]))\n",
    "\n",
    "percentiles[0] = np.percentile(diff, 68, axis = 0)\n",
    "percentiles[1] = np.percentile(diff, 95, axis = 0)\n",
    "percentiles[2] = np.percentile(diff, 99, axis = 0)\n",
    "percentiles[3] = np.percentile(diff, 99.9, axis = 0)\n",
    "\n",
    "plt.figure(figsize=(12, 9))\n",
    "plt.fill_between(new_ells, 0, percentiles[2,:], color = 'salmon', label = '99%', alpha=0.8)\n",
    "plt.fill_between(new_ells, 0, percentiles[1,:], color = 'red', label = '95%', alpha = 0.7)\n",
    "plt.fill_between(new_ells, 0, percentiles[0,:], color = 'darkred', label = '68%', alpha = 1)\n",
    "\n",
    "# plt.ylim(0, 0.2)\n",
    "\n",
    "plt.legend(frameon=False, fontsize=30, loc='upper left')\n",
    "plt.ylabel(r'$\\frac{| C_{\\ell, \\rm{emulated}}^{\\rm{TT}} - C_{\\ell, \\rm{true}}^{\\rm{TT}}|} {\\sigma_{\\ell, \\rm{CMB}}^{\\rm{TT}}}$', fontsize=50)\n",
    "plt.xlabel(r'$\\ell$',  fontsize=50)\n",
    "\n",
    "ax = plt.gca()\n",
    "ax.xaxis.set_major_locator(plt.MaxNLocator(10))\n",
    "ax.yaxis.set_major_locator(plt.MaxNLocator(5))\n",
    "\n",
    "plt.setp(ax.get_xticklabels(), fontsize=25)\n",
    "plt.setp(ax.get_yticklabels(), fontsize=25)\n",
    "plt.tight_layout()\n",
    "plt.savefig('./accuracy_emu_TT_wide_2e5_test_t0.pdf')\n",
    "\n",
    "diff_=np.sum(diff,axis = 1)\n",
    "print('diff_.shape is ', diff_.shape)\n",
    "sort_index = np.argsort(diff_)\n",
    "bad_params = dict()\n",
    "for key in test_parameters:\n",
    "    bad_params[key] = []\n",
    "for i in sort_index[len(sort_index)-100:]:\n",
    "    for key in test_parameters:\n",
    "        bad_params[key].append(test_parameters[key][i])\n",
    "fig, ax = plt.subplots(nrows=1, ncols=8, figsize=(50,10))\n",
    "for i in range(8):\n",
    "    pa =list(test_parameters.keys())[i]\n",
    "    h = bad_params[pa]\n",
    "    ax[i].plot(h, 'o')\n",
    "    ax[i].set_ylabel(pa, fontsize='x-large')\n",
    "plt.savefig('bad_params_2e5_0.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "include_colab_link": true,
   "machine_shape": "hm",
   "provenance": []
  },
  "gpuClass": "standard",
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}