# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
import numpy as np
import glob, os, json
import pickle

import matplotlib.pyplot as plt
import corner
# %matplotlib inline
# %config InlineBackend.figure_format = 'retina'

import enterprise
from enterprise.pulsar import Pulsar
import enterprise.signals.parameter as parameter
from enterprise.signals import utils
from enterprise.signals import signal_base
from enterprise.signals import selections
from enterprise.signals.selections import Selection
from enterprise.signals import white_signals
from enterprise.signals import gp_signals
from enterprise.signals import deterministic_signals

from PTMCMCSampler.PTMCMCSampler import PTSampler as ptmcmc

import libstempo as libs
import libstempo.plot as LP, libstempo.toasim as LT

# +
# import the par and tim files
datadir = './data/partim/'

parfiles = sorted(glob.glob(datadir + '*.par'))
timfiles = sorted(glob.glob(datadir + '*.tim'))

psrs = []
for p, t in zip(parfiles, timfiles):
    psrname = parfiles[0].split('/')[-1].split('_')[0]
    psr = Pulsar(p, t, ephem='DE438')
    psrs.append(psr)
# -

for psr in psrs:
    print(psr.name)

# +
## Get parameter noise dictionary
noisedir = os.getcwd() + '/data/noisefiles/'

noise_params = {}
for filepath in glob.iglob(noisedir + '/*.json'):
    with open(filepath, 'r') as f:
        noise_params.update(json.load(f))
# print(noise_params)
# -

print(len(parfiles))

# cut down to pulsars that have been measured for 6 years
psrcut = []
for psr in psrs:
    tmax = max(psr.toas)
    tmin = min(psr.toas)
    Tspan = tmax - tmin
    if Tspan / 525600 / 60 > 6:
        psrcut.append(psr)
    # print(psr.name)
print(len(psrcut))


pickle.dump(psrcut, open("psrs.p", "wb"))

psrcut = pickle.load(open("psrs.p", "rb"))

print(len(psrcut))

# import each pulsar into tempo2
tempopsr = []
for psr in psrcut:
    parfiles = sorted(glob.glob(datadir + psr.name + '*.par'))
    timfiles = sorted(glob.glob(datadir + psr.name + '*.tim'))
    psr = libs.tempopulsar(parfile = parfiles[0],
                           timfile = timfiles[0], maxobs=50000)
    tempopsr.append(psr)

# remove all noise from the pulsars
# add in EFAC noise
for psr in tempopsr:
    LT.make_ideal(psr)
    LT.add_efac(psr, efac=1.0)
    # psr.fit()
# LT.createGWB(tempopsr, Amp=2e-15, gam=13./3.)

for psr in tempopsr:
    print(psr.name)
    rednoise_toadd = noise_params[str(psr.name) + '_log10_A']
    # print(psr.name)
    gamma = noise_params[str(psr.name) + '_gamma']
    LT.add_rednoise(psr, 10**rednoise_toadd, gamma, components=30)
    psr.fit()

# create and add GWB to all the pulsars
LT.createGWB(tempopsr, Amp=3e-15, gam=13./3.)
for psr in tempopsr:
    psr.fit()

# +
# save the pulsars
pardir = './fakes_gwb3/par/'
timdir = './fakes_gwb3/tim/'
if not os.path.exists(pardir):
    os.makedirs(pardir)
    os.makedirs(timdir)


for psr in tempopsr:
    psr.savepar(pardir + psr.name + '-sim.par')
    psr.savetim(timdir + psr.name + '-sim.tim')
    T.purgetim(timdir + psr.name + '-sim.tim')  # fix the tim files
# -



StarcoderdataJupyterScriptsDedupFiltered