tesswrapper.py

import os
os.environ['MKL_NUM_THREADS'] = '1'
os.environ['OPENBLAS_NUM_THREADS'] = '1'
os.environ['NUMEXPR_NUM_THREADS'] = '1'
os.environ['OMP_NUM_THREADS'] = '1'

import numpy as np
import lightkurve as lk
import matplotlib.pyplot as plt
import matplotlib as mpl
import pandas as pd
import scipy.signal as sps
import astropy.timeseries as ts
from astropy.io import fits
import sys, time, argparse
from spinneret import *

parser = argparse.ArgumentParser(description='Process a Kepler quarter and TESSified sector using spinneret')
parser.add_argument('-f', '--filename', required=True, type=str, help='FITS file name')
# parser.add_argument('-t', '--tic', required=True, type=int, help='TESS ID')
parser.add_argument('-s', '--sector', required=True, type=int, help='TESS sector')
parser.add_argument('-p', '--gtp', required=False, type=float, default=0, help='Ground truth period')

params = parser.parse_args()

fn = params.filename
# tid = params.tic
sec = params.sector
gtp = params.gtp

printsec = ''
if sec < 10:
    printsec = f'0{sec}'
else:
    printsec = str(sec)

# filename = f'/shared_data/vanir/TESS/LightCurves/sector{printsec}/{fn}_lc.fits'
filename = f'{fn}_lc.fits'

cadence = 1/24/30 # 2min

try:
    hdu = fits.open(filename)
except FileNotFoundError:
    print(f'NO DATA: {filename}')
    sys.exit()

table = hdu[1].data
time = table['TIME']
flux = table['PDCSAP_FLUX']

head = hdu[0].header
tid = head['TICID']
teff = head['TEFF']
logg = head['LOGG']
tmag = head['TESSMAG']

head2 = hdu[1].header
cr = head2['CROWDSAP']

hdu.close()

ticstr = ''
if len(str(tid)) == 7:
    ticstr = f'00{tid}'
elif len(str(tid)) == 8:
    ticstr = f'0{tid}'
else:
    ticstr = str(tid)

# directorymaker('figs')

dir_name = f'sector{printsec}'
directorymaker(dir_name)

time, flux = nancleaner2d(time, flux)
time, flux = clip(time, flux, 3) #3 sigma clip
flux = lk.LightCurve(time=time, flux=flux).normalize().flux.value - 1

p_grid = np.linspace(0.001, 30, 5000)
freq = 1/p_grid

target = Spinner(time, flux, teff=teff, logg=logg, crowding=cr, tmag=tmag)

# freq, ps = ts.LombScargle(time, flux).autopower(method='fastchi2', maximum_frequency=200, samples_per_peak=20)
ps = ts.LombScargle(time, flux).power(freq)
target.ls_one_term(freq, ps)

# freq, ps = ts.LombScargle(time, flux, nterms=2).autopower(method='fastchi2', maximum_frequency=200, samples_per_peak=20)
ps = ts.LombScargle(time, flux).power(freq)
target.ls_two_term(freq, ps)

lags_raw, acf_raw, lags, acf, _x, _y = simple_acf(time, flux, cadence, width=16)
target.acf(lags, acf)

fig1 = target.diagnostic_plot(heading=f'TIC {tid}')
# figsaver(fig1, f'TIC{tid}.png', '/home/isy/Documents/Work/rotation/figs')
figsaver(fig1, f'TIC{tid}.png')
filemaker(target, tid, 0, filename=f'TIC{tid}.csv', filepath=f'./{dir_name}')

# print(f'{tid} done')