208 lines
7.9 KiB
Python
208 lines
7.9 KiB
Python
import asyncio
|
|
import numpy as np
|
|
from astropy.io import fits
|
|
import matplotlib.pyplot as plt
|
|
from astropy.wcs import WCS
|
|
import tqdm
|
|
from multiprocessing.pool import ThreadPool
|
|
from chan_psf import solve_for_locations
|
|
|
|
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_path"], return_inverse=True)
|
|
data = np.array([np.load(p)[::-1,::-1] 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
|
|
print(sizex, sizey)
|
|
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 make_srccount_and_detmap(emap, evt, h, wcs=None):
|
|
psfdata, ui = prepare_psf(evt)
|
|
if wcs is None:
|
|
wcs = read_wcs(h)
|
|
x, y = evt["x"], evt["y"]
|
|
else:
|
|
ewcs = read_wcs(h)
|
|
x, y = wcs.all_world2pix(ewcs.all_pix2world(np.array([x, y]).T, 0), 0).T
|
|
|
|
sizex, sizey = select_xychunksize(wcs)
|
|
iidx, xyu, cts = create_neighboring_blocks(x, y, sizex, sizey)
|
|
cc = np.zeros(cts.size + 1, int)
|
|
cc[1:] = np.cumsum(cts)
|
|
cmap, pmap = np.zeros(emap.shape, float), np.zeros(emap.shape, float)
|
|
#xe, ye, pk, roll, psfi = np.copy(evt["x"][iidx]), np.copy(evt["y"][iidx]), np.copy((evt["quant_eff"]/evt["bkg_model"])[iidx]), np.copy(evt["roll_pnt"][iidx]), np.copy(ui[iidx])
|
|
xe = np.copy(x[iidx]).astype(float)
|
|
ye = np.copy(y[iidx]).astype(float)
|
|
pk = np.copy((2e-3/evt["bkg_model"])[iidx]).astype(float)
|
|
roll = np.copy(np.deg2rad(evt["roll_pnt"][iidx])).astype(float)
|
|
psfi = np.copy(ui[iidx])
|
|
|
|
yg, xg = np.mgrid[0:sizey:1, 0:sizex:1]
|
|
|
|
def worker(ixys):
|
|
i, (xs, ys) = ixys
|
|
eloc = emap[ys*sizey:ys*sizey+sizey, xs*sizex:xs*sizex+sizex]
|
|
mask = eloc > 0.
|
|
xl = (xg[mask] + xs*sizex).astype(float)
|
|
yl = (yg[mask] + ys*sizey).astype(float)
|
|
ell = (eloc[mask]).astype(float)
|
|
if np.any(mask):
|
|
"""
|
|
mask = np.any([(np.abs(x - xl) < 20) & (np.abs(y - yl) < 20) for x, y in zip(xe[cc[i]:cc[i+1]], ye[cc[i]:cc[i+1]])], axis=0)
|
|
if mask.sum() == 0:
|
|
return None
|
|
xl, yl, ell = xl[mask], yl[mask], ell[mask]
|
|
cr, pr = solve_for_locations(psfi[cc[i]:cc[i+1]], xe[cc[i]:cc[i+1]], ye[cc[i]:cc[i+1]], roll[cc[i]:cc[i+1]], pk[cc[i]:cc[i+1]], xl, yl, ell, psfdata)
|
|
print(pr.min(), pr.max())
|
|
plt.scatter(xl, yl, c=pr)
|
|
plt.scatter(xe[cc[i]: cc[i + 1]], ye[cc[i]: cc[i+1]], marker="x", c="r", s=50, zorder=20)
|
|
plt.show()
|
|
"""
|
|
cr, pr = solve_for_locations(psfi[cc[i]:cc[i+1]], xe[cc[i]:cc[i+1]], ye[cc[i]:cc[i+1]], roll[cc[i]:cc[i+1]], pk[cc[i]:cc[i+1]], xl, yl, ell, psfdata)
|
|
else:
|
|
xl, yl, cr, pr = np.empty(0, int),np.empty(0, int),np.empty(0, float),np.empty(0, float)
|
|
return xl.astype(int), yl.astype(int), cr, pr
|
|
|
|
|
|
"""
|
|
for r in enumerate(xyu.T):
|
|
worker(r)
|
|
|
|
"""
|
|
tpool = ThreadPool(8)
|
|
for xl, yl, cr, pr in tqdm.tqdm(tpool.imap_unordered(worker, enumerate(xyu.T)), total=xyu.shape[1]):
|
|
cmap[yl, xl] = cr
|
|
pmap[yl, xl] = pr
|
|
|
|
"""
|
|
for i, (xs, ys) in enumerate(xyu.T):
|
|
eloc = emap[xs:xs+sizex,ys:ys+sizey]
|
|
mask = eloc > 0.
|
|
xl = (xg[mask] + xs*sizex).astype(float)
|
|
yl = (yg[mask] + ys*sizey).astype(float)
|
|
ell = (eloc[mask]).astype(float)
|
|
print(i, xs, ys, xl.min(), xl.max(), yl.min(), yl.max())
|
|
cr, pr = solve_for_locations(psfi[cc[i]:cc[i+1]], xe[cc[i]:cc[i+1]], ye[cc[i]:cc[i+1]], roll[cc[i]:cc[i+1]], pk[cc[i]:cc[i+1]], xl, yl, ell, psfdata)
|
|
xl, yl = xl.astype(int), yl.astype(int)
|
|
cmap[xl, yl] = cr
|
|
pmap[xl, yl] = pr
|
|
i = 0
|
|
xs, ys = xyu[:, i]
|
|
eloc = emap[xs:xs+sizex,ys:ys+sizey]
|
|
mask = eloc > 0.
|
|
xl = (xg[mask] + xs*sizex).astype(float)
|
|
yl = (yg[mask] + ys*sizey).astype(float)
|
|
ell = (eloc[mask]).astype(float)
|
|
print(i, xs, ys, xl.min(), xl.max(), yl.min(), yl.max())
|
|
cr, pr = solve_for_locations(psfi[cc[i]:cc[i+1]], xe[cc[i]:cc[i+1]], ye[cc[i]:cc[i+1]], roll[cc[i]:cc[i+1]], pk[cc[i]:cc[i+1]], xl, yl, ell, psfdata)
|
|
xl, yl = xl.astype(int), yl.astype(int)
|
|
cmap[xl, yl] = cr
|
|
pmap[xl, yl] = pr
|
|
"""
|
|
return wcs, cmap, pmap
|
|
|
|
|
|
if __name__ == "__main__":
|
|
p1 = fits.open("updated_test.fits")
|
|
emap = fits.getdata("exp.fits.gz") #np.full((8192, 8192), 10000.)
|
|
|
|
wcs, cmap, pmap = make_srccount_and_detmap(emap, p1[1].data, p1[1].header)
|
|
fits.ImageHDU(data=pmap, header=wcs.to_header()).writeto("tmap3.fits.gz", overwrite=True)
|
|
"""
|
|
|
|
if __name__ == "__main__":
|
|
p1 = fits.open("updated_test.fits")
|
|
evt = p1[1].data
|
|
|
|
emap = np.full((8192, 8192), 14139.9)
|
|
psfdata, ui = prepare_psf(evt)
|
|
#psfdata = np.zeros(psfdata.shape, float)
|
|
#psfdata[:, 50, 50] = 1.
|
|
wcs = read_wcs(p1[1].header)
|
|
|
|
sizex, sizey = select_xychunksize(wcs)
|
|
iidx, xyu, cts = create_neighboring_blocks(evt["x"], evt["y"], sizex, sizey)
|
|
print(xyu.shape)
|
|
cc = np.zeros(cts.size + 1, int)
|
|
cc[1:] = np.cumsum(cts)
|
|
cmap, pmap = np.zeros(emap.shape, float), np.zeros(emap.shape, float)
|
|
#xe, ye, pk, roll, psfi = np.copy(evt["x"][iidx]), np.copy(evt["y"][iidx]), np.copy((evt["quant_eff"]/evt["bkg_model"])[iidx]), np.copy(evt["roll_pnt"][iidx]), np.copy(ui[iidx])
|
|
xe, ye, pk, roll, psfi = np.copy(evt["x"][iidx]), np.copy(evt["y"][iidx]), np.copy((2e-3/evt["bkg_model"])[iidx]), np.copy(np.deg2rad(evt["roll_pnt"][iidx])), np.copy(ui[iidx])
|
|
xe = xe.astype(float)
|
|
ye = ye.astype(float)
|
|
roll = roll.astype(float)
|
|
pk = pk.astype(float)
|
|
print("pk.dtype", pk.dtype, "xe.dtype", xe.dtype)
|
|
|
|
#plt.hist(pk, 100)
|
|
#plt.show()
|
|
xg, yg = np.mgrid[0:sizex:1, 0:sizey:1]
|
|
|
|
i = 60# np.argmax(cts)
|
|
xs, ys = xyu[:, i]
|
|
|
|
eloc = emap[xs:xs+sizex,ys:ys+sizey]
|
|
mask = eloc > 0.
|
|
|
|
xl = (xg[mask] + xs*sizex).astype(float)
|
|
yl = (yg[mask] + ys*sizey).astype(float)
|
|
plt.scatter(xe[cc[i]:cc[i+1]], ye[cc[i]: cc[i + 1]], marker="+", color="r")
|
|
plt.scatter(xl, yl, marker="x", color="k", zorder=20)
|
|
plt.show()
|
|
|
|
ell = np.copy(eloc[mask])
|
|
print("check pk in python", pk[[cc[i], cc[i] + 10, cc[i] + 20]])
|
|
print("check pk in python", psfi[[cc[i], cc[i] + 10, cc[i] + 20]])
|
|
print("check pk in python", xe[[cc[i], cc[i] + 10, cc[i] + 20]])
|
|
cr, pr = solve_for_locations(psfi[cc[i]:cc[i+1]], xe[cc[i]:cc[i+1]], ye[cc[i]:cc[i+1]], roll[cc[i]:cc[i+1]], pk[cc[i]:cc[i+1]], xl, yl, ell, psfdata)
|
|
plt.scatter(xl, yl, c=pr)
|
|
plt.show()
|
|
|
|
|
|
"""
|