#!/usr/bin/env python

"""

Created by stevertaylor

Copyright (c) 2014 Stephen R. Taylor

Code contributions by Rutger van Haasteren (piccard) and Justin Ellis (PAL/PAL2).

"""

from __future__ import division

import os, math, optparse, time, cProfile

import json, sys

import cPickle as pickle

from time import gmtime, strftime

from collections import OrderedDict

import h5py as h5

import numpy as np

from numpy import *

from numpy import random

from scipy import integrate

from scipy import optimize

from scipy import constants as sc

from scipy import special as ss

from scipy import linalg as sl

from scipy.interpolate import interp1d

import numexpr as ne

import ephem

from ephem import *

import libstempo as T2

import NX01_AnisCoefficients as anis

import NX01_utils as utils

import NX01_psr

try:

import NX01_jitter as jitter

except ImportError:

print "You do not have NX01_jitter.so. " \

"Trying to make the .so file now..."

import pyximport

pyximport.install(setup_args={"include_dirs":np.get_include()},

reload_support=True)

try:

import NX01_jitter as jitter

except ImportError:

error_warning = """\

_____ __ __ _____ ____ _____ _______ ______ _____ _____ ____ _____ _ _

|_ _| \/ | __ \ / __ \| __ \__ __| | ____| __ \| __ \ / __ \| __ \| | |

| | | \ / | |__) | | | | |__) | | | | |__ | |__) | |__) | | | | |__) | | |

| | | |\/| | ___/| | | | _ / | | | __| | _ /| _ /| | | | _ /| | |

_| |_| | | | | | |__| | | \ \ | | | |____| | \ \| | \ \| |__| | | \ \|_|_|

|_____|_| |_|_| \____/|_| \_\ |_| |______|_| \_\_| \_\\____/|_| \_(_|_)

_____ ____ __ __ _____ _____ _ ______ _ _____ _______ _______ ______ _____

/ ____/ __ \| \/ | __ \_ _| | | ____| | |_ _|__ __|__ __| ____| __ \

| | | | | | \ / | |__) || | | | | |__ | | | | | | | | | |__ | |__) |

| | | | | | |\/| | ___/ | | | | | __| _ | | | | | | | | | __| | _ /

| |___| |__| | | | | | _| |_| |____| |____ | |__| |_| |_ | | | | | |____| | \ \

\_____\____/|_| |_|_| |_____|______|______| \____/|_____| |_| |_| |______|_| \_\

"""

print error_warning

print "You need to run: " \

"python setup-cython.py build_ext --inplace"

sys.exit()

try:

from mpi4py import MPI

comm = MPI.COMM_WORLD

rank = comm.Get_rank()

except ImportError:

print 'Do not have mpi4py package.'

import nompi4py as MPI

comm = MPI.COMM_WORLD

rank = comm.Get_rank()

parser = optparse.OptionParser(description = "NX01 - It's been a long road, getting from there to here")

############################

############################

parser.add_option('--sampler', dest='sampler', action='store', type=str, default='ptmcmc',

help='Which sampler do you want to use: PTMCMC (ptmcmc), MultiNest (mnest), or Polychord (pchord) (default = ptmcmc)')

parser.add_option('--ins', dest='ins', action='store_true', default=False,

help='Switch on importance nested sampling for MultiNest (default = False)')

parser.add_option('--nlive', dest='nlive', action='store', type=int, default=500,

help='Number of live points for MultiNest or Polychord (default = 500)')

parser.add_option('--sampleEff', dest='sampleEff', action='store', type=float, default=0.3,

help='Sampling efficiency for MultiNest (default = 0.3)')

parser.add_option('--constEff', dest='constEff', action='store_true', default=False,

help='Run MultiNest in constant efficiency mode? (default = False)')

parser.add_option('--resume', dest='resume', action='store_true', default=False,

help='Do you want to resume the sampler (default = False)')

parser.add_option('--nmodes', dest='nmodes', action='store', type=int,

help='Number of modes in low-rank time-frequency approximation')

parser.add_option('--cadence', dest='cadence', action='store', type=float, default=14.0,

help='Number days between successive observations (default = 14 days)')

parser.add_option('--parfile', dest='parfile', action='store', type=str,

help='Full path to parfile')

parser.add_option('--timfile', dest='timfile', action='store', type=str,

help='Full path to timfile')

parser.add_option('--efacequad-sysflag', dest='systarg', action='store', type=str, default='group',

help='Which system flag should the EFACs/EQUADs target? (default = \'group\')')

parser.add_option('--redPrior', dest='redPrior', action='store', type=str, default='loguniform',

help='What kind of prior to place on the red noise amplitude? (default = \'loguniform\')')

parser.add_option('--dmPrior', dest='dmPrior', action='store', type=str, default='loguniform',

help='What kind of prior to place on the DM variation amplitude? (default = \'loguniform\')')

parser.add_option('--incDM', dest='incDM', action='store_true', default=False,

help='Search for DM variations in the data (False)? (default=False)')

parser.add_option('--fullN', dest='fullN', action='store_true', default=False,

help='Search for EFAC/EQUAD/ECORR over all systems (True), or just apply a GEFAC (False)? (default=False)')

parser.add_option('--grab_planets', dest='grab_planets', action='store_true', default=False,

help='Grab the planet position vectors at the TOA timestamps? (default=False)')

parser.add_option('--incGlitch', dest='incGlitch', action='store_true', default=False,

help='Search for a glitch in the pulsar? (default=False)')

parser.add_option('--jitterbin', dest='jitterbin', action='store', type=float, default=1.0,

help='What time duration (in seconds) do you want a jitter bin to be? (default = 1.0)')

parser.add_option('--mnest', dest='mnest', action='store_true', default=False,

help='Do you want to sample using MultiNest? (default=False)')

parser.add_option('--writeHotChains', dest='writeHotChains', action='store_true', default=False,

help='Do you want to store hot PTMCMC chains? (default=False)')

parser.add_option('--hotChain', dest='hotChain', action='store_true', default=False,

help='Do you want to include a very hot chain? (default=False)')

parser.add_option('--dirExt', dest='dirExt', action='store', type=str, default='./chains_nanoAnalysis/',

help='What master directory name do you want to put this run into? (default = ./chains_nanoAnalysis/)')

(args, x) = parser.parse_args()

header = """\

"""

if rank == 0:

print header

if args.nmodes:

print ("\n You've given me the number of frequencies",

"to include in the low-rank time-frequency approximation, got it?\n")

else:

print ("\n You've given me the sampling cadence for the observations,",

"which determines the upper frequency limit and the number of modes, got it?\n")

if args.sampler == 'mnest':

import pymultinest

elif args.sampler == 'ptmcmc':

import PTMCMCSampler

from PTMCMCSampler import PTMCMCSampler as ptmcmc

#####################################################################

# PASSING THROUGH TEMPO2 VIA libstempo

#####################################################################

t2psr = T2.tempopulsar(parfile=args.parfile, timfile=args.timfile)

t2psr.fit(iters=10)

if np.any(np.isfinite(t2psr.residuals())==False)==True:

t2psr = T2.tempopulsar(parfile=args.parfile,timfile=args.timfile)

psr = NX01_psr.PsrObj(t2psr)

psr.grab_all_vars(jitterbin=args.jitterbin, makeGmat=False,

fastDesign=True, planetssb=args.grab_planets)

#############################################################################

# GETTING MAXIMUM TIME, COMPUTING FOURIER DESIGN MATRICES, AND GETTING MODES

#############################################################################

Tmax = psr.toas.max() - psr.toas.min()

if args.nmodes:

psr.makeTe(args.nmodes, Tmax, makeDM=args.incDM)

# get GW frequencies

fqs = np.linspace(1/Tmax, args.nmodes/Tmax, args.nmodes)

nmode = args.nmodes

else:

nmode = int(round(Tmax/args.cadence))

psr.makeTe(nmode, Tmax, makeDM=args.incDM)

# get GW frequencies

fqs = np.linspace(1/Tmax, nmode/Tmax, nmode)

################################################################################################################################

# FORM A LIST COMPOSED OF NP ARRAYS CONTAINING THE INDEX POSITIONS WHERE EACH UNIQUE SYSTEM IS APPLIED

################################################################################################################################

if args.fullN:

systems = psr.sysflagdict[args.systarg]

else:

systems = OrderedDict.fromkeys([psr.name])

systems[psr.name] = np.arange(len(psr.toas))

################################################################################################################################

# SETTING UP PRIOR RANGES

################################################################################################################################

pmin = np.array([-20.0,0.0])

pmax = np.array([-11.0,7.0])

if args.incDM:

pmin = np.append(pmin,np.array([-20.0,0.0]))

pmax = np.append(pmax,np.array([-8.0,7.0]))

pmin = np.append(pmin,0.001*np.ones(len(systems)))

pmax = np.append(pmax,10.0*np.ones(len(systems)))

if args.fullN:

pmin = np.append(pmin,-10.0*np.ones(len(systems)))

pmax = np.append(pmax,-3.0*np.ones(len(systems)))

if 'pta' in t2psr.flags():

if 'NANOGrav' in list(set(t2psr.flagvals('pta'))):

if len(psr.sysflagdict['nano-f'].keys())>0:

pmin = np.append(pmin, -10.0*np.ones(len(psr.sysflagdict['nano-f'].keys())))

pmax = np.append(pmax, -3.0*np.ones(len(psr.sysflagdict['nano-f'].keys())))

if args.incGlitch:

pmin = np.append(pmin,[np.min(psr.toas),-18.0])

pmax = np.append(pmax,[np.max(psr.toas),-11.0])

################################################################################################################################

# PRIOR AND LIKELIHOOD

################################################################################################################################

def my_prior(xx):

return logp

def ln_prob(xx):

return logLike + prior_fac

#########################

#########################

parameters = ["log(A_red)","gam_red"]

if args.incDM:

parameters.append("log(A_dm)")

parameters.append("gam_dm")

for ii in range(len(systems)):

parameters.append('EFAC_'+systems.keys()[ii])

if args.fullN:

for ii in range(len(systems)):

parameters.append('EQUAD_'+systems.keys()[ii])

if 'pta' in t2psr.flags():

if 'NANOGrav' in list(set(t2psr.flagvals('pta'))):

if len(psr.sysflagdict['nano-f'].keys())>0:

if rank == 0:

print "\n You have some NANOGrav ECORR parameters..."

for ii,nano_sysname in enumerate(psr.sysflagdict['nano-f'].keys()):

parameters.append('ECORR_'+nano_sysname)

if args.incGlitch:

parameters += ["glitch_epoch","glitch_lamp"]

n_params = len(parameters)

if rank == 0:

print "\n You are searching for the following single-pulsar parameters: {0}\n".format(parameters)

print "\n The total number of parameters is {0}\n".format(n_params)

# Define a unique file tag

file_tag = 'pta_'+psr.name

file_tag += '_red{0}'.format(args.redPrior)

if args.incDM:

file_tag += '_dm{0}'.format(args.dmPrior)

if args.incGlitch:

file_tag += '_glitch'

file_tag += '_nmodes{0}'.format(nmode)

#####################

# Now, we sample....

#####################

if rank == 0:

print "\n Now, we sample... \n"

print """\

_______ .__ __. _______ ___ _______ _______ __

| ____|| \ | | / _____| / \ / _____|| ____|| |

| |__ | \| | | | __ / ^ \ | | __ | |__ | |

| __| | . ` | | | |_ | / /_\ \ | | |_ | | __| | |

| |____ | |\ | | |__| | / _____ \ | |__| | | |____ |__|

|_______||__| \__| \______| /__/ \__\ \______| |_______|(__)

"""

###########################

# Define function wrappers

###########################

if args.sampler == 'mnest':

#dir_name = './chains_nanoAnalysis/nano_singlePsr/'+file_tag+'_mnest'

#dir_name = './chn_eptapsr/'+file_tag+'_mnest'

dir_name = args.dirExt+file_tag+'_mnest'

if rank == 0:

if not os.path.exists(dir_name):

os.makedirs(dir_name)

if rank == 0:

# Printing out the list of searched parameters

fil = open(dir_name+'/parameter_list.txt','w')

for ii,parm in enumerate(parameters):

print >>fil, ii, parm

fil.close()

# Saving command-line arguments to file

with open(dir_name+'/run_args.json', 'w') as fp:

json.dump(vars(args), fp)

def prior_func(xx,ndim,nparams):

for ii in range(nparams):

xx[ii] = pmin[ii] + xx[ii]*(pmax[ii]-pmin[ii])

def like_func(xx,ndim,nparams):

xx = np.array([xx[ii] for ii in range(nparams)])

return ln_prob(xx)

pymultinest.run(like_func, prior_func, n_params,

importance_nested_sampling = args.ins,

resume = args.resume, verbose = True,

n_live_points=args.nlive,

outputfiles_basename=u'{0}/mnest_'.format(dir_name),

sampling_efficiency=args.sampleEff,

const_efficiency_mode=args.constEff)

if args.sampler == 'ptmcmc':

x0 = np.array([-15.0,2.0])

cov_diag = np.array([0.5,0.5])

if args.incDM:

x0 = np.append(x0,np.array([-15.0,2.0]))

cov_diag = np.append(cov_diag,np.array([0.5,0.5]))

x0 = np.append(x0,np.random.uniform(0.75,1.25,len(systems)))

cov_diag = np.append(cov_diag,0.5*np.ones(len(systems)))

if args.fullN:

x0 = np.append(x0,np.random.uniform(-10.0,-5.0,len(systems)))

cov_diag = np.append(cov_diag,0.5*np.ones(len(systems)))

if 'pta' in t2psr.flags():

if 'NANOGrav' in list(set(t2psr.flagvals('pta'))):

if len(psr.sysflagdict['nano-f'].keys())>0:

x0 = np.append(x0, np.random.uniform(-8.5,-5.0,len(psr.sysflagdict['nano-f'].keys())))

cov_diag = np.append(cov_diag,0.5*np.ones(len(psr.sysflagdict['nano-f'].keys())))

if rank == 0:

print "\n Your initial parameters are {0}\n".format(x0)

print "\n Running a quick profile on the likelihood to estimate evaluation speed...\n"

cProfile.run('ln_prob(x0)')

########################################

# Creating parameter sampling groupings

ind = []

param_ct = 0

##### red noise #####

ids = [[0,1]]

[ind.append(id) for id in ids]

param_ct += 2

##### DM noise #####

if args.incDM:

ids = [[2,3]]

[ind.append(id) for id in ids]

param_ct += 2

##### White noise #####

if args.fullN:

efacs = [param_ct+ii for ii in range(len(systems))]

ids = [efacs]

[ind.append(id) for id in ids if len(id) > 0]

param_ct += len(systems)

##

equads = [param_ct+ii for ii in range(len(systems))]

ids = [equads]

[ind.append(id) for id in ids if len(id) > 0]

param_ct += len(systems)

##

if 'pta' in t2psr.flags():

if 'NANOGrav' in list(set(t2psr.flagvals('pta'))):

ecorrs = [param_ct+ii for ii

in range(len(psr.sysflagdict['nano-f'].keys()))]

ids = [ecorrs]

[ind.append(id) for id in ids if len(id) > 0]

param_ct += len(psr.sysflagdict['nano-f'].keys())

##### all parameters #####

all_inds = range(len(x0))

ind.insert(0, all_inds)

if rank == 0:

print "Your parameter index groupings for sampling are {0}".format(ind)

########

dir_name = args.dirExt+file_tag+'_ptmcmc'

if rank == 0:

if not os.path.exists(dir_name):

os.makedirs(dir_name)

if rank == 0:

# Printing out the list of searched parameters

fil = open(dir_name+'/parameter_list.txt','w')

for ii,parm in enumerate(parameters):

print >>fil, ii, parm

fil.close()

# Saving command-line arguments to file

with open(dir_name+'/run_args.json', 'w') as fp:

json.dump(vars(args), fp)

sampler = ptmcmc.PTSampler(ndim = n_params, logl = ln_prob,

logp = my_prior, cov = np.diag(cov_diag),

outDir='./{0}'.format(dir_name),

resume = args.resume, groups = ind)

def drawFromRedNoisePrior(parameters, iter, beta):

return q, qxy

def drawFromDMNoisePrior(parameters, iter, beta):

return q, qxy

def drawFromEfacPrior(parameters, iter, beta):

return q, qxy

def drawFromEquadPrior(parameters, iter, beta):

return q, qxy

def drawFromEcorrPrior(parameters, iter, beta):

return q, qxy

# add jump proposals

sampler.addProposalToCycle(drawFromRedNoisePrior, 10)

if args.incDM:

sampler.addProposalToCycle(drawFromDMNoisePrior, 10)

if args.fullN:

sampler.addProposalToCycle(drawFromEfacPrior, 10)

sampler.addProposalToCycle(drawFromEquadPrior, 10)

if 'pta' in t2psr.flags():

if 'NANOGrav' in list(set(t2psr.flagvals('pta'))):

if len(psr.sysflagdict['nano-f'].keys())>0:

sampler.addProposalToCycle(drawFromEcorrPrior, 10)

sampler.sample(p0=x0, Niter=int(5e6), thin=10,

covUpdate=1000, AMweight=20,

SCAMweight=30, DEweight=50,

writeHotChains=args.writeHotChains,

hotChain=args.hotChain)