ipta_gwa_tutorial.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 setup an analysis of indiviudual pulsar noise properties,\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": {
    "ExecuteTime": {
     "end_time": "2018-05-11T16:11:20.468597Z",
     "start_time": "2018-05-11T16:11:20.454565Z"
    }
   },
   "source": [
    "# Pre-requisites (installation etc.)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- **Install `miniconda` locally**\n",
    "\n",
    "    `wget -q https://repo.continuum.io/miniconda/Miniconda2-latest-Linux-x86_64.sh`\n",
    "    \n",
    "    `bash Miniconda2-latest-Linux-x86_64.sh -b -p ~/.local/opt/miniconda2`\n",
    "    \n",
    "    `rm Miniconda2-latest-Linux-x86_64.sh`\n",
    "\n",
    "\n",
    "- **Add miniconda’s `python`  to the front of your `$PATH`**\n",
    "\n",
    "    `echo \"export PATH=$HOME/.local/opt/miniconda2/bin:$PATH\" >> .bashrc`\n",
    "    \n",
    "    `source .bashrc`\n",
    "\n",
    "\n",
    "- **Install the basic python packages**\n",
    "\n",
    "    `conda install -y numpy==1.13.3 cython scipy`\n",
    "\n",
    "\n",
    "- **Install latest `libstempo` from GitHub with `pip`.  `tempo2` should be installed automatically.  Add extra ephemerides if needed**\n",
    "\n",
    "    `pip install git+https://github.com/vallis/libstempo@master`\n",
    "\n",
    "\n",
    "- **Install more python packages**\n",
    "\n",
    "    `conda install -y matplotlib ipython h5py mpi4py numexpr statsmodels astropy ephem`\n",
    "\n",
    "\n",
    "- **Install non-conda packages with `pip`**\n",
    "\n",
    "    `pip install healpy acor line_profiler jplephem corner numdifftools`\n",
    "\n",
    "\n",
    "- [optional] **Install `scikit-sparse`**\n",
    "\n",
    "    `conda install -c menpo scikit-sparse`\n",
    "   \n",
    "   \n",
    "- [optional] **Alternatively install `suite sparse` and then use pip to install `scikit-sparse` (maybe needed with python 3.6)**\n",
    "\n",
    "    `conda install -c conda-forge suitesparse`\n",
    "    \n",
    "    `pip install git+https://github.com/scikit-sparse/scikit-sparse.git@master`\n",
    "\n",
    "\n",
    "- **Install PTMCMC for sampling**\n",
    "\n",
    "    `pip install git+https://github.com/jellis18/PTMCMCSampler@master`\n",
    "\n",
    "\n",
    "- **Finally, install enterprise**\n",
    "\n",
    "    `pip install git+https://github.com/nanograv/enterprise@master`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T17:42:33.177433Z",
     "start_time": "2018-05-15T17:42:31.657125Z"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "% matplotlib inline\n",
    "%config InlineBackend.figure_format = 'retina'\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "from __future__ import division\n",
    "\n",
    "import numpy as np\n",
    "import os, glob, json \n",
    "import matplotlib.pyplot as plt\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": [
    "# Define a python dictionary of pulsar names and PTAs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T17:42:33.594100Z",
     "start_time": "2018-05-15T17:42:33.556901Z"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# The pulsars we'll be analyzing\n",
    "psrdict = {'J1713+0747': [{'pta': ['NANOGrav', 'PPTA']}], \n",
    "           'J1909-3744': [{'pta': ['NANOGrav', 'PPTA']}], \n",
    "           'J1640+2224': [{'pta': ['NANOGrav']}], \n",
    "           'J1600-3053': [{'pta': ['NANOGrav']}],\n",
    "           'J2317+1439': [{'pta': ['NANOGrav']}], \n",
    "           'J1918-0642': [{'pta': ['NANOGrav']}], \n",
    "           'J1614-2230': [{'pta': ['NANOGrav']}], \n",
    "           'J1744-1134': [{'pta': ['NANOGrav', 'PPTA']}],\n",
    "           'J0030+0451': [{'pta': ['NANOGrav']}], \n",
    "           'J2145-0750': [{'pta': ['NANOGrav']}], \n",
    "           'J1857+0943': [{'pta': ['NANOGrav']}], \n",
    "           'J1853+1303': [{'pta': ['NANOGrav']}], \n",
    "           'J0613-0200': [{'pta': ['NANOGrav']}],\n",
    "           'J1455-3330': [{'pta': ['NANOGrav']}], \n",
    "           'J1741+1351': [{'pta': ['NANOGrav']}], \n",
    "           'J2010-1323': [{'pta': ['NANOGrav']}], \n",
    "           'J1024-0719': [{'pta': ['NANOGrav']}], \n",
    "           'J1012+5307': [{'pta': ['NANOGrav']}],\n",
    "           'J0437-4715': [{'pta': ['PPTA']}]\n",
    "          }\n",
    "psrlist = psrdict.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T17:42:33.922845Z",
     "start_time": "2018-05-15T17:42:33.891909Z"
    }
   },
   "outputs": [],
   "source": [
    "psrlist"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get par, tim, and noise files\n",
    "Here we collect the tim and par files as well as noise files made from the `PAL2` code. These are the same par, tim, and noise files used in the 9-year analysis papers. We use the convienience function above to convert from `PAL2` noise files to `enterprise` parameter dictionaries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T17:42:39.433040Z",
     "start_time": "2018-05-15T17:42:39.407641Z"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "datadir = './partim_filtered_ppta_ng/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T17:42:40.533083Z",
     "start_time": "2018-05-15T17:42:40.500941Z"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "parfiles = sorted(glob.glob(datadir + '/*.par'))\n",
    "timfiles = sorted(glob.glob(datadir + '/*.tim'))\n",
    "\n",
    "# filter\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]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T17:42:40.937284Z",
     "start_time": "2018-05-15T17:42:40.907923Z"
    }
   },
   "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. `DE436`). 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(2015)`. 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": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T17:44:54.738924Z",
     "start_time": "2018-05-15T17:42:43.023316Z"
    }
   },
   "outputs": [],
   "source": [
    "psrs = []\n",
    "for p, t in zip(parfiles, timfiles):\n",
    "    psr = Pulsar(p, t, ephem='DE436', clk=None)\n",
    "    psrs.append(psr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Single pulsar analysis"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-11T19:00:23.677385Z",
     "start_time": "2018-05-11T19:00:23.570716Z"
    }
   },
   "source": [
    "* `enterprise` is structured so that one first creates `parameters`, then `signals` that these `parameters` belong to, then finally a `model` that is the union of all `signals` and the `data`.\n",
    "\n",
    "* We will show this explciitly below, then introduce some model shortcut code that will make your life easier.\n",
    "* We test on `J1713+0747`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T17:49:30.300840Z",
     "start_time": "2018-05-15T17:49:30.261826Z"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "psr = [p for p in psrs if p.name == 'J1713+0747'][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T17:49:30.750433Z",
     "start_time": "2018-05-15T17:49:30.721856Z"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# find the maximum time span to set red-noise/DM-variation frequency sampling\n",
    "tmin = psr.toas.min()\n",
    "tmax = psr.toas.max()\n",
    "Tspan = np.max(tmax) - np.min(tmin)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T17:49:31.267060Z",
     "start_time": "2018-05-15T17:49:31.239006Z"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# define selection by observing backend\n",
    "selection = selections.Selection(selections.by_backend)\n",
    "\n",
    "# special selection for ECORR only use wideband NANOGrav data\n",
    "selection2 = selections.Selection(selections.nanograv_backends)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T18:05:15.402712Z",
     "start_time": "2018-05-15T18:05:15.370248Z"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# white noise parameters\n",
    "#efac = parameter.Uniform(0.01, 10.0)\n",
    "#equad = parameter.Uniform(-8.5, -5)\n",
    "#ecorr = parameter.Uniform(-8.5, -5)\n",
    "efac = parameter.Constant() \n",
    "equad = parameter.Constant() \n",
    "ecorr = parameter.Constant() \n",
    "\n",
    "# red noise parameters\n",
    "log10_A = parameter.Uniform(-20, -11)\n",
    "gamma = parameter.Uniform(0, 7)\n",
    "\n",
    "# dm-variation parameters\n",
    "log10_A_dm = parameter.Uniform(-20, -11)\n",
    "gamma_dm = parameter.Uniform(0, 7)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T18:05:16.054660Z",
     "start_time": "2018-05-15T18:05:16.013299Z"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# [Optional] If fixing white-noise, read-in previously computed noisefiles\n",
    "noisefiles = sorted(glob.glob('./partim_filtered_ppta_ng/noisefiles_ppta_ng_normal/*.json'))\n",
    "\n",
    "params = {}\n",
    "for nf in noisefiles:\n",
    "    with open(nf, 'r') as fin:\n",
    "        params.update(json.load(fin))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create signals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T18:05:18.620288Z",
     "start_time": "2018-05-15T18:05:18.583614Z"
    },
    "collapsed": true
   },
   "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",
    "# DM-variations (powerlaw with 30 frequencies)\n",
    "dm_basis = utils.createfourierdesignmatrix_dm(nmodes=30, Tspan=Tspan)\n",
    "dm_pl = utils.powerlaw(log10_A=log10_A_dm, gamma=gamma_dm)\n",
    "dm_gp = gp_signals.BasisGP(dm_pl, dm_basis, name='dm_gp')\n",
    "\n",
    "# timing model\n",
    "tm = gp_signals.TimingModel(use_svd=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Piece the full model together"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T18:05:19.245655Z",
     "start_time": "2018-05-15T18:05:19.214962Z"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# full model\n",
    "s = ef + eq + rn + dm_gp + tm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T18:05:29.701318Z",
     "start_time": "2018-05-15T18:05:20.295739Z"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# intialize a single-pulsar pta model\n",
    "models = []\n",
    "   \n",
    "if 'NANOGrav' in psr.flags['pta']:\n",
    "    s2 = s + ec # ecorr only applied to NANOGrav data\n",
    "    models.append(s2(psr))\n",
    "else:\n",
    "    models.append(s(psr))\n",
    "    \n",
    "pta = signal_base.PTA(models)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T18:05:29.793302Z",
     "start_time": "2018-05-15T18:05:29.703707Z"
    }
   },
   "outputs": [],
   "source": [
    "# [Optional] Set white-noise parmeters from previous analysis\n",
    "pta.set_default_params(params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T18:05:29.821629Z",
     "start_time": "2018-05-15T18:05:29.795298Z"
    }
   },
   "outputs": [],
   "source": [
    "pta.params"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Draw initial sample from model parameter space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T18:05:48.625620Z",
     "start_time": "2018-05-15T18:05:48.595761Z"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "x0 = np.hstack(p.sample() for p in pta.params)\n",
    "ndim = len(x0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup sampler (simple, with no tricks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T18:05:49.721570Z",
     "start_time": "2018-05-15T18:05:49.690910Z"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# initial jump covariance matrix\n",
    "cov = np.diag(np.ones(ndim) * 0.01**2) # helps to tune MCMC proposal distribution\n",
    "\n",
    "# where chains will be written to\n",
    "outdir = './chains/ipta_dr2_ng_ppta_{}/'.format(str(psr.name))\n",
    "\n",
    "# sampler object\n",
    "sampler = ptmcmc(ndim, pta.get_lnlikelihood, pta.get_lnprior, cov,\n",
    "                 outDir=outdir, \n",
    "                 resume=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sample the parameter space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T18:07:34.298903Z",
     "start_time": "2018-05-15T18:05:51.594933Z"
    }
   },
   "outputs": [],
   "source": [
    "# sampler for N steps\n",
    "N = int(1e6)\n",
    "\n",
    "# SCAM = Single Component Adaptive Metropolis\n",
    "# AM = Adaptive Metropolis\n",
    "# DE = Differential Evolution\n",
    "## You can keep all these set at default values\n",
    "sampler.sample(x0, N, SCAMweight=30, AMweight=15, DEweight=50, )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simple post-processing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T21:14:29.205383Z",
     "start_time": "2018-05-15T21:14:28.983891Z"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "chain = np.loadtxt(outdir + 'chain_1.txt')\n",
    "burn = int(0.25 * chain.shape[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T21:14:29.448022Z",
     "start_time": "2018-05-15T21:14:29.400142Z"
    }
   },
   "outputs": [],
   "source": [
    "# Find column of chain file corresponding to a parameter\n",
    "ind = list(pta.param_names).index('J1713+0747_log10_A')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T18:08:16.630039Z",
     "start_time": "2018-05-15T18:08:16.283915Z"
    }
   },
   "outputs": [],
   "source": [
    "# Make trace-plot to diagnose sampling\n",
    "plt.plot(chain[burn:, ind])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T18:08:17.503339Z",
     "start_time": "2018-05-15T18:08:17.323209Z"
    }
   },
   "outputs": [],
   "source": [
    "# Plot a histogram of the marginalized posterior distribution\n",
    "plt.hist(chain[burn:,ind], 50, normed=True, histtype='stepfilled', \n",
    "         lw=2, color='C0', alpha=0.5);\n",
    "plt.xlabel('J1713+0747_log10_A')\n",
    "plt.ylabel('PDF')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T18:10:43.625371Z",
     "start_time": "2018-05-15T18:10:43.252841Z"
    }
   },
   "outputs": [],
   "source": [
    "# Make 2d histogram plot\n",
    "ind_redA = list(pta.param_names).index('J1713+0747_log10_A')\n",
    "ind_redgam = list(pta.param_names).index('J1713+0747_gamma')\n",
    "fig = corner.corner(chain[burn:, [ind_redA, ind_redgam]], \n",
    "                    labels=['J1713+0747_log10_A', 'J1713+0747_gamma'],\n",
    "                   levels=[0.68,0.95]);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-11T19:38:14.459181Z",
     "start_time": "2018-05-11T19:38:14.437071Z"
    }
   },
   "source": [
    "## Add a new custom function"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this example, we add a new custom function that takes the form of a dispersive ($\\nu^{-2}$) exponential dip in the residuals of `J1713+0747` due to a void in the ISM plasma. This is a real effect that has been observed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T18:10:52.540488Z",
     "start_time": "2018-05-15T18:10:52.509558Z"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "@signal_base.function\n",
    "def chrom_exp_decay(toas, freqs, log10_Amp=-7,\n",
    "                    t0=54000, log10_tau=1.7, idx=2):\n",
    "    \"\"\"\n",
    "    Chromatic exponential-dip delay term in TOAs.\n",
    "\n",
    "    :param t0: time of exponential minimum [MJD]\n",
    "    :param tau: 1/e time of exponential [s]\n",
    "    :param log10_Amp: amplitude of dip\n",
    "    :param idx: index of chromatic dependence\n",
    "\n",
    "    :return wf: delay time-series [s]\n",
    "    \"\"\"\n",
    "    t0 *= const.day\n",
    "    tau = 10**log10_tau * const.day\n",
    "    wf = -10**log10_Amp * np.heaviside(toas - t0, 1) * \\\n",
    "        np.exp(- (toas - t0) / tau)\n",
    "\n",
    "    return wf * (1400 / freqs) ** idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T18:10:53.048214Z",
     "start_time": "2018-05-15T18:10:53.016148Z"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def dm_exponential_dip(tmin, tmax, idx=2, name='dmexp'):\n",
    "    \"\"\"\n",
    "    Returns chromatic exponential dip (i.e. TOA advance):\n",
    "\n",
    "    :param tmin, tmax:\n",
    "        search window for exponential dip time.\n",
    "    :param idx:\n",
    "        index of radio frequency dependence (i.e. DM is 2). If this is set\n",
    "        to 'vary' then the index will vary from 1 - 6\n",
    "    :param name: Name of signal\n",
    "\n",
    "    :return dmexp:\n",
    "        chromatic exponential dip waveform.\n",
    "    \"\"\"\n",
    "    t0_dmexp = parameter.Uniform(tmin,tmax)\n",
    "    log10_Amp_dmexp = parameter.Uniform(-10, -2)\n",
    "    log10_tau_dmexp = parameter.Uniform(np.log10(5), np.log10(100))\n",
    "    wf = chrom_exp_decay(log10_Amp=log10_Amp_dmexp, t0=t0_dmexp,\n",
    "                         log10_tau=log10_tau_dmexp, idx=idx)\n",
    "    dmexp = deterministic_signals.Deterministic(wf, name=name)\n",
    "\n",
    "    return dmexp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T18:11:13.001040Z",
     "start_time": "2018-05-15T18:11:03.015658Z"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# full model\n",
    "s = ef + eq + rn + dm_gp + tm\n",
    "\n",
    "# intialize a single-pulsar pta model\n",
    "models = []\n",
    "    \n",
    "if psr.name == 'J1713+0747' and 'NANOGrav' in psr.flags['pta']:\n",
    "    s_prime = s + dm_exponential_dip(tmin=(tmin/86400.0), tmax=(tmax/86400.0)) + ec\n",
    "    models.append(s_prime(psr))\n",
    "elif psr.name == 'J1713+0747':\n",
    "    s_prime = s + dm_exponential_dip(tmin=(tmin/86400.0), tmax=(tmax/86400.0))\n",
    "    models.append(s_prime(psr))\n",
    "elif 'NANOGrav' in psr.flags['pta']:\n",
    "    s_prime = s + ec\n",
    "    models.append(s_prime(psr))\n",
    "else:\n",
    "    models.append(s(psr))\n",
    "    \n",
    "pta = signal_base.PTA(models)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T18:11:28.201694Z",
     "start_time": "2018-05-15T18:11:28.116873Z"
    }
   },
   "outputs": [],
   "source": [
    "# [Optional] Set white-noise parmeters from previous analysis\n",
    "pta.set_default_params(params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T18:11:30.699073Z",
     "start_time": "2018-05-15T18:11:30.667735Z"
    }
   },
   "outputs": [],
   "source": [
    "pta.params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T18:22:17.885688Z",
     "start_time": "2018-05-15T18:13:14.957548Z"
    }
   },
   "outputs": [],
   "source": [
    "# Sample as before\n",
    "x0 = np.hstack(p.sample() for p in pta.params)\n",
    "ndim = len(x0)\n",
    "\n",
    "# initial jump covariance matrix\n",
    "cov = np.diag(np.ones(ndim) * 0.01**2) # helps to tune MCMC proposal distribution\n",
    "\n",
    "# where chains will be written to\n",
    "outdir = './chains/ipta_dr2_ng_ppta_{}_dmexp/'.format(str(psr.name))\n",
    "\n",
    "# sampler object\n",
    "sampler = ptmcmc(ndim, pta.get_lnlikelihood, pta.get_lnprior, cov,\n",
    "                 outDir=outdir, \n",
    "                 resume=False)\n",
    "\n",
    "# sampler for N steps\n",
    "N = int(1e6)\n",
    "\n",
    "# SCAM = Single Component Adaptive Metropolis\n",
    "# AM = Adaptive Metropolis\n",
    "# DE = Differential Evolution\n",
    "## You can keep all these set at default values\n",
    "sampler.sample(x0, N, SCAMweight=30, AMweight=15, DEweight=50, )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T18:59:39.789629Z",
     "start_time": "2018-05-15T18:59:39.689206Z"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "chain = np.loadtxt(outdir + 'chain_1.txt')\n",
    "burn = int(0.5 * chain.shape[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T19:00:10.404441Z",
     "start_time": "2018-05-15T19:00:10.199024Z"
    }
   },
   "outputs": [],
   "source": [
    "# Plot a histogram of the marginalized posterior distribution\n",
    "\n",
    "ind = list(pta.param_names).index('J1713+0747_dmexp_log10_tau')\n",
    "\n",
    "plt.hist(chain[burn:,ind], 50, normed=True, histtype='stepfilled', \n",
    "         lw=2, color='C0', alpha=0.5);\n",
    "plt.xlabel('J1713+0747_dmexp_t0')\n",
    "plt.ylabel('PDF')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Now, the easy way to do all of this\n",
    "\n",
    "Many of us have created shortcuts to carry out these tasks. You will find them in `enterprise_extensions`: https://github.com/stevertaylor/enterprise_extensions. Clone this repo, go into the cloned repo directory, then execute `pip install .`\n",
    "\n",
    "This will install the package on your local machine."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T21:14:51.711971Z",
     "start_time": "2018-05-15T21:14:51.659585Z"
    },
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import enterprise_extensions\n",
    "from enterprise_extensions import models, model_utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T21:15:02.109579Z",
     "start_time": "2018-05-15T21:14:51.916806Z"
    }
   },
   "outputs": [],
   "source": [
    "# Create a single pulsar model\n",
    "pta = models.model_singlepsr_noise(psr, psd='powerlaw', \n",
    "                                   noisedict=params, \n",
    "                                   white_vary=False,\n",
    "                                   dm_var=True, \n",
    "                                   dm_psd='powerlaw', dm_annual=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T21:15:20.657591Z",
     "start_time": "2018-05-15T21:15:20.614041Z"
    }
   },
   "outputs": [],
   "source": [
    "# Setup a sampler instance.\n",
    "# This will add some fanicer stuff than before, like prior draws, \n",
    "# and custom sample groupings.\n",
    "sampler = model_utils.setup_sampler(pta, outdir=outdir, resume=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T21:22:08.883008Z",
     "start_time": "2018-05-15T21:15:23.402604Z"
    }
   },
   "outputs": [],
   "source": [
    "# sampler for N steps\n",
    "N = int(1e6)\n",
    "x0 = x0 = np.hstack(p.sample() for p in pta.params)\n",
    "\n",
    "# SCAM = Single Component Adaptive Metropolis\n",
    "# AM = Adaptive Metropolis\n",
    "# DE = Differential Evolution\n",
    "## You can keep all these set at default values\n",
    "sampler.sample(x0, N, SCAMweight=30, AMweight=15, DEweight=50, )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T21:22:12.699593Z",
     "start_time": "2018-05-15T21:22:12.447631Z"
    },
    "collapsed": true
   },
   "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-15T21:22:13.789123Z",
     "start_time": "2018-05-15T21:22:12.701917Z"
    }
   },
   "outputs": [],
   "source": [
    "pp.plot_trace()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Single-pulsar Model Selection\n",
    "\n",
    "We want to be able to create custom noise descriptions per pulsar. Some pulsars will have funky noise features (like the dispersive dip in `J1713+0747`) while others will have more vanilla noise models. The sure-fire way to exlore this is through model selection. We want to incorporate how well a model describes the data, and the Occam penalty given through too many unneccessary parameters. A simple way to do this is given below."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create list of pta models for our model selection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T21:22:54.018955Z",
     "start_time": "2018-05-15T21:22:33.317872Z"
    }
   },
   "outputs": [],
   "source": [
    "nmodels = 2\n",
    "mod_index = np.arange(nmodels)\n",
    "\n",
    "# Make dictionary of single-pulsar PTAs.\n",
    "# Select between a power-law or t-process red-noise spectrum\n",
    "pta = dict.fromkeys(mod_index)\n",
    "pta[0] = models.model_singlepsr_noise(psr, psd='powerlaw', noisedict=params, \n",
    "                                      white_vary=False,\n",
    "                                      dm_var=True, dm_psd='powerlaw', dm_annual=True)\n",
    "pta[1] = models.model_singlepsr_noise(psr, psd='tprocess', noisedict=params, \n",
    "                                      white_vary=False,\n",
    "                                      dm_var=True, dm_psd='powerlaw', dm_annual=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Instanciate a collection of models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T21:23:37.525841Z",
     "start_time": "2018-05-15T21:23:37.492145Z"
    },
    "collapsed": true
   },
   "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:23:38.800138Z",
     "start_time": "2018-05-15T21:23:38.761623Z"
    }
   },
   "outputs": [],
   "source": [
    "sampler = super_model.setup_sampler(resume=False, outdir=outdir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T21:23:41.025568Z",
     "start_time": "2018-05-15T21:23:40.986219Z"
    },
    "collapsed": true
   },
   "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:32:46.451287Z",
     "start_time": "2018-05-15T21:23:42.741331Z"
    }
   },
   "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-15T21:33:38.302769Z",
     "start_time": "2018-05-15T21:33:37.905553Z"
    },
    "collapsed": true
   },
   "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-15T21:34:22.340106Z",
     "start_time": "2018-05-15T21:34:22.139093Z"
    }
   },
   "outputs": [],
   "source": [
    "# Histogram of model-indexing variable\n",
    "plt.hist(chain[burn:,-5],);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Baysian Odds-ratio"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The odds-ratio is given by the relative number of sampling iterations spent in each sub-likelihood."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2018-05-15T21:34:23.417916Z",
     "start_time": "2018-05-15T21:34:23.382652Z"
    }
   },
   "outputs": [],
   "source": [
    "print model_utils.odds_ratio(chain[burn:,-5], models=[0,1])"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.13"
  },
  "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
}