gen_FoM_dataset.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generate Figure of Merit based dataset\n",
    "\n",
    "This notebook is to generate a set of DR2-lite style `.par` and `.tim` files based on a figure of merit (FoM).\n",
    "We include functions for a burst with memory (BWM) and stochastic gravitational wave background (gwb) FoM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from __future__ import division, print_function, unicode_literals\n",
    "\n",
    "import numpy as np\n",
    "import os, glob\n",
    "import json\n",
    "import matplotlib.pyplot as plt\n",
    "import libstempo as t2\n",
    "from enterprise.signals import utils\n",
    "\n",
    "import dr2lite_utils as dr2u\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Update paths to reflect your usage."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "DR2dir = os.path.abspath('/Users/ptb/Projects/pulsar/data/DR2')\n",
    "datadir = os.path.join(DR2dir, 'release/VersionB')\n",
    "\n",
    "lite_dir = os.path.join(DR2dir, 'DR2lite_partim')  # cleaned/combined par/tim files\n",
    "os.system('mkdir -p {}'.format(lite_dir));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1) clean up `.par` and `.tim` files"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## make list of pulsar names for later"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "psrlist = []\n",
    "psrs = glob.glob(datadir + '/J*')\n",
    "for psr in psrs:\n",
    "    name = psr.split('/')[-1]\n",
    "    psrlist.append(name)\n",
    "psrlist.sort()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## clean `.par` files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "parfiles = glob.glob(datadir + '/J*/*IPTADR2.par')\n",
    "\n",
    "for p in parfiles:\n",
    "    name = p.split('/')[-2]\n",
    "    outfile = os.path.join(lite_dir, '{}.par'.format(name))\n",
    "    dr2u.clean_par(p, outfile)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## combine `.tim` files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "timfiles = glob.glob(datadir + '/J*/*IPTADR2.tim')\n",
    "\n",
    "for t in timfiles:\n",
    "    name = t.split('/')[-2]\n",
    "    outfile = os.path.join(lite_dir, '{}.tim'.format(name))\n",
    "    dr2u.combine_tim(t, outfile)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2) Figure of Merit Calculation\n",
    "\n",
    "Calculate the FoM for each pulsar for each PTA that observes it.\n",
    "Both of these FoMs are for white noise dominated pulsars, which is not true... You may want to modify them.\n",
    "\n",
    "## Burst with Memory FoM\n",
    "From van Haasteren & Levin (2010) the uncertainty on the measurement of the BWM amplitude is estimated for even cadence, white-noise dominated observations.  Madison *et al*. (2014) interpret this as the minimum amplitude for a 1-$\\sigma$ detection.\n",
    "\n",
    "$$ \\sigma_{h_0} = h_{min} \\sim \\sigma_\\star \\,\\Delta t^{1/2}\\, T^{-3/2}, $$\n",
    "\n",
    "where $\\sigma_\\star$ is the TOA measurment uncertainty for a single pulsar or the response weighted root-mean-squared of uncertainty for all pulsars in a network, $\\Delta t$ is the observing cadence, and $T$ the observation baseline.\n",
    "\n",
    "To maximize the sensitivity of our array we want to minimize $\\sigma_{h_0}$.\n",
    "We define our figure of merit for pulsar selection as:\n",
    "$$ \\mathrm{FoM}_{bwm} = \\frac{1}{\\sigma_{h_0}} = \\sqrt{\\frac{T^{3}}{\\left<{\\sigma_{TOA}}\\right>^2 \\, \\Delta t}} $$\n",
    "\n",
    "## Gravitational Wave Background FoM\n",
    "\n",
    "From GWB signal-to-noise ratio (SNR) scaling laws in Siemens, et al. 2013. \n",
    "\n",
    "The weak signal scaling law, used at the December 2017 busyweek at Flatiron CCA in NYC gives a FoM of:\n",
    "\n",
    "$$ \\mathrm{FoM}_{gwb} = \\frac{T^\\frac{13}{3}}{\\left<{\\sigma_{TOA}}\\right>^2 \\, \\Delta t} $$\n",
    "\n",
    "And the intermediate regime sclaing gives the following Figure of Merit:\n",
    "\n",
    "$$ \\mathrm{FoM}_{gwb} = T^\\frac{1}{2}\\left(\\frac{1}{\\left<{\\sigma_{TOA}}\\right>^2 \\, \\Delta t}\\right)^\\frac{1}{2\\beta} $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def bwm_FoM(Tobs, dt, TOAerr):\n",
    "    \"\"\"BWM FoM = 1/sig_h0\n",
    "    \"\"\"\n",
    "    return np.sqrt(Tobs**3/dt) / TOAerr\n",
    "\n",
    "def gwb_FoM_weak(Tobs, dt, TOAerr):\n",
    "    \"\"\"GWB FoM from SNR scaling laws\n",
    "    \"\"\"\n",
    "    return Tobs**(13/3) / TOAerr**2 / dt \n",
    "\n",
    "def gwb_FoM_inter(Tobs, dt, TOAerr):\n",
    "    \"\"\"GWB FoM from SNR scaling laws\n",
    "    \"\"\"\n",
    "    return Tobs**(1/2) * (TOAerr**2 * dt) ** (-1/(2*13/3)) \n",
    "\n",
    "my_FoM = bwm_FoM  # function alias for FoM of choice"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# analysis tag to use in filenames\n",
    "tag = \"bwm_10d\"\n",
    "\n",
    "# filter parameters\n",
    "BW = 1.1\n",
    "DM_window = 10\n",
    "\n",
    "# best PSR selection parameters\n",
    "Tmin = 5.0  # yrs\n",
    "sigmax = 10.0  # us"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "ptas = ['PPTA', 'EPTA', 'NANOGrav']\n",
    "backends = {'NANOGrav': \n",
    "                ['327_ASP', '430_ASP', 'L-wide_ASP', 'S-wide_ASP',\n",
    "                 '327_PUPPI', '430_PUPPI', 'L-wide_PUPPI',  'S-wide_PUPPI',\n",
    "                 'Rcvr_800_GASP', 'Rcvr1_2_GASP',\n",
    "                 'Rcvr_800_GUPPI', 'Rcvr1_2_GUPPI',\n",
    "                ],\n",
    "            'PPTA': \n",
    "                ['PDFB_10CM', 'PDFB_20CM', 'PDFB_40CM',\n",
    "                 'CPSR2_20CM', 'CPSR2_50CM',\n",
    "                 'WBCORR_10CM', 'WBCORR_20CM',\n",
    "                ],\n",
    "                \n",
    "            'EPTA': \n",
    "                ['EFF.EBPP.1360', 'EFF.EBPP.1410', 'EFF.EBPP.2639',\n",
    "                 'JBO.DFB.1400', 'JBO.DFB.1520',\n",
    "                 'NRT.BON.1400', 'NRT.BON.1600', 'NRT.BON.2000',\n",
    "                 'WSRT.P1.328', 'WSRT.P1.328.C', 'WSRT.P1.323.C',\n",
    "                 'WSRT.P1.382', 'WSRT.P1.382.C', 'WSRT.P1.367.C',\n",
    "                 'WSRT.P1.840', 'WSRT.P1.840.C',\n",
    "                 'WSRT.P1.1380', 'WSRT.P1.1380.C',\n",
    "                 'WSRT.P1.1380.1',\n",
    "                 'WSRT.P1.1380.2', 'WSRT.P1.1380.2.C',\n",
    "                 'WSRT.P1.2273.C',\n",
    "                ]\n",
    "           }  # backends to use for each PTA, values for `-group` flag"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Backend Notes:\n",
    "#### NANOGrav\n",
    "\n",
    "legacy, not used (`-group` flag)\n",
    "```\n",
    "kaspi\n",
    "kaspi23 (-be mk3_23m, mk2_23)\n",
    "kaspi14 (-be mk3_14m, mk2_14, mk3_14r)\n",
    "L-band_Mark3a\n",
    "L-band_Mark3b\n",
    "L-band_Mark4\n",
    "S-band_Mark4\n",
    "L-band_ABPP\n",
    "S-band_ABPP\n",
    "```\n",
    "\n",
    "#### PPTA\n",
    "\n",
    "backend list for `-f` flag\n",
    "```\n",
    "['1050CM_PDFB1', '1050CM_PDFB2', '1050CM_PDFB3', '1050CM_PDFB4',\n",
    " '1050CM_WBCORR',\n",
    " 'H-OH_CPSR2m', 'H-OH_CPSR2n', 'H-OH_PDFB1',\n",
    " 'MULTI_CPSR2m', 'MULTI_CPSR2n',\n",
    " 'MULTI_PDFB1', 'MULTI_PDFB2', 'MULTI_PDFB3', 'MULTI_PDFB4',\n",
    " 'MULTI_WBCORR',\n",
    "]\n",
    "```\n",
    "small number of TOAs\n",
    "```\n",
    "CPSR2_40CM -- J0437 & J0711 & J1045 & J2129 TOA total <3 TOAs / psr\n",
    "PDFB1_early_10CM -- J1824-2452A -- 7 TOAs total\n",
    "PDFB1_early_20CM -- J1824-2452A -- 1 TOA\n",
    "PDFB_50CM -- J0437 & J1045, 2 TOAs total\n",
    "PDFB1_20CM -- J1824-2452A, 4 TOAs total\n",
    "PDFB1_10CM -- J1824-2452A, 13 TOAs total\n",
    "WBCORR_20CM -- 11 Parkes PSRs, 30 TOAs for J0437, rest with <10 (still used... should it be?)\n",
    "```\n",
    "there are some backends with `legacy` in the name... I'm not using those.  They have lots of missing metadata.\n",
    "```\n",
    "PKS.fptm.20cm-legacy\n",
    "PKS.s2.20cm-legacy\n",
    "PKS.cpsr1.20cm_legacy\n",
    "PKS.cpsr2.20cm_legacy\n",
    "PKS.AFB.20cm_legacy\n",
    "PKS.cpsr2.50cm_legacy\n",
    "```\n",
    "\n",
    "#### EPTA\n",
    "\n",
    "possible legacy, only used for single freq J1939\n",
    "```\n",
    "NRT.DDS.1400\n",
    "```\n",
    "small number of TOAs\n",
    "```\n",
    "JBO.DFB.5000 -- J1713, 4 TOAs\n",
    "\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Construct `FoM_dict` (takes a long time, skip if already saved)\n",
    "\n",
    "Runs `filter_psr()` for each pulsar for each PTA, computing the FoM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "redo = False\n",
    "\n",
    "dict_file = 'FoM_output/FoM_{}_dict.json'.format(tag)  # where to save FoM dictionary\n",
    "data_file = 'FoM_output/FoM_{}.dat'.format(tag)  # plain text file of FoM data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if redo or not os.path.isfile(dict_file):\n",
    "    FoM_dict = {}\n",
    "    fout = open(data_file, 'w')\n",
    "    fout.write('{:11}{:9}{:9}{:9}{:9}{:9}{:12}\\n'\n",
    "               .format(' Pulsar', ' PTA', ' raw_Tobs', '  mf_Tobs', '  TOAerr', '  cadence', '      FoM'))\n",
    "\n",
    "    for psrName in psrlist:\n",
    "        FoM_dict[psrName] = {}\n",
    "    \n",
    "        parfile = os.path.join(lite_dir, '{}.par'.format(psrName))\n",
    "        timfile = os.path.join(lite_dir, '{}.tim'.format(psrName))\n",
    "        psr = t2.tempopulsar(parfile, timfile, maxobs=30000)\n",
    "        for pta in ptas:\n",
    "            if pta in psr.flagvals('pta'):\n",
    "                FoM_dict[psrName][pta] = {}\n",
    "                filt = {'pta':[pta], 'group':backends[pta]}\n",
    "                psr = dr2u.filter_psr(psr, bw=BW, dt=DM_window, filter_dict=filt, plot=False)\n",
    "                \n",
    "                idx = ~psr.deletedmask()\n",
    "                Tobs_raw = (psr.toas().max()-psr.toas().min()) / 365.25\n",
    "                if len(psr.toas()[idx]) > 0:\n",
    "                    U, _ = utils.create_quantization_matrix(psr.toas()[idx]*86400, dt=86400, nmin=1)\n",
    "                    cadence = (psr.toas()[idx].max() - psr.toas()[idx].min()) / U.shape[1] # mean dt\n",
    "                    Tobs = (psr.toas()[idx].max()-psr.toas()[idx].min()) / 365.25\n",
    "                    sigma = 1 / np.mean(1/psr.toaerrs[idx])  # harmonic mean TOAerr\n",
    "                    FoM = my_FoM(Tobs, cadence, sigma)\n",
    "                else:\n",
    "                    Tobs = 0\n",
    "                    sigma = np.inf\n",
    "                    cadence = np.inf\n",
    "                    FoM = 0\n",
    "                fout.write('{:11}{:9}{:9.2f}{:9.2f}{:9.2f}{:9.2f}{:9.2f}\\n'\n",
    "                            .format(psrName, pta, Tobs_raw, Tobs, sigma, cadence, FoM))\n",
    "                FoM_dict[psrName][pta]['Tobs'] = Tobs\n",
    "                FoM_dict[psrName][pta]['dt'] = cadence\n",
    "                FoM_dict[psrName][pta]['sigma'] = sigma\n",
    "                FoM_dict[psrName][pta]['FoM'] = FoM\n",
    "        del psr\n",
    "    fout.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### save `FoM_dict` as `json` database for later usage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "if redo or not os.path.isfile(dict_file):\n",
    "\n",
    "    # convert FoM_dict to python float (json can't handle numpy.float128)\n",
    "    for psr in FoM_dict.keys():\n",
    "        for pta in FoM_dict[psr].keys():\n",
    "            for key in FoM_dict[psr][pta].keys():\n",
    "                FoM_dict[psr][pta][key] = float(FoM_dict[psr][pta][key])\n",
    "\n",
    "    with open(dict_file, 'w') as fout:\n",
    "        json.dump(FoM_dict, fout, sort_keys=True, indent=4, separators=(',', ': '))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3) Select pulsars to use\n",
    "\n",
    "Use only PSRs with `Tmin` yrs of multi-frequency data.  Pick one PTA per pulsar."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#dict_file = 'FoM_output/FoM_{}_dict.json'.format(tag)  # FoM dictionary, same as above\n",
    "with open(dict_file, 'r') as fin:\n",
    "    FoM_dict = json.load(fin)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "psrdict = {}\n",
    "\n",
    "for psrname in psrlist:\n",
    "    best = {'pta':None, 'Tobs':0, 'FoM':0}\n",
    "    for pta in ptas:\n",
    "        try:\n",
    "            this = FoM_dict[psrname][pta]\n",
    "            if this['Tobs'] > Tmin and this['sigma'] < sigmax and this['FoM'] > best['FoM']:\n",
    "                best = {'pta':pta, 'Tobs':this['Tobs'], 'FoM':this['FoM']}\n",
    "        except KeyError:\n",
    "            pass\n",
    "    if best['pta'] is not None:\n",
    "        psrdict.update({psrname: {'pta':[best['pta']], 'group':backends[best['pta']]}})\n",
    "        print('{}: {pta:8}, T = {Tobs:5.2f} yr, FoM = {FoM:6.2f}'.format(psrname, **best))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "E,N,P = 0,0,0\n",
    "for val in psrdict.values():\n",
    "    val = val['pta']\n",
    "    if 'EPTA' in val:\n",
    "        E += 1\n",
    "    elif 'NANOGrav' in val:\n",
    "        N += 1\n",
    "    elif 'PPTA' in val:\n",
    "        P += 1\n",
    "print(' EPTA - {:2d}\\n   NG - {:2d}\\n PPTA - {:2d}'.format(E,N,P))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# save `psrdict` to file\n",
    "fname = 'psrdicts/FoM_{}_psrs.json'.format(tag)\n",
    "with open(fname, 'w') as fout:\n",
    "    json.dump(psrdict, fout, sort_keys=True, indent=4, separators=(',', ': '))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3) Create filtered par and tim files for DR2-lite analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# directory where you want the dataset saved... change for your usage\n",
    "rundir = '/Users/ptb/Projects/pulsar/11yr/bwm_results/dr2lite/partim_bwm'\n",
    "\n",
    "dr2u.make_dataset(psrdict, indir=lite_dir, outdir=rundir,\n",
    "                  bw=BW, dt=DM_window, plot=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "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.6.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}