optstat_tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Computing the optimal statistic with enterprise\n",
    "\n",
    "In this notebook you will learn how to compute the optimal statistic. The optimal statistic is a frequentist detection statistic for the stochastic background. It assesses the significance of the cross-correlations, and compares them to the Hellings-Downs curve.\n",
    "\n",
    "For more information, see Anholm et al. 2009, Demorest et al. 2013, Chamberlin et al. 2015, Vigeland et al. 2018.\n",
    "\n",
    "This notebook shows you how to compute the optimal statistic for the 12.5yr data set. You can download a pickle of the pulsars and the noisefiles here: https://paper.dropbox.com/doc/NG-12.5yr_v3-GWB-Analysis--A2vs2wHh5gR4VTgm2DeODR2zAg-DICJei6NxsPjxnO90mGMo. You will need the following files:\n",
    " * Channelized Pickled Pulsars (DE438) - Made in Py3\n",
    " * Noisefiles (make sure you get the one that says it contains all the pulsar parameters)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import (absolute_import, division,\n",
    "                        print_function, unicode_literals)\n",
    "import numpy as np\n",
    "import pickle\n",
    "import json\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "from enterprise.signals import signal_base\n",
    "from enterprise.signals import gp_signals\n",
    "\n",
    "from enterprise_extensions import model_utils, blocks\n",
    "from enterprise_extensions.frequentist import optimal_statistic as opt_stat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load up the pulsars from the pickle file\n",
    "# Change the picklefile to point to where you have saved the pickle of the pulsars that you downloaded\n",
    "picklefile = '/Users/vigeland/Documents/Research/NANOGrav/nanograv_data/12p5yr/channelized_v3_DE438_45psrs.pkl'\n",
    "\n",
    "with open(picklefile, 'rb') as f:\n",
    "    psrs = pickle.load(f)\n",
    "len(psrs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load up the noise dictionary to get values for the white noise parameters\n",
    "# Change the noisefile to point to where you have saved the noisefile\n",
    "noisefile = '/Users/vigeland/Documents/Research/NANOGrav/nanograv_data/12p5yr/channelized_12p5yr_v3_full_noisedict.json'\n",
    "\n",
    "with open(noisefile, 'r') as f:\n",
    "    noisedict = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the optimal statistic object\n",
    "# You can give it a list of pulsars and the noise dictionary, and it will create the pta object for you\n",
    "# Alternatively, you can make the pta object yourself and give it to the OptimalStatistic object as an argument\n",
    "\n",
    "# find the maximum time span to set GW frequency sampling\n",
    "Tspan = model_utils.get_tspan(psrs)\n",
    "\n",
    "# Here we build the signal model\n",
    "# First we add the timing model\n",
    "s = gp_signals.TimingModel()\n",
    "\n",
    "# Then we add the white noise\n",
    "# There are three types of white noise: EFAC, EQUAD, and ECORR\n",
    "# We use different white noise parameters for every backend/receiver combination\n",
    "# The white noise parameters are held constant\n",
    "s += blocks.white_noise_block(vary=False, inc_ecorr=True, select='backend')\n",
    "\n",
    "# Next comes the individual pulsar red noise\n",
    "# We model the red noise as a Fourier series with 30 frequency components, \n",
    "# with a power-law PSD\n",
    "s += blocks.red_noise_block(prior='log-uniform', Tspan=Tspan, components=30)\n",
    "\n",
    "# Finally, we add the common red noise, which is modeled as a Fourier series with 5 frequency components\n",
    "# The common red noise has a power-law PSD with spectral index of 4.33\n",
    "s += blocks.common_red_noise_block(psd='powerlaw', prior='log-uniform', Tspan=Tspan, \n",
    "                                   components=5, gamma_val=4.33, name='gw')\n",
    "\n",
    "# We set up the PTA object using the signal we defined above and the pulsars\n",
    "pta = signal_base.PTA([s(p) for p in psrs])\n",
    "\n",
    "# We need to set the white noise parameters to the values in the noise dictionary\n",
    "pta.set_default_params(noisedict)\n",
    "\n",
    "os = opt_stat.OptimalStatistic(psrs, pta=pta)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load up the maximum-likelihood values for the pulsars' red noise parameters and the common red process\n",
    "# These values come from the results of a Bayesian search (model 2A)\n",
    "# Once you have done your own Bayesian search, \n",
    "# you can make your own parameter dictionary of maximum-likelihood values\n",
    "\n",
    "with open('data/12p5yr_maxlike.json', 'r') as f:\n",
    "    ml_params = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute the optimal statistic\n",
    "# The optimal statistic returns five quantities:\n",
    "#  - xi: an array of the angular separations between the pulsar pairs (in radians)\n",
    "#  - rho: an array of the cross-correlations between the pulsar pairs\n",
    "#  - sig: an array of the uncertainty in the cross-correlations\n",
    "#  - OS: the value of the optimal statistic\n",
    "#  - OS_sig: the uncertainty in the optimal statistic\n",
    "\n",
    "xi, rho, sig, OS, OS_sig = os.compute_os(params=ml_params)\n",
    "print(OS, OS_sig, OS/OS_sig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the cross-correlations and compare to the Hellings-Downs curve\n",
    "# Before plotting, we need to bin the cross-correlations\n",
    "\n",
    "def weightedavg(rho, sig):\n",
    "    weights, avg = 0., 0.\n",
    "    for r,s in zip(rho,sig):\n",
    "        weights += 1./(s*s)\n",
    "        avg += r/(s*s)\n",
    "        \n",
    "    return avg/weights, np.sqrt(1./weights)\n",
    "\n",
    "def bin_crosscorr(zeta, xi, rho, sig):\n",
    "    \n",
    "    rho_avg, sig_avg = np.zeros(len(zeta)), np.zeros(len(zeta))\n",
    "    \n",
    "    for i,z in enumerate(zeta[:-1]):\n",
    "        myrhos, mysigs = [], []\n",
    "        for x,r,s in zip(xi,rho,sig):\n",
    "            if x >= z and x < (z+10.):\n",
    "                myrhos.append(r)\n",
    "                mysigs.append(s)\n",
    "        rho_avg[i], sig_avg[i] = weightedavg(myrhos, mysigs)\n",
    "        \n",
    "    return rho_avg, sig_avg\n",
    "\n",
    "# sort the cross-correlations by xi\n",
    "idx = np.argsort(xi)\n",
    "\n",
    "xi_sorted = xi[idx]\n",
    "rho_sorted = rho[idx]\n",
    "sig_sorted = sig[idx]\n",
    "\n",
    "# bin the cross-correlations so that there are the same number of pairs per bin\n",
    "npairs = 66\n",
    "\n",
    "xi_mean = []\n",
    "xi_err = []\n",
    "\n",
    "rho_avg = []\n",
    "sig_avg = []\n",
    "\n",
    "i = 0\n",
    "while i < len(xi_sorted):\n",
    "    \n",
    "    xi_mean.append(np.mean(xi_sorted[i:npairs+i]))\n",
    "    xi_err.append(np.std(xi_sorted[i:npairs+i]))\n",
    "\n",
    "    r, s = weightedavg(rho_sorted[i:npairs+i], sig_sorted[i:npairs+i])\n",
    "    rho_avg.append(r)\n",
    "    sig_avg.append(s)\n",
    "    \n",
    "    i += npairs\n",
    "    \n",
    "xi_mean = np.array(xi_mean)\n",
    "xi_err = np.array(xi_err)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_HD_curve(zeta):\n",
    "    \n",
    "    coszeta = np.cos(zeta*np.pi/180.)\n",
    "    xip = (1.-coszeta) / 2.\n",
    "    HD = 3.*( 1./3. + xip * ( np.log(xip) -1./6.) )\n",
    "    \n",
    "    return HD/2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# now make the plot\n",
    "\n",
    "(_, caps, _) = plt.errorbar(xi_mean*180/np.pi, rho_avg, xerr=xi_err*180/np.pi, yerr=sig_avg, marker='o', ls='', \n",
    "                            color='0.1', fmt='o', capsize=4, elinewidth=1.2)\n",
    "\n",
    "zeta = np.linspace(0.01,180,100)\n",
    "HD = get_HD_curve(zeta+1)\n",
    "\n",
    "plt.plot(zeta, OS*HD, ls='--', label='Hellings-Downs', color='C0', lw=1.5)\n",
    "\n",
    "plt.xlim(0, 180);\n",
    "#plt.ylim(-4e-30, 5e-30);\n",
    "plt.ylabel(r'$\\hat{A}^2 \\Gamma_{ab}(\\zeta)$')\n",
    "plt.xlabel(r'$\\zeta$ (deg)');\n",
    "\n",
    "plt.legend(loc=4);\n",
    "\n",
    "plt.tight_layout();\n",
    "plt.show();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To compute the noise-marginalized optimal statistic (Vigeland et al. 2018), you will need the chain from a Bayesian search for a common red process without spatial correlations (model 2A)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Change chaindir to point to where you have the chain from your Bayesian search\n",
    "chaindir = 'chains/model_2a/'\n",
    "params = list(np.loadtxt(chaindir + '/params.txt', dtype='str'))\n",
    "chain = np.loadtxt(chaindir + '/chain_1.txt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 1000   # number of times to compute the optimal statistic\n",
    "burn = int(0.25*chain.shape[0])   # estimate of when the chain has burned in\n",
    "\n",
    "noisemarg_OS, noisemarg_OS_err = np.zeros(N), np.zeros(N)\n",
    "\n",
    "for i in range(N):\n",
    "    \n",
    "    # choose a set of noise values from the chain\n",
    "    # make sure that you pull values from after the chain has burned in\n",
    "    idx = np.random.randint(burn, chain.shape[0])\n",
    "    \n",
    "    # construct a dictionary with these parameter values\n",
    "    param_dict = {}\n",
    "    for p in params:\n",
    "        param_dict.update({p: chain[idx, params.index(p)]})\n",
    "    \n",
    "    # compute the optimal statistic at this set of noise values and save in an array\n",
    "    _, _, _, noisemarg_OS[i], noisemarg_OS_err[i] = os.compute_os(params=param_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(noisemarg_OS)\n",
    "\n",
    "plt.figure();\n",
    "plt.hist(noisemarg_OS/noisemarg_OS_err)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}