ridge/scripts/02_grxe_resid.py

#!/usr/bin/env python

__author__      = "Roman Krivonos"
__copyright__   = "Space Research Institute (IKI)"

import numpy as np
import pandas as pd
from astropy.io import fits
from astropy.table import Table, Column
import matplotlib.pyplot as plt
import math, sys
import pickle

from sklearn.linear_model import LinearRegression
from sklearn.linear_model import HuberRegressor
from sklearn.linear_model import RANSACRegressor
from sklearn.linear_model import TheilSenRegressor
from sklearn.model_selection import cross_val_score
from sklearn.model_selection import RepeatedKFold

#from statsmodels.robust.scale import huber
from astropy.stats import sigma_clip
from numpy import absolute
from numpy import arange

from ridge.utils import *
from ridge.config import *

enkey = sys.argv[1]
outkey = sys.argv[2]

fn="detcnts.{}.fits".format(enkey)

with open(proddir+fn.replace(".fits",".pkl"), 'rb') as fp:
    bgdmodel = pickle.load(fp)

with open(proddir+fn.replace(".fits",".ignored_scw.pkl"), 'rb') as fp:
    ignored_scw = pickle.load(fp)

with open(proddir+fn.replace(".fits",".crabmodel.pkl"), 'rb') as fp:
    crabmodel, z = pickle.load(fp)
    p = np.poly1d(z)
    #print(crabmodel)

crab_rev_max = np.max(list(crabmodel.keys()))
print("Crab is defined untill orbit {}".format(crab_rev_max))

with fits.open(datadir+fn) as hdul:
    data=hdul[1].data

#print(data.columns)

rev   = data.field('rev')
mjd   = data.field('mjd')
clean = data.field('clean')
phase = data.field('phase')


rev0=[]
phase0=[]
clean0=[]
model0=[]
resid0=[] # residuals in cts/s
grxe0=[] # mCrab
crab0=[] # Crab count rate
mjd0=[]
a0=[]
b0=[]
err0=[]
lon0=[]
lat0=[]

d  = fits.getdata(datadir+fn)
df = pd.DataFrame(np.array(d).byteswap().newbyteorder())

# BKG
if(outkey == 'BKG'):
    df = df.query('CLEAN > 0.0 & ( abs(LAT) > {} | abs(LON) > {}) & PHASE > {} & PHASE < {}'.format(bmax,lmax,phmin,phmax))

# GAL
if(outkey=='GAL'): 
    df = df.query('CLEAN > 0.0 & abs(LAT) < {} & abs(LON) < {} & PHASE > {} & PHASE < {}'.format(bmax,lmax,phmin,phmax))

# ALL
if(outkey=='ALL'): 
    df = df.query('CLEAN > 0.0 & PHASE > {} & PHASE < {}'.format(phmin,phmax))
    
for i, row in df.iterrows():
    orbit=row['REV']
    obsid=row['OBSID'].decode("UTF-8")
    
    if not (orbit > revmin and orbit < revmax):
        print("Skip orbit",orbit,row['OBSID'])
        continue

    if not (orbit < crab_rev_max):
        print("Skip orbit",orbit,row['OBSID'])
        continue

    if (obsid in ignored_scw):
        print("Skip ScW",obsid)
        continue

    a = bgdmodel[orbit]['a']
    b = bgdmodel[orbit]['b']
    err = bgdmodel[orbit]['err']
    m = a*row['PHASE']+b

    clean0.append(clean[i])
    mjd0.append(mjd[i])
    model0.append(m)
    resid0.append(clean[i]-m)
    grxe0.append(1000*(clean[i]-m)/p(orbit))
    crab0.append(p(orbit))
        
    a0.append(a)
    b0.append(b)
    err0.append(err)
    phase0.append(row['PHASE'])
    rev0.append(orbit)
    lon0.append(row['LON'])
    lat0.append(row['LAT'])


indices = sorted(
    range(len(mjd0)),
    key=lambda index: mjd0[index]
)

coldefs = fits.ColDefs([
    #fits.Column(name='OBSID', format='11A', array=[obs_id[index] for index in indices]), 
    #fits.Column(name='RA', format='D', unit='deg',  array=[ra[index] for index in indices]),
    #fits.Column(name='DEC', format='D', unit='deg', array=[dec[index] for index in indices]),
    fits.Column(name='LON', format='D', unit='deg', array=[lon0[index] for index in indices]),
    fits.Column(name='LAT', format='D', unit='deg', array=[lat0[index] for index in indices]),
    fits.Column(name='REV', format='J', unit='', array=[rev0[index] for index in indices]),
    fits.Column(name='MJD', format='D', unit='', array=[mjd0[index] for index in indices]),
    fits.Column(name='PHASE', format='D', unit='', array=[phase0[index] for index in indices]),
    fits.Column(name='CLEAN', format='D', unit='cts/s', array=[clean0[index] for index in indices]),
    fits.Column(name='MODEL', format='D', unit='cts/s', array=[model0[index] for index in indices]),
    fits.Column(name='RESID', format='D', unit='cts/s', array=[resid0[index] for index in indices]),
    fits.Column(name='GRXE', format='D', unit='mCrab', array=[grxe0[index] for index in indices]),
    fits.Column(name='CRAB', format='D', unit='cts/s', array=[crab0[index] for index in indices]),
    fits.Column(name='A', format='D', unit='', array=[a0[index] for index in indices]),
    fits.Column(name='B', format='D', unit='', array=[b0[index] for index in indices]),
    fits.Column(name='ERR', format='D', unit='', array=[err0[index] for index in indices]),
])

fout = fn.replace(".fits",".{}.resid.fits".format(outkey))
hdu = fits.BinTableHDU.from_columns(coldefs, name='GRXE')
hdu.header['MISSION'] = ('INTEGRAL', '')
hdu.header['TELESCOP'] = (outkey, '')
hdu.header['INSTITUT'] = ('IKI', 'Affiliation')
hdu.header['AUTHOR'] = ('Roman Krivonos', 'Responsible person')
hdu.header['EMAIL'] = ('krivonos@cosmos.ru', 'E-mail')
#hdu.add_checksum()
print(hdu.columns)
hdu.writeto(proddir+fout, overwrite=True)

with fits.open(proddir+fout, mode='update') as hdus:
    hdus[1].add_checksum()