gwa_tutorial_part2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using enterprise to analyze PTA data\n",
    "\n",
    "In this notebook you will learn:\n",
    "* How to use `enterprise` to interact with IPTA data,\n",
    "* How to search in PTA data for GWs,\n",
    "* How to perform Bayesian model selection,\n",
    "* How to post-process your results."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T21:48:39.150382Z",
     "start_time": "2018-05-15T21:48:39.094542Z"
    }
   },
   "outputs": [],
   "source": [
    "from __future__ import division\n",
    "\n",
    "%matplotlib inline\n",
    "%config InlineBackend.figure_format = 'retina'\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import os, glob, json, pickle\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import scipy.linalg as sl\n",
    "\n",
    "import enterprise\n",
    "from enterprise.pulsar import Pulsar\n",
    "import enterprise.signals.parameter as parameter\n",
    "from enterprise.signals import utils\n",
    "from enterprise.signals import signal_base\n",
    "from enterprise.signals import selections\n",
    "from enterprise.signals.selections import Selection\n",
    "from enterprise.signals import white_signals\n",
    "from enterprise.signals import gp_signals\n",
    "from enterprise.signals import deterministic_signals\n",
    "import enterprise.constants as const\n",
    "\n",
    "import corner\n",
    "from PTMCMCSampler.PTMCMCSampler import PTSampler as ptmcmc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get par, tim, and noise files (this is not the preferred method when we have supplied pickled enterprise Pulsar files; see below)\n",
    "Here we collect the tim and par files. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T21:48:42.658570Z",
     "start_time": "2018-05-15T21:48:42.631742Z"
    }
   },
   "outputs": [],
   "source": [
    "psrlist = None # define a list of pulsar name strings that can be used to filter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "datadir = '/Users/taylosr8/Research/NANOGrav/NANOGrav_12y/' # set your data directory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T21:48:43.043752Z",
     "start_time": "2018-05-15T21:48:43.007076Z"
    }
   },
   "outputs": [],
   "source": [
    "parfiles = sorted(glob.glob(datadir + 'par/*.par'))\n",
    "timfiles = sorted(glob.glob(datadir + 'tim/*.tim'))\n",
    "\n",
    "# filter\n",
    "if psrlist is not None:\n",
    "    parfiles = [x for x in parfiles if x.split('/')[-1].split('.')[0] in psrlist]\n",
    "    timfiles = [x for x in timfiles if x.split('/')[-1].split('.')[0] in psrlist]   \n",
    "    \n",
    "# Make sure you use the tempo2 parfile for J1713+0747!!\n",
    "# ...filtering out the tempo parfile... \n",
    "parfiles = [x for x in parfiles if 'J1713+0747_NANOGrav_12yv2.gls.par' not in x]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T21:48:43.452613Z",
     "start_time": "2018-05-15T21:48:43.421953Z"
    }
   },
   "outputs": [],
   "source": [
    "len(parfiles)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load into Pulsar class list"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* The `enterprise` Pulsar class uses `libstempo` to read in `par` and `tim` files, then stores all pulsar data into a `Pulsar` object. This object contains all data and meta-data needed for the ensuing pulsar and PTA analysis. You no longer to reference the `par` and `tim` files after this cell.\n",
    "* Note below that you can explicitly declare which version of the JPL solar-system ephemeris model that will be used to compute the Roemer delay between the geocenter and the barycenter (e.g. `DE438`). Otherwise the default values will be taken from the `par` files. Explicitly declaring the version here is good practice.\n",
    "* You can also explicitly set the clock file to a version of `BIPM`, e.g. `BIPM(2018)`. This is less important, and you can let the code take the value from the `par` file.\n",
    "* When you execute the following cell, you will get warnings like `WARNING: Could not find pulsar distance for PSR ...`. Don't worry! This is expected, and fine. Not all pulsars have well constrained distances, and will be set to `1 kpc` with a `20%` uncertainty."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Read par and tim files into enterprise Pulsar objects"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T21:49:00.790991Z",
     "start_time": "2018-05-15T21:48:53.703630Z"
    }
   },
   "outputs": [],
   "source": [
    "psrs = []\n",
    "for p, t in zip(parfiles, timfiles):\n",
    "    psr = Pulsar(p, t, ephem='DE438', clk='BIPM(2018)')\n",
    "    psrs.append(psr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## OR... load in enterprise pickled Pulsar instances that we've prepared!\n",
    "Go here for full details: https://paper.dropbox.com/doc/NG-12.5yr_v3-GWB-Analysis--A2zJbxQU704Oq9jU1oqAWJCHAQ-DICJei6NxsPjxnO90mGMo\n",
    "\n",
    "Pickled 12.5yr pulsars: https://drive.google.com/file/d/1GUcmdj9OMf7-hrAOydeHEO4ylQxgC9Kb/edit\n",
    "\n",
    "Noise files: https://drive.google.com/file/d/1V7bu2y5hxFSj_7KWO3uNM7Q_0dciXfSX/edit\n",
    "\n",
    "Empirical red noise proposal distributions: https://drive.google.com/file/d/19odsqZ93Wh8og1AGdU7SfvXzB0DJwipG/edit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## set your data directory\n",
    "datadir = '/Users/taylosr8/Downloads/' # set your data directory\n",
    "\n",
    "## read in pickles\n",
    "psrs = pickle.load(open(datadir + 'channelized_12yr_v3_partim_DE438.pkl', 'rb'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* We can read-in some previously computed noise properties from single-pulsar analyses. These are things like `EFAC`, `EQUAD`, and (for `NANOGrav`) `ECORR`. \n",
    "* In practice, we set these white-noise properties as fixed in the low-frequency noise / GW searches.\n",
    "* The noise properties have been stored as `json` files, and are read in to a big parameter dictionary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Get parameter noise dictionary\n",
    "noise_ng12 = datadir + 'channelized_12p5yr_v3_full_noisedict.json'\n",
    "\n",
    "params = {}\n",
    "with open(noise_ng12, 'r') as fp:\n",
    "    params.update(json.load(fp))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Load in empirical distributions\n",
    "#emp_dists = pickle.load(open(datadir + '12yr_emp_dist_RNonly_py3.pkl', 'rb'))\n",
    "#emp_dists = '/home/stephen.taylor/NANOGrav/nanograv_12p5yr_analysis/nanograv_12p5yr_analysis_mar2020/data/12yr_emp_dist_RNonly_py3.pkl'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PTA Parameter Estimation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set up `enterprise` model for PTA upper-limit (*verbose version*)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* Usually, in a full PTA analysis we fix all of the white noise (EFAC, EQUAD, and ECORR) parameters to the values obtained from the noise files. This is done by using `Constant` parameters. In this case we do not specify a default value for all instances of that parameter but instead will set them, based on their initialized pulsar and backend specific name, later via the `set_default_params` method of `PTA`. \n",
    "\n",
    "* We use the `Selection` object to define which noise parameters are assigned to which chunks of TOAs. This selection is based on unique combination of backends and receivers."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-11T16:39:50.133639Z",
     "start_time": "2018-05-11T16:39:50.102972Z"
    }
   },
   "source": [
    "* Another feature to notice is that **for upper limits** we do not use a uniform prior on the log of the red-noise or GWB amplitude. Instead we use a `LinearExp` prior (short for linear-exponent prior), that is a prior of the form $p(x)\\propto 10^x$. This is how we can still use the log of the parameter to sample but place a uniform prior on the parameter itself. We do this for both the red noise and GWB amplitude parameters. **For detection analyses** we still use a `Uniform` prior on the log of the red-noise and GWB amplitude. \n",
    "\n",
    "* In order to save on computing time we do not include spatial correlations here. Instead we model the GWB as a common red process across all pulsars. In `enterprise` we can do this with a simple trick. We pre-initialize the parameters before passing them to the `Signal` model. In this way the *same* parameter instance is used for all pulsars. Lastly, we fixt the spectral index of the GWB to be 13/3 (4.33) using the `Constant` parameter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T21:49:45.238037Z",
     "start_time": "2018-05-15T21:49:45.209705Z"
    }
   },
   "outputs": [],
   "source": [
    "# find the maximum time span to set GW frequency sampling\n",
    "tmin = [p.toas.min() for p in psrs]\n",
    "tmax = [p.toas.max() for p in psrs]\n",
    "Tspan = np.max(tmax) - np.min(tmin)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T21:49:48.496954Z",
     "start_time": "2018-05-15T21:49:48.470085Z"
    }
   },
   "outputs": [],
   "source": [
    "# define selection by observing backend\n",
    "selection = selections.Selection(selections.by_backend)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T21:49:50.042751Z",
     "start_time": "2018-05-15T21:49:50.005677Z"
    }
   },
   "outputs": [],
   "source": [
    "# white noise parameters\n",
    "# since we are fixing these to values from the noise file we set\n",
    "# them as constant parameters\n",
    "efac = parameter.Constant()\n",
    "equad = parameter.Constant()\n",
    "ecorr = parameter.Constant()\n",
    "\n",
    "# red noise parameters\n",
    "log10_A = parameter.LinearExp(-20, -11)\n",
    "gamma = parameter.Uniform(0, 7)\n",
    "\n",
    "# GW parameters (initialize with names here to use parameters in common across pulsars)\n",
    "log10_A_gw = parameter.LinearExp(-18,-12)('log10_A_gw')\n",
    "gamma_gw = parameter.Constant(4.33)('gamma_gw')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-11T16:44:44.713491Z",
     "start_time": "2018-05-11T16:44:44.685587Z"
    }
   },
   "source": [
    "### Signals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T21:49:51.557605Z",
     "start_time": "2018-05-15T21:49:51.499006Z"
    }
   },
   "outputs": [],
   "source": [
    "# white noise\n",
    "ef = white_signals.MeasurementNoise(efac=efac, selection=selection)\n",
    "eq = white_signals.EquadNoise(log10_equad=equad, selection=selection)\n",
    "ec = white_signals.EcorrKernelNoise(log10_ecorr=ecorr, selection=selection2)\n",
    "\n",
    "# red noise (powerlaw with 30 frequencies)\n",
    "pl = utils.powerlaw(log10_A=log10_A, gamma=gamma)\n",
    "rn = gp_signals.FourierBasisGP(spectrum=pl, components=30, Tspan=Tspan)\n",
    "\n",
    "# gwb (no spatial correlations)\n",
    "cpl = utils.powerlaw(log10_A=log10_A_gw, gamma=gamma_gw)\n",
    "gw = gp_signals.FourierBasisGP(spectrum=cpl, components=30, Tspan=Tspan, name='gw')\n",
    "\n",
    "# for spatial correltions you can do...\n",
    "#orf = utils.hd_orf()\n",
    "#crn = gp_signals.FourierBasisCommonGP(cpl, orf, components=30, Tspan=Tspan, name='gw')\n",
    "\n",
    "# to add solar system ephemeris modeling...\n",
    "bayesephem=False\n",
    "if bayesephem:\n",
    "    eph = deterministic_signals.PhysicalEphemerisSignal(use_epoch_toas=True)\n",
    "\n",
    "# timing model\n",
    "tm = gp_signals.TimingModel(use_svd=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T21:49:52.662732Z",
     "start_time": "2018-05-15T21:49:52.638142Z"
    }
   },
   "outputs": [],
   "source": [
    "# full model\n",
    "if bayesephem:\n",
    "    s = ef + eq + rn + tm + eph + gw\n",
    "else:\n",
    "    s = ef + eq + rn + tm + gw"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T21:49:55.063952Z",
     "start_time": "2018-05-15T21:49:53.602147Z"
    }
   },
   "outputs": [],
   "source": [
    "# intialize PTA\n",
    "models = []\n",
    "        \n",
    "for p in psrs:    \n",
    "    models.append(s(p))\n",
    "    \n",
    "pta = signal_base.PTA(models)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T21:49:56.940424Z",
     "start_time": "2018-05-15T21:49:56.910938Z"
    }
   },
   "outputs": [],
   "source": [
    "pta.params"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set white noise parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T21:50:00.319108Z",
     "start_time": "2018-05-15T21:50:00.244934Z"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "pta.set_default_params(params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set initial parameters drawn from prior"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T21:50:03.822379Z",
     "start_time": "2018-05-15T21:50:03.788843Z"
    }
   },
   "outputs": [],
   "source": [
    "x0 = np.hstack([p.sample() for p in pta.params])\n",
    "ndim = len(x0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set up sampler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T21:50:06.902347Z",
     "start_time": "2018-05-15T21:50:06.866018Z"
    }
   },
   "outputs": [],
   "source": [
    "# initial jump covariance matrix\n",
    "cov = np.diag(np.ones(ndim) * 0.01**2)\n",
    "\n",
    "sampler = ptmcmc(ndim, pta.get_lnlikelihood, pta.get_lnprior, cov, \n",
    "                 outDir='./chains_pta_test/', resume=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sample!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T21:51:39.157421Z",
     "start_time": "2018-05-15T21:50:07.962002Z"
    }
   },
   "outputs": [],
   "source": [
    "# sampler for N steps\n",
    "N = int(5e6)\n",
    "x0 = np.hstack(p.sample() for p in pta.params)\n",
    "sampler.sample(x0, N, SCAMweight=30, AMweight=15, DEweight=50, )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T21:51:44.119807Z",
     "start_time": "2018-05-15T21:51:44.060350Z"
    }
   },
   "outputs": [],
   "source": [
    "chain = np.loadtxt('./chains_pta_test/chain_1.txt')\n",
    "burn = int(0.25 * chain.shape[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T21:51:46.102261Z",
     "start_time": "2018-05-15T21:51:45.927388Z"
    }
   },
   "outputs": [],
   "source": [
    "plt.hist(chain[burn:,-5], 50, normed=True, histtype='step', lw=2);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "### Upper limit value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T21:51:49.718217Z",
     "start_time": "2018-05-15T21:51:49.682886Z"
    }
   },
   "outputs": [],
   "source": [
    "upper = 10**np.percentile(chain[burn:, -5], q=95)\n",
    "print(upper)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Now, the easy way to do all of this...\n",
    "\n",
    "We use `enterprise_extensions` as in the single-pulsar analysis tutorial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T21:53:11.676320Z",
     "start_time": "2018-05-15T21:53:11.629096Z"
    }
   },
   "outputs": [],
   "source": [
    "import enterprise_extensions\n",
    "from enterprise_extensions import models, model_utils, hypermodel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T21:56:18.738032Z",
     "start_time": "2018-05-15T21:56:17.309346Z"
    }
   },
   "outputs": [],
   "source": [
    "pta = models.model_general(psrs, red_psd='powerlaw', common_psd='powerlaw', \n",
    "                           red_components=30, common_components=30, \n",
    "                           white_vary=False, noisedict=params, tm_svd=True, \n",
    "                           gamma_common=4.33, upper_limit_red=True, \n",
    "                           upper_limit_common=True, bayesephem=True)\n",
    "# note that bayesephem is switched on here, but you can also switch it off."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T21:56:20.663093Z",
     "start_time": "2018-05-15T21:56:20.632086Z"
    }
   },
   "outputs": [],
   "source": [
    "# Setup an instance of a HyperModel.\n",
    "# This doesn't mean we are doing model selection (yet!), but the \n",
    "# hypermodel module gives access to some nifty sampling schemes.\n",
    "super_model = hypermodel.HyperModel({0: pta})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T21:56:20.823333Z",
     "start_time": "2018-05-15T21:56:20.790108Z"
    }
   },
   "outputs": [],
   "source": [
    "outdir = './chains_pta_test/'\n",
    "sampler = super_model.setup_sampler(resume=False, outdir=outdir, sample_nmodel=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T21:56:21.506309Z",
     "start_time": "2018-05-15T21:56:21.474529Z"
    }
   },
   "outputs": [],
   "source": [
    "# sampler for N steps\n",
    "N = int(5e6)\n",
    "x0 = super_model.initial_sample()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T21:58:54.315884Z",
     "start_time": "2018-05-15T21:56:25.820406Z"
    }
   },
   "outputs": [],
   "source": [
    "# sample\n",
    "sampler.sample(x0, N, SCAMweight=30, AMweight=15, DEweight=50, )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T21:58:58.146382Z",
     "start_time": "2018-05-15T21:58:58.013143Z"
    }
   },
   "outputs": [],
   "source": [
    "# Read in chains and parameters\n",
    "\n",
    "chain = np.loadtxt(outdir + '/chain_1.txt')\n",
    "burn = int(0.25*chain.shape[0])\n",
    "pars = np.loadtxt(outdir + '/pars.txt', dtype=np.unicode_)\n",
    "\n",
    "pp = model_utils.PostProcessing(chain, pars)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T21:58:58.777011Z",
     "start_time": "2018-05-15T21:58:58.555011Z"
    }
   },
   "outputs": [],
   "source": [
    "# Plot GW amplitude posterior\n",
    "ind = list(pars).index('log10_A_gw')\n",
    "plt.hist(chain[burn:,ind], bins=40);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T21:59:00.397744Z",
     "start_time": "2018-05-15T21:59:00.366793Z"
    }
   },
   "outputs": [],
   "source": [
    "# Compute upper limit\n",
    "print model_utils.ul(chain[burn:, ind], q=95.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PTA Model Selection\n",
    "\n",
    "We want to be able to compute the Bayesian odds for a GWB in the data. This can be done using the hypermodel class, where we choose between a model with a common (but uncorrelated) red process in the pulsars, and a GWB affecting all pulsars.\n",
    "\n",
    "We typically perform detection-type analyses with uniform-in-log priors on all amplitude parameters for low-frequency processes. This is implemented below whenever we switch `upper_limit` to be equal to `False`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup dictionary of PTA models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T21:59:25.363204Z",
     "start_time": "2018-05-15T21:59:22.414521Z"
    }
   },
   "outputs": [],
   "source": [
    "nmodels = 2\n",
    "mod_index = np.arange(nmodels)\n",
    "\n",
    "# Make dictionary of PTAs.\n",
    "pta = dict.fromkeys(mod_index)\n",
    "pta[0] = models.model_general(psrs, red_psd='powerlaw', common_psd='powerlaw', \n",
    "                           red_components=30, common_components=30, \n",
    "                           white_vary=False, noisedict=params, tm_svd=True, \n",
    "                           gamma_common=4.33, upper_limit_red=False, \n",
    "                           upper_limit_common=False, bayesephem=True, orf=None)\n",
    "pta[1] = models.model_general(psrs, red_psd='powerlaw', common_psd='powerlaw', \n",
    "                           red_components=30, common_components=30, \n",
    "                           white_vary=False, noisedict=params, tm_svd=True, \n",
    "                           gamma_common=4.33, upper_limit_red=False, \n",
    "                           upper_limit_common=False, bayesephem=True, orf='hd')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T21:59:38.043468Z",
     "start_time": "2018-05-15T21:59:38.006027Z"
    }
   },
   "outputs": [],
   "source": [
    "super_model = model_utils.HyperModel(pta)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T21:59:40.020837Z",
     "start_time": "2018-05-15T21:59:39.985755Z"
    }
   },
   "outputs": [],
   "source": [
    "sampler = super_model.setup_sampler(resume=False, outdir=outdir, sample_nmodel=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T21:59:40.335389Z",
     "start_time": "2018-05-15T21:59:40.301829Z"
    }
   },
   "outputs": [],
   "source": [
    "# sampler for N steps\n",
    "N = int(5e6)\n",
    "x0 = super_model.initial_sample()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T22:07:38.858293Z",
     "start_time": "2018-05-15T21:59:41.689205Z"
    }
   },
   "outputs": [],
   "source": [
    "# sample\n",
    "sampler.sample(x0, N, SCAMweight=30, AMweight=15, DEweight=50, )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Post-process"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T22:07:42.274771Z",
     "start_time": "2018-05-15T22:07:41.979583Z"
    }
   },
   "outputs": [],
   "source": [
    "chain = np.loadtxt(outdir + '/chain_1.txt')\n",
    "burn = int(0.25*chain.shape[0])\n",
    "pars = np.loadtxt(outdir + '/pars.txt', dtype=np.unicode_)\n",
    "\n",
    "pp = model_utils.PostProcessing(chain, pars)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T22:07:44.047020Z",
     "start_time": "2018-05-15T22:07:43.807128Z"
    }
   },
   "outputs": [],
   "source": [
    "# Plot histgram for GW amplitude\n",
    "chain_burn = chain[burn:,:]\n",
    "\n",
    "ind_model = list(pars).index('nmodel')\n",
    "ind_gwamp = list(pars).index('log10_A_gw')\n",
    "\n",
    "# ORF = None\n",
    "#plt.hist(chain_burn[chain_burn[:, ind_model] < 0.5, ind_gwamp], bins=40);\n",
    "\n",
    "# ORF = Hellings & Downs\n",
    "plt.hist(chain_burn[chain_burn[:, ind_model] > 0.5, ind_gwamp], bins=40);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T22:07:46.481391Z",
     "start_time": "2018-05-15T22:07:46.244859Z"
    }
   },
   "outputs": [],
   "source": [
    "# Plot histogram for GWB model selection\n",
    "plt.hist(chain_burn[:, ind_model], bins=40);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bayes factors"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Savage-Dickey Bayes factor\n",
    "\n",
    "This gives the signal-vs-noise Bayes factor for a common red process in the pulsars plus intrisnc noise, versus intrinsic noise alone."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T22:08:38.949894Z",
     "start_time": "2018-05-15T22:08:38.920953Z"
    }
   },
   "outputs": [],
   "source": [
    "print(model_utils.bayes_fac(chain_burn[chain_burn[:, ind_model] < 0.5, ind_gwamp], ntol=1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Posterior odds ratio\n",
    "\n",
    "This gives the Bayesian odds between a model with a Hellings & Downs correlated red process between pulsars, and a common (but uncorrelated) red process between pulsars. This is the smoking-gun detection statsitic for a GWB signal."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T22:08:40.590729Z",
     "start_time": "2018-05-15T22:08:40.559995Z"
    }
   },
   "outputs": [],
   "source": [
    "print(model_utils.odds_ratio(chain_burn[:, ind_model], models=[0,1]))"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "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"
  },
  "toc": {
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "toc_cell": false,
   "toc_position": {
    "height": "743px",
    "left": "0px",
    "right": "1458px",
    "top": "107px",
    "width": "212px"
   },
   "toc_section_display": "block",
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}