ridge/scripts/01_crabmodel.py
2024-10-31 13:56:16 +03:00

315 lines
8.7 KiB
Python
Executable File

#!/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 import wcs
import matplotlib.pyplot as plt
import math, sys, os
import pickle
from numpy.polynomial import Polynomial
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 astropy.coordinates import SkyCoord # High-level coordinates
from astropy.coordinates import ICRS, Galactic, FK4, FK5 # Low-level frames
from astropy.coordinates import Angle, Latitude, Longitude # Angles
import astropy.units as u
#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 *
if not os.path.exists(proddir):
os.makedirs(proddir)
enkey = sys.argv[1]
sigma=3
# for these bands, slope is taken from stacked profile
fixed_slope = ['A02','A03',
'B02','B03','B04','B05','B06','B07','B08','B09','B10','B11','B12','B13','B14','B15','B16','B17','B18','B19','B20','B21']
# for these bands, slope is free for each revolution
free_slope = ['A01', 'B01']
# for these bands, slope is fixed at constant (or if positive, which is not allowed)
const_slope = ['A02','A03','B18','B19','B20','B21']
# for stacked profile, skip orbits>800 for energy channels <30 keV
skip800 = ['B01',]
# some static revs/scws to be removed
ignore_orbits=[352,834,912,1019,1021,1028,2275,2405,2493]
ignore_scws=['066600420020','066600420030','132800350010','090200390010','269500190010']
fn="detcnts.{}.fits".format(enkey)
with open(proddir+fn.replace(".fits",".ignored_scw.pkl"), 'rb') as fp:
ignored_scw = pickle.load(fp)
with fits.open(datadir+fn) as data:
df = pd.DataFrame(data[1].data)
crab_crd = SkyCoord(crab_ra, crab_dec, frame=FK5(), unit="deg")
plotme=False
npix = 50
sw = 30.0 # deg
pix = sw/npix # pixel size in degrees
mean_map = [[0.0 for i in range(npix)] for j in range(npix)]
count_map = [[0 for i in range(npix)] for j in range(npix)]
crabmodel={}
rota_arr=[]
a0=[]
b0=[]
a_full0=[]
b_full0=[]
b_est0=[]
err0=[]
rev0=[]
totx=[]
toty=[]
for i,rec in df.iterrows():
obsid = rec['OBSID']#.decode("utf-8")
if(obsid in ignore_scws):
print("Remove {}".format(obsid))
df.drop(index=i, inplace=True)
if (obsid in ignored_scw):
print("Skip ScW",obsid)
continue
# accumulate full data set
for rev in range(revmin,revmax):
if(rev in ignore_orbits):
continue
if(enkey in skip800):
if(rev > 800):
continue
df0 = df.query('SRC > 0.0 & REV == {} & PHASE > {} & PHASE < {} & CRAB_SEP < {}'.format(rev,phmin,phmax,crab_sep_max))
nobs=len(df0)
if not (nobs> crab_nmax):
continue
print(rev,nobs)
for n in df0['CRAB_SEP'].values:
totx.append(n)
for n in df0['SRC'].values:
toty.append(n)
#if(np.min(df0['SRC'])<1e-3):
# print(rev)
# sys.exit()
x = np.array(totx)
y = np.array(toty)
x = x.reshape((-1, 1))
model = LinearRegression()
#model = TheilSenRegressor()
results = evaluate_model(x, y, model)
a_full,b_full,err_full = plot_best_fit(x, y, model)
if(plotme):
plot_ab(x, y, a_full, b_full, err_full, title="REGRESSION")
#sys.exit()
# go over orbits
poly_x=[]
poly_y=[]
ntotal=0
nrev=0
for rev in range(revmin,revmax):
if(rev in ignore_orbits):
continue
df0 = df.query('SRC > 0.0 & REV == {} & PHASE > {} & PHASE < {} & CRAB_SEP < {}'.format(rev,phmin,phmax,crab_sep_max))
nobs=len(df0)
if not (nobs> crab_nmax):
continue
cen_ra = np.array(df0['RA'].values)
cen_dec = np.array(df0['DEC'].values)
print("*** Orbit ",rev)
x = np.array(df0['CRAB_SEP'].values)
y = np.array(df0['SRC'].values)
rota_deg = np.array(df0['ROTA'].values)
rota = np.array(df0['ROTA'].values) * np.pi / 180. # in radians
detx = np.cos(rota)*x
dety = np.sin(rota)*x
for i in range(rota_deg.shape[0]):
rota_arr.append(rota_deg[i])
w = wcs.WCS(naxis=2)
w.wcs.crpix = [npix/2, npix/2]
w.wcs.cdelt = np.array([-pix, pix])
w.wcs.crota = [rota_deg[i], rota_deg[i]]
w.wcs.crval = [cen_ra[i], cen_dec[i]]
w.wcs.ctype = ["RA---TAN", "DEC--TAN"]
#w.wcs.set_pv([(2, 1, 45.0)])
#header = w.to_header()
#hdu = fits.PrimaryHDU(header=header)
pixcrd = w.wcs_world2pix([(crab_ra,crab_dec)], 1)
ix=int(np.round(pixcrd[0][0]))
iy=int(np.round(pixcrd[0][1]))
#print(ix,iy)
mean_map[ix][iy] += y[i]
count_map[ix][iy] += 1
x = x.reshape((-1, 1))
model = LinearRegression()
#model = TheilSenRegressor()
results = evaluate_model(x, y, model)
a,b,err = plot_best_fit(x, y, model)
b_est = np.mean(y - a_full*x)
if(enkey in const_slope or a > 0.0):
filtered_data = sigma_clip(y, sigma=sigma, maxiters=10, return_bounds=True)
filtered_y = filtered_data[0]
filtered_min = filtered_data[1]
filtered_max = filtered_data[2]
b = np.mean(filtered_y)
a = 0.0
err = np.sqrt(np.sum((b-filtered_y)**2))/len(filtered_y)
b_est = b
a_full = a
a_full0.append(a_full)
b_full0.append(b_full)
b_est0.append(b_est)
#if(b>2.0e-4):
# print(rev)
# sys.exit()
a0.append(a)
b0.append(b)
err0.append(err)
rev0.append(rev)
poly_x.append(rev)
poly_y.append(b_est)
if(enkey in fixed_slope):
crabmodel[rev]={'a':a_full, 'b':b_est, 'err':err}
if(plotme):
plot_ab(x, y, a_full, b_est, err, title="REGRESSION rev {}".format(rev))
if(enkey in free_slope):
crabmodel[rev]={'a':a, 'b':b, 'err':err}
if(plotme):
plot_ab(x, y, a, b, err, title="REGRESSION rev {}".format(rev))
print("ax+b: a={:.2e}, b={:.2e}, std={:.2e}\n".format(a,b,err))
ntotal=ntotal+nobs
nrev=nrev+1
print("Orbits: {}, obs: {}".format(nrev,ntotal))
for i in range(npix):
for j in range(npix):
if(count_map[i][j]>0):
mean_map[i][j]=mean_map[i][j]/count_map[i][j]
w = wcs.WCS(naxis=2)
w.wcs.crpix = [npix/2, npix/2]
w.wcs.cdelt = np.array([-pix, pix])
w.wcs.crval = [0.0, 0.0]
w.wcs.ctype = ["RA---TAN", "DEC--TAN"]
header = w.to_header()
hdu = fits.PrimaryHDU(header=header)
hdu.data=mean_map
hdu.writeto(proddir+fn.replace(".fits",".crab_mean_map.fits"), overwrite=True)
hdu.data=count_map
hdu.writeto(proddir+fn.replace(".fits",".crab_count_map.fits"), overwrite=True)
npoly=4
z = np.polyfit(poly_x, poly_y, npoly)
p = np.poly1d(z)
poly_z=[]
for t in poly_x:
poly_z.append(p(t))
plt.scatter(poly_x, poly_y)
plt.plot(poly_x, poly_z, color='r')
plt.title("Crab detector count rate evolution")
plt.ylabel("Crab count rate cts/s/pix")
plt.xlabel("INTEGRAL orbit")
#plt.show()
plt.savefig(proddir+fn.replace(".fits",".crab_rate.png"))
indices = sorted(
range(len(a0)),
key=lambda index: a0[index]
)
coldefs = fits.ColDefs([
fits.Column(name='REV', format='J', unit='', array=[rev0[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]),
fits.Column(name='A_FULL', format='D', unit='', array=[a_full0[index] for index in indices]),
fits.Column(name='B_FULL', format='D', unit='', array=[b_full0[index] for index in indices]),
fits.Column(name='B_EST', format='D', unit='', array=[b_est0[index] for index in indices]),
fits.Column(name='B_POLY', format='D', unit='', array=[poly_z[index] for index in indices]),
])
fout = fn.replace(".fits",".crabmodel.fits")
hdu = fits.BinTableHDU.from_columns(coldefs, name='GRXE')
hdu.header['MISSION'] = ('INTEGRAL', '')
hdu.header['TELESCOP'] = ('IBIS', '')
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)
if not (len(crabmodel)):
print("Crab model is not saved")
sys.exit()
with open(proddir+fn.replace(".fits",".crabmodel.pkl"), 'wb') as fp:
pickle.dump([crabmodel, z], fp, protocol=pickle.HIGHEST_PROTOCOL)
with open(proddir+fn.replace(".fits",".rota.dat"), 'w') as fp:
for i in range(len(rota_arr)):
fp.write("{} {}\n".format(i,rota_arr[i]))