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() """