102 lines
3.8 KiB
Python
102 lines
3.8 KiB
Python
import numpy as np
|
|
from astropy.io import fits
|
|
import matplotlib.pyplot as plt
|
|
import pickle
|
|
from astropy.wcs import WCS
|
|
import tqdm
|
|
from multiprocessing.pool import ThreadPool
|
|
from chan_psf import solve_for_locations, solve_for_locations_eintp, solve_for_rates
|
|
|
|
psfe = np.array([1.8, 1.9, 3.0, 4.0, 6.0, 7.0, 8.0, 9.0])
|
|
|
|
def prepare_psf(evt):
|
|
"""
|
|
find all unique psf for observation and load in single 3d data cuve
|
|
return data cube with events slices indexes
|
|
"""
|
|
u, ui = np.unique(evt["psf_cube"], return_inverse=True)
|
|
data = np.array([np.load(p[3:])[:, ::-1,::-1].copy() for p in u])
|
|
return data, ui
|
|
|
|
def select_xychunksize(wcs, halfpsfsize=36./3600.):
|
|
"""
|
|
get wcs and find wcs pixel size of psf reach
|
|
"""
|
|
sizex = int(abs(halfpsfsize/wcs.wcs.cdelt[1])) + 2
|
|
sizey = int(abs(halfpsfsize/wcs.wcs.cdelt[0])) + 2
|
|
return sizex, sizey
|
|
|
|
def read_wcs(h):
|
|
"""
|
|
read events wcs header
|
|
"""
|
|
w = WCS(naxis=2)
|
|
w.wcs.ctype = [h["TCTYP11"], h["TCTYP12"]]
|
|
w.wcs.crval = [h["TCRVL11"], h["TCRVL12"]]
|
|
w.wcs.cdelt = [h["TCDLT11"], h["TCDLT12"]]
|
|
w.wcs.crpix = [h["TCRPX11"], h["TCRPX12"]]
|
|
w = WCS(w.to_header())
|
|
return w
|
|
|
|
|
|
def create_neighboring_blocks(x, y, sizex, sizey):
|
|
"""
|
|
schematically all sky is splitted on squares, which are approximatelly ~ 10 times greater then the psf
|
|
events for corresponding square are joined :: squer + diluttaion of psf reach
|
|
|
|
coordinate system with events and all required coefficiets are fed to psf solver
|
|
current psf size is 25*0.5 arcsec (with up to \sqrt(2) factor in case of worst rolls -> 36''
|
|
|
|
"""
|
|
"""
|
|
event list already contains x and y for each event
|
|
"""
|
|
iix = (x//sizex + 0.5).astype(int)
|
|
iiy = (y//sizey + 0.5).astype(int)
|
|
isx, isy = np.mgrid[-1:2:1, -1:2:1]
|
|
oidx = np.repeat(np.arange(x.size), 9)
|
|
xyu, iixy, xyc = np.unique(np.array([np.repeat(iix, 9) + np.tile(isx.ravel(), x.size),
|
|
np.repeat(iiy, 9)+ np.tile(isy.ravel(), x.size)]), axis=1, return_counts=True, return_inverse=True)
|
|
|
|
sord = np.argsort(iixy)
|
|
return oidx[sord], xyu, xyc
|
|
|
|
|
|
def lkls_for_rates(evt, expv, wcs, srcx, srcy, rates):
|
|
sizex, sizey = select_xychunksize(wcs)
|
|
x, y = evt["x"].astype(float), evt["y"].astype(float)
|
|
mask = np.logical_and.reduce([x > srcx - sizex//2, y > srcy - sizey//2, x < srcx + sizex//2, y < srcy + sizey//2], axis=0)
|
|
print("mask sum", srcx, srcy, mask.sum())
|
|
eloc = evt[mask]
|
|
pickle.dump(eloc, open("eloc.pkl", "wb"))
|
|
psfdata, ui = prepare_psf(eloc)
|
|
xe, ye = np.copy(x[mask]), np.copy(y[mask])
|
|
eidx = np.maximum(np.searchsorted(psfe*1e3, eloc["ENERGY"]) - 1, 0)
|
|
ee = np.maximum((eloc["ENERGY"]/1000. - psfe[eidx])/(psfe[eidx + 1] - psfe[eidx]), 0.).astype(float) + eidx
|
|
pk = np.copy(eloc["src_spec"]/eloc["bkg_spec"]).astype(float)
|
|
roll = np.copy(np.deg2rad(eloc["roll_pnt"])).astype(float)
|
|
#"OOOOOOOdddO", &psfi, &eidx, &x, &y, &roll, &pk, &rates, &xc, &yc, &eloc, &smat
|
|
# O O O O O O O d d d O"
|
|
print(ui, ee, xe, ye, roll, pk)
|
|
lkls = solve_for_rates(ui, ee, xe, ye, roll, pk, rates, srcx, srcy, expv, psfdata)
|
|
return lkls
|
|
|
|
|
|
if __name__ == "__main__":
|
|
p1 = fits.open("test.fits")
|
|
ewcs = read_wcs(p1[1].header)
|
|
wcs = WCS(fits.getheader("eR_spec_asp_0.fits.gz", 0))
|
|
xc, yc = 4290, 4147
|
|
xc, yc = 4643, 4223.7
|
|
#xc, yc = 4147,4290
|
|
xc, yc = ewcs.all_world2pix(wcs.all_pix2world([[xc, yc],], 0), 0).T
|
|
print(xc, yc)
|
|
eloc = 0.025 #0.0283
|
|
#rates = np.array([4.2/eloc,]) #np.logspace(-0.5, 0.5, 129)*4.2/eloc
|
|
rates = np.logspace(-0.5, 0.5, 129)*1352/eloc #*4.2/eloc
|
|
lkls = lkls_for_rates(p1[1].data, eloc, ewcs, xc, yc, rates)
|
|
plt.plot(rates, lkls)
|
|
plt.axvline(rates[64])
|
|
plt.show()
|
|
|