import unittest import sys,os,math from itertools import islice import mpi4py.MPI as MPI import numpy as np import matplotlib.pyplot as plt import matplotlib as mpl mpl.use('Agg') import scipy.io from scipy import ndimage #sys.path.append("/public/home/weichengliang/lnData/CSST_new_framwork/csstPSF_v4_20210326") # sys.path.append("../") from ObservationSim.PSF.PSFInterp import PSFConfig as myConfig from ObservationSim.PSF.PSFInterp import PSFUtil as myUtil NPSF = 400 ############################## ###计算PSF椭率### def psfSecondMoments(psfMat, cenX, cenY, pixSize=1): apr = 0.5 #arcsec, 0.5角秒内测量 fl = 28. #meters pxs = 2.5 #microns apr = np.deg2rad(apr/3600.)*fl*1e6 apr = apr/pxs apr = np.int(np.ceil(apr)) I = psfMat ncol = I.shape[1] nrow = I.shape[0] w = 0.0 w11 = 0.0 w12 = 0.0 w22 = 0.0 for icol in range(ncol): for jrow in range(nrow): x = icol*pixSize - cenX y = jrow*pixSize - cenY rr = np.sqrt(x*x + y*y) wgt= 0.0 if rr <= apr: wgt = 1.0 w += I[jrow, icol]*wgt w11 += x*x*I[jrow, icol]*wgt w12 += x*y*I[jrow, icol]*wgt w22 += y*y*I[jrow, icol]*wgt w11 /= w w12 /= w w22 /= w sz = w11 + w22 e1 = (w11 - w22)/sz e2 = 2.0*w12/sz return sz, e1, e2 ############################## ############################## ''' def assignTasks(npsf, NTasks, ThisTask): npsfPerTasks = int(npsf/NTasks) iStart= 0 + npsfPerTasks*ThisTask iEnd = npsfPerTasks + npsfPerTasks*ThisTask if ThisTask == NTasks: iEnd = npsf return iStart, iEnd ''' #def test_psfIDW(iccd, iwave, psfPath, ThisTask, NTasks): def test_psfIDW(iccd, iwave, psfPath): nccd = 30 npsfA = 900 npsfB = 400 #iStart, iEnd = assignTasks(400, NTasks, ThisTask) psfPathA = psfPath+'_30x30field' psfPathB = psfPath+'_20x20field' imxA = np.zeros(npsfA) imyA = np.zeros(npsfA) psfA = np.zeros([npsfA, 512, 512]) imxB = np.zeros(npsfB) imyB = np.zeros(npsfB) psfB = np.zeros([npsfB, 512, 512]) #for ipsf in range(iStart+1, iEnd+1): for ipsf in range(1, npsfA+1): print('ipsfA:', ipsf, end='\r', flush=True) psfInfo = myConfig.LoadPSF(iccd, iwave, ipsf, psfPathA, InputMaxPixelPos=True, PSFCentroidWgt=True, VPSF=False) imxA[ipsf-1] = psfInfo['image_x'] + psfInfo['centroid_x'] imyA[ipsf-1] = psfInfo['image_y'] + psfInfo['centroid_y'] psfA[ipsf-1, :, :] = psfInfo['psfMat'] for ipsf in range(1, npsfB+1): print('ipsfB:', ipsf, end='\r', flush=True) psfInfo = myConfig.LoadPSF(iccd, iwave, ipsf, psfPathB, InputMaxPixelPos=True, PSFCentroidWgt=True, VPSF=False) imxB[ipsf-1] = psfInfo['image_x'] + psfInfo['centroid_x'] imyB[ipsf-1] = psfInfo['image_y'] + psfInfo['centroid_y'] psfB[ipsf-1, :, :] = psfInfo['psfMat'] #myConfig.psfMaker_IDW(px, py, PSFMat, cen_col, cen_row, IDWindex=2, OnlyNeighbors=True, hoc=None, hoclist=None, PSFCentroidWgt=False) for ipsf in range(npsfB): print('ipsf:', ipsf, end='\r', flush=True) px = imxB[ipsf] py = imyB[ipsf] cen_col = imxA cen_row = imyA PSFMat = psfA psfIDW = myConfig.psfMaker_IDW(px, py, PSFMat, cen_col, cen_row, IDWindex=2, OnlyNeighbors=True, hoc=None, hoclist=None, PSFCentroidWgt=False) np.save('figs/psfIDW/psfIDW_{:}_{:}_{:}'.format(iccd, iwave, ipsf+1), psfIDW) #def test_psfEll(iccd, iwave, ThisTask, NTasks): def test_psfEll(iccd, iwave): nccd = 30 npsf = 400 #iStart, iEnd = assignTasks(npsf, NTasks, ThisTask) #imx = np.zeros(npsf) #imy = np.zeros(npsf) psf_e1 = np.zeros(npsf) psf_e2 = np.zeros(npsf) psf_sz = np.zeros(npsf) #for ipsf in range(iStart+1, iEnd+1): for ipsf in range(1, 401): if ipsf > 1: continue print('ipsf-{:}'.format(ipsf), end='\r') #psfInfo = myConfig.LoadPSF(iccd, iwave, ipsf, psfPath, InputMaxPixelPos=True, PSFCentroidWgt=True, VPSF=False) #psfMat = psfInfo['psfMat'] #imx[ipsf-1] = psfInfo['image_x']+psfInfo['centroid_x'] #imy[ipsf-1] = psfInfo['image_y']+psfInfo['centroid_y'] psfMat = np.load('figs/psfIDW/psfIDW_{:}_{:}_{:}.npy'.format(iccd, iwave, ipsf)) #psfInfo['psfMat'] cenX = 256 cenY = 256 sz, e1, e2 = psfSecondMoments(psfMat, cenX, cenY, pixSize=1) psf_e1[ipsf-1] = e1 psf_e2[ipsf-1] = e2 psf_sz[ipsf-1] = sz ####### #comm.barrier() #imx = comm.allreduce(imx, op=MPI.SUM) #imy = comm.allreduce(imy, op=MPI.SUM) #psf_e1 = comm.allreduce(psf_e1, op=MPI.SUM) #psf_e2 = comm.allreduce(psf_e2, op=MPI.SUM) #psf_sz = comm.allreduce(psf_sz, op=MPI.SUM) #comm.barrier() #if ThisTask == 0: # arr = [psf_e1, psf_e2, psf_sz] # np.save('/public/home/weichengliang/lnData/CSST_new_framwork/csstPSF_v4_20210326/test4report/data/psfEll20IDW_{:}_{:}'.format(iccd, iwave), arr) ''' def test_psfResidualCalc(iccd, iwave, ipsf, psfPath): psfInfo = myConfig.LoadPSF(iccd, iwave, ipsf, psfPath+'_20x20field', InputMaxPixelPos=True, PSFCentroidWgt=True, VPSF=False) psfMatORG = psfInfo['psfMat'] psfMatIDW = np.load('/public/home/weichengliang/lnData/CSST_new_framwork/csstPSF_v4_20210326/test4report/figs/psfIDW/psfIDW_{:}_{:}_{:}.npy'.format(iccd, iwave, ipsf)) _,tREE80 = myUtil.psfEncircle(psfMatORG, fraction=0.8, psfSampleSizeInMicrons=2.5, focalLengthInMeters=28, cenPix=None) tREE80_pix = np.int32(np.ceil(tREE80/(0.074/2/2))[0]) #print(tREE80, np.ceil(tREE80/(0.074/2/2)), tREE80_pix) timg0 = psfMatORG[256-tREE80_pix:256+tREE80_pix, 256-tREE80_pix:256+tREE80_pix] timg1 = psfMatIDW[256-tREE80_pix:256+tREE80_pix, 256-tREE80_pix:256+tREE80_pix] #print("residual::", np.max((timg1-timg0)/timg0), np.min((timg1-timg0)/timg0), np.mean((timg1-timg0)/timg0)) return np.mean((timg1-timg0)/timg0) ''' def test_psfResidualPlot(iccd, iwave, ipsf, psfPath): psfInfo = myConfig.LoadPSF(iccd, iwave, ipsf, psfPath+'_20x20field', InputMaxPixelPos=True, PSFCentroidWgt=True, VPSF=False) psfMatORG = psfInfo['psfMat'] psfMatIDW = np.load('figs/psfIDW/psfIDW_{:}_{:}_{:}.npy'.format(iccd, iwave, ipsf)) npix = psfMatORG.shape[0] pixCutEdge= int(npix/2-15) img0 = psfMatORG[pixCutEdge:npix-pixCutEdge, pixCutEdge:npix-pixCutEdge] img1 = psfMatIDW[pixCutEdge:npix-pixCutEdge, pixCutEdge:npix-pixCutEdge] imgX = (img1 - img0)/img0 img0 = np.log10(img0) img1 = np.log10(img1) imgX = np.log10(np.abs(imgX)) fig = plt.figure(figsize=(18,4)) ax = plt.subplot(1,3,1) plt.imshow(img0, origin='lower', vmin=-7, vmax=-1.3) plt.plot([npix/2-pixCutEdge, npix/2-pixCutEdge],[0, (npix/2-pixCutEdge)*2-1],'w--') plt.plot([0, (npix/2-pixCutEdge)*2-1],[npix/2-pixCutEdge, npix/2-pixCutEdge],'w--') plt.annotate('ORG', [0,(npix/2-pixCutEdge)*2-5], c='w', size=15) cticks=[-7, -6, -5, -4, -3, -2, -1] cbar = plt.colorbar(ticks=cticks) cbar.ax.set_yticklabels(['$10^{-7}$', '$10^{-6}$', '$10^{-5}$','$10^{-4}$','$10^{-3}$','$10^{-2}$', '$10^{-1}$']) print(img0.min(), img0.max()) ax = plt.subplot(1,3,2) plt.imshow(img1, origin='lower', vmin=-7, vmax=-1.3) plt.plot([npix/2-pixCutEdge, npix/2-pixCutEdge],[0, (npix/2-pixCutEdge)*2-1],'w--') plt.plot([0, (npix/2-pixCutEdge)*2-1],[npix/2-pixCutEdge, npix/2-pixCutEdge],'w--') plt.annotate('IDW', [0,(npix/2-pixCutEdge)*2-5], c='w', size=15) cticks=[-7, -6, -5, -4, -3, -2, -1] cbar = plt.colorbar(ticks=cticks) cbar.ax.set_yticklabels(['$10^{-7}$', '$10^{-6}$', '$10^{-5}$','$10^{-4}$','$10^{-3}$','$10^{-2}$', '$10^{-1}$']) print(img1.min(), img1.max()) ax = plt.subplot(1,3,3) plt.imshow(imgX, origin='lower', vmin =-3, vmax =np.log10(3e-1)) plt.plot([npix/2-pixCutEdge, npix/2-pixCutEdge],[0, (npix/2-pixCutEdge)*2-1],'w--') plt.plot([0, (npix/2-pixCutEdge)*2-1],[npix/2-pixCutEdge, npix/2-pixCutEdge],'w--') #plt.annotate('(IDW-ORG)/ORG', [0,(npix/2-pixCutEdge)*2-5], c='w', size=15) cticks=[-5, -4, -3, -2, -1] cbar = plt.colorbar(ticks=cticks) cbar.ax.set_yticklabels(['$10^{-5}$','$10^{-4}$','$10^{-3}$','$10^{-2}$', '$10^{-1}$']) print(np.max((psfMatORG-psfMatIDW))) plt.savefig('figs/psfResidual_iccd{:}.pdf'.format(iccd)) def test_psfEllPlot(OVERPLOT=False): #if ThisTask == 0: if True: prefix = 'psfEll20' iccd = 1 iwave= 1 data = np.load('data/'+prefix+'_1_1.npy') imx= data[0] imy= data[1] psf_e1 = data[2] psf_e2 = data[3] print(np.shape(imx)) npsf = np.shape(imx)[0] plt.cla() plt.close("all") fig = plt.figure(figsize=(12,12)) plt.plot(imx, imy, 'r.') plt.savefig('figs/psfPos.pdf') ####### fig = plt.figure(figsize=(12, 12)) plt.subplots_adjust(wspace=0.1, hspace=0.1) ax = plt.subplot(1, 1, 1) for ipsf in range(npsf): plt.plot(imx[ipsf], imy[ipsf], 'b.') ang = np.arctan2(psf_e2[ipsf], psf_e1[ipsf])/2 ell = np.sqrt(psf_e1[ipsf]**2 + psf_e2[ipsf]**2) ell *= 15 lcos = ell*np.cos(ang) lsin = ell*np.sin(ang) plt.plot([imx[ipsf]-lcos, imx[ipsf]+lcos],[imy[ipsf]-lsin, imy[ipsf]+lsin],'b', lw=2) ########### ang = 0. ell = 0.05 ell*= 15 lcos = ell*np.cos(ang) lsin = ell*np.sin(ang) #plt.plot([imx[898]-lcos, imx[898]+lcos],[imy[898]+5.-lsin, imy[898]+5.+lsin],'k', lw=2) #plt.annotate('{:}'.format(ell/15), (imx[898]-2., imy[898]+6.), xycoords='data', fontsize=10) plt.xlabel('CCD X (mm)') plt.ylabel('CCD Y (mm)') if OVERPLOT == True: prefix = 'psfEll20IDW' data = np.load('data/'+prefix+'_1_1.npy') #imx= data[0] #imy= data[1] psf_e1 = data[0] psf_e2 = data[1] npsf = np.shape(imx)[0] for ipsf in range(npsf): #plt.plot(imx[ipsf], imy[ipsf], 'r.') ang = np.arctan2(psf_e2[ipsf], psf_e1[ipsf])/2 ell = np.sqrt(psf_e1[ipsf]**2 + psf_e2[ipsf]**2) ell *= 15 lcos = ell*np.cos(ang) lsin = ell*np.sin(ang) plt.plot([imx[ipsf]-lcos, imx[ipsf]+lcos],[imy[ipsf]-lsin, imy[ipsf]+lsin],'r', lw=1) plt.gca().set_aspect(1) if OVERPLOT == True: prefix = 'psfEllOPIDW' plt.savefig('figs/'+prefix+'_iccd{:}.pdf'.format(iccd)) ''' def test_psfdEllPlot(): if ThisTask == 0: prefix = 'psfEll20' iccd = 1 iwave= 1 data = np.load('data/'+prefix+'_1_1.npy') imx= data[0] imy= data[1] psf_e1 = data[2] psf_e2 = data[3] psf_sz = data[4] print(np.shape(imx)) npsf = np.shape(imx)[0] ellX = np.sqrt(psf_e1**2 + psf_e2**2) angX = np.arctan2(psf_e2, psf_e1)/2 angX = np.rad2deg(angX) szX = psf_sz ############################## prefix = 'psfEll20IDW' data = np.load('data/'+prefix+'_1_1.npy') #imx= data[0] #imy= data[1] psf_e1 = data[0] psf_e2 = data[1] psf_sz = data[2] ellY = np.sqrt(psf_e1**2 + psf_e2**2) angY = np.arctan2(psf_e2, psf_e1)/2 angY = np.rad2deg(angY) szY = psf_sz ############################## fig=plt.figure(figsize=(15, 4)) ax = plt.subplot(1,3,1) plt.hist(ellX, bins=20, color='b', alpha=0.5) plt.hist(ellY, bins=20, color='r', alpha=0.5) plt.xlabel('$\epsilon$') plt.ylabel('PDF') ax = plt.subplot(1,3,2) plt.hist((ellY-ellX)/ellX, bins=20, color='r', alpha=0.5) plt.xlabel('$(\epsilon_{\\rm IDW}-\epsilon_{\\rm ORG})/\epsilon_{\\rm ORG}$') plt.ylabel('PDF') ax = plt.subplot(1,3,3) plt.hist((angY-angX)/angX, bins=20, color='r', alpha=0.5, range=[-0.1, 0.1]) plt.xlabel('$(\\alpha_{\\rm IDW}-\\alpha_{\\rm ORG})/\\alpha_{\\rm ORG}$') plt.ylabel('PDF') plt.savefig('figs/psfEllOPIDWPDF.pdf') fig=plt.figure(figsize=(4, 4)) plt.hist((szY-szX)/szX, bins=20, color='r', alpha=0.5) plt.xlabel('$(R_{\\rm IDW}-R_{\\rm ORG})/R_{\\rm ORG}$') plt.ylabel('PDF') plt.savefig('figs/psfEllOPIDWPDF_dsz.pdf') ''' def test_psfdEllabsPlot(iccd): #if ThisTask == 0: if True: prefix = 'psfEll20' #iccd = 1 #iwave= 1 data = np.load('data/'+prefix+'_{:}_1.npy'.format(iccd)) imx= data[0] imy= data[1] psf_e1 = data[2] psf_e2 = data[3] psf_sz = data[4] print(np.shape(imx)) npsf = np.shape(imx)[0] ellX = np.sqrt(psf_e1**2 + psf_e2**2) angX = np.arctan2(psf_e2, psf_e1)/2 angX = np.rad2deg(angX) szX = psf_sz ############################## prefix = 'psfEll20IDW' data = np.load('data/'+prefix+'_{:}_1.npy'.format(iccd)) #imx= data[0] #imy= data[1] psf_e1 = data[0] psf_e2 = data[1] psf_sz = data[2] ellY = np.sqrt(psf_e1**2 + psf_e2**2) angY = np.arctan2(psf_e2, psf_e1)/2 angY = np.rad2deg(angY) szY = psf_sz ############################## fig=plt.figure(figsize=(6, 5)) grid = plt.GridSpec(3,1,left=0.15, right=0.95, wspace=None, hspace=0.02) #plt.subplots_adjust(left=None,bottom=None,right=None,top=None,wspace=None,hspace=0.02) ax = plt.subplot(grid[0:2,0]) plt.plot([0.01,0.1],[0.01,0.1], 'k--', lw=1. ) plt.scatter(ellX, ellY, color='b', alpha=1., s=3., edgecolors='None') plt.xlim([0.015, 0.085]) plt.ylim([0.015, 0.085]) plt.ylabel('$\epsilon_{\\rm IDW}$') plt.gca().axes.get_xaxis().set_visible(False) ax = plt.subplot(grid[2,0]) plt.plot([0.015,0.085],[0.,0.], 'k--', lw=1. ) plt.scatter(ellX, (ellY-ellX), color='b', s=3., edgecolors='None') plt.xlim([0.015, 0.085]) plt.ylim([-0.0018, 0.0018]) plt.xlabel('$\epsilon_{\\rm ORG}$') plt.ylabel('$\Delta$') plt.savefig('figs/psfEllOPIDWPDF_{:}.pdf'.format(iccd)) fig=plt.figure(figsize=(4, 4)) plt.hist((szY-szX)/szX, bins=20, color='r', alpha=0.5) plt.xlabel('$(R_{\\rm IDW}-R_{\\rm ORG})/R_{\\rm ORG}$') plt.ylabel('PDF') plt.savefig('figs/psfEllOPIDWPDF_dsz_{:}.pdf'.format(iccd)) class PSFMatsIDW_coverage(unittest.TestCase): def test_psfIDW_(self): #comm = MPI.COMM_WORLD #ThisTask = comm.Get_rank() #NTasks = comm.Get_size() iccd = 1 iwave= 1 ipsf = 400 psfPath = '/data/simudata/CSSOSDataProductsSims/data/csstPSFdata/CSSOS_psf_20210326/CSST_psf_ciomp' #test_psfIDW(iccd, iwave, psfPath, ThisTask, NTasks) test_psfIDW(iccd, iwave, psfPath) ''' for iccd in range(7, 10): res = np.zeros(400) for ipsf in range(1,401): print(ipsf, end="\r") res[ipsf-1] = test_psfResidualCalc(iccd, iwave, ipsf, psfPath) #fig = plt.figure(figsize=(6,6)) #plt.hist(np.abs(res), bins=50) #plt.xlim([0,1]) #plt.savefig('figs/psfResidualREE80PDF.pdf') print("{:}:".format(iccd), res[res<=0.01].size/400*100) ''' test_psfResidualPlot(iccd, iwave, ipsf, psfPath) #test_psfEll(iccd, iwave, ThisTask, NTasks) test_psfEll(iccd, iwave) test_psfEllPlot(OVERPLOT=True) #test_psfdEllPlot() test_psfdEllabsPlot(iccd) ############################## ############################## ############################## if __name__=='__main__': unittest.main()