Skip to content
PSFMatsREE80.py 7.27 KiB
Newer Older
Fang Yuedong's avatar
Fang Yuedong committed
import sys
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
#import xlrd
from scipy import ndimage

sys.path.append("/public/home/weichengliang/lnData/CSST_new_framwork/csstPSF_20210108")
import PSFConfig as myConfig
import PSFUtil as myUtil

NPSF = 900
##############################
##############################
##############################
def test_psfREE80(psfPath, ThisTask, NTasks):
    nccd = 30
    npsf = NPSF

    npsfPerTasks = int(npsf/NTasks)
    iStart= 0 + npsfPerTasks*ThisTask
    iEnd  = npsfPerTasks + npsfPerTasks*ThisTask
    if ThisTask == NTasks:
        iEnd = npsf

    CENPIXUSED  = True
    wvREE80 = np.zeros([4, nccd])  #psf in different waves-4
    ttREE80 = np.zeros(nccd)       #stacked psf

    for iccd in range(1, nccd+1):
        psf_wvREE80 = np.zeros([4, npsf])
        psf_ttREE80 = np.zeros(npsf)

        #for ipsf in range(1, npsf+1):
        for ipsf in range(iStart+1, iEnd+1):
            psf4iwave = []
            for iwave in range(1, 5):        
                if ThisTask == 0:
                    print('iccd-ipsf-iwave: {:} {:} {:}'.format(iccd, ipsf, iwave), end='\r')
                psfInfo = myConfig.LoadPSF(iccd, iwave, ipsf, psfPath, InputMaxPixelPos=True, PSFCentroidWgt=False)
                
                cenPix = None
                if CENPIXUSED:
                    psfInfoX= myConfig.LoadPSF(iccd, iwave, ipsf, psfPath, InputMaxPixelPos=True, PSFCentroidWgt=True)
                    deltX= psfInfoX['centroid_x'] #in mm
                    deltY= psfInfoX['centroid_y'] #in mm
                    pixsize  = 2.5*1e-3 #mm, will use binningPSF
                    cenPix_X = 512/2 + deltX/pixsize
                    cenPix_Y = 512/2 + deltY/pixsize
                    cenPix = [cenPix_X, cenPix_Y]

                ipsfMat = psfInfo['psfMat']
                cenX, cenY, sz, e1, e2, REE80 = myUtil.psfSizeCalculator(ipsfMat, CalcPSFcenter=True, SigRange=True, TailorScheme=2, cenPix=cenPix)
                psf_wvREE80[iwave-1, ipsf-1] = REE80
                psf4iwave.append(ipsfMat)
            tt = myUtil.psfStack(psf4iwave[0], psf4iwave[1], psf4iwave[2], psf4iwave[3])
            cenX, cenY, sz, e1, e2, REE80 = myUtil.psfSizeCalculator(tt, CalcPSFcenter=True, SigRange=True, TailorScheme=2)
            psf_ttREE80[ipsf-1] = REE80

        if iccd == 1 and iwave ==1:
            print('iccd-{:}:'.format(iccd), flush=True)
            print('psfSet has been loaded.', flush=True)
            #print('Usage: psfSet[i][keys]', flush=True)
            #print('psfSet.keys:', psfSet[0].keys(), flush=True)
        else:
            print('iccd-{:}, iwave-{:}'.format(iccd, iwave), end='\r', flush=True)
    
        comm.barrier()
        psf_ttREE80 = comm.allreduce(psf_ttREE80, op=MPI.SUM)
        psf_wvREE80[0, :] = comm.allreduce(psf_wvREE80[0, :], op=MPI.SUM)
        psf_wvREE80[1, :] = comm.allreduce(psf_wvREE80[1, :], op=MPI.SUM)
        psf_wvREE80[2, :] = comm.allreduce(psf_wvREE80[2, :], op=MPI.SUM)
        psf_wvREE80[3, :] = comm.allreduce(psf_wvREE80[3, :], op=MPI.SUM)

        ttREE80[iccd-1]    = np.mean(psf_ttREE80)
        wvREE80[0, iccd-1] = np.mean(psf_wvREE80[0, :])
        wvREE80[1, iccd-1] = np.mean(psf_wvREE80[1, :])
        wvREE80[2, iccd-1] = np.mean(psf_wvREE80[2, :])
        wvREE80[3, iccd-1] = np.mean(psf_wvREE80[3, :])
        ##############################

    comm.barrier()
    #ttREE80 = comm.allreduce(ttREE80, op=MPI.SUM)
    #wvREE80 = comm.allreduce(wvREE80, op=MPI.SUM)

    #plot-test
    if ThisTask == 0:
        REE80W1 = wvREE80[0, :]
        REE80W2 = wvREE80[1, :]
        REE80W3 = wvREE80[2, :]
        REE80W4 = wvREE80[3, :]

        np.savetxt('REE80_w1.txt',REE80W1)
        np.savetxt('REE80_w2.txt',REE80W2)
        np.savetxt('REE80_w3.txt',REE80W3)
        np.savetxt('REE80_w4.txt',REE80W4)
        np.savetxt('REE80_tt.txt',ttREE80)


        ccdFilterLayout = ['GV', 'GV', 'GU', 'GU', 'GI', 'y', 'i', 'g', 'r', 'GI', 'z', 'NUV', 'NUV', 'u', 'y', 'y','u', 'NUV', 'NUV', 'z', 'GI', 'r', 'g', 'i', 'y', 'GI', 'GU', 'GU','GV', 'GV']

        fig = plt.figure(figsize=(18,10))
        for iccd in range(0,30):
            plt.arrow(iccd+1, REE80W1[iccd], 0, REE80W4[iccd]-REE80W1[iccd], width = 0.05, head_length=0.002, ec='None', color='k')
            plt.plot([iccd+1], [REE80W1[iccd]], 'o',c='k')
            plt.plot([iccd+1.1], [REE80W2[iccd]], 'o',c='b')
            plt.plot([iccd+1.2], [REE80W3[iccd]], 'o',c='g')
            plt.plot([iccd+1.3], [REE80W4[iccd]], 'o',c='r')
            plt.plot([iccd+1, iccd+1.1, iccd+1.2, iccd+1.3], [REE80W1[iccd], REE80W2[iccd], REE80W3[iccd], REE80W4[iccd]], '--',c='k')
            if REE80W1[iccd] < REE80W4[iccd]:
                plt.text(iccd+1-0.2, REE80W1[iccd]-0.005, ccdFilterLayout[iccd], fontsize=15)
            if REE80W1[iccd] > REE80W4[iccd]:
                plt.text(iccd+1-0.2, REE80W1[iccd]+0.003, ccdFilterLayout[iccd], fontsize=15)

        plt.fill_betweenx([0.078, 0.145], [0.5,0.5], [5.5,5.5], color='gray',alpha=0.5)
        plt.fill_betweenx([0.078, 0.145], [25.5,25.5], [30.5,30.5], color='gray',alpha=0.5)

        plt.fill_betweenx([0.078, 0.145], [9.5,9.5], [10.5,10.5], color='gray',alpha=0.5)
        plt.fill_betweenx([0.078, 0.145], [20.5,20.5], [21.5,21.5], color='gray',alpha=0.5)

        plt.plot([5.5, 5.5], [0.078, 0.5], ':')
        plt.plot([10.5, 10.5], [0.078, 0.5], 'k:')
        plt.plot([15.5, 15.5], [0.078, 0.5], 'k:')
        plt.plot([20.5, 20.5], [0.078, 0.5], 'k:')
        plt.plot([25.5, 25.5], [0.078, 0.5], 'k:')


        plt.ylim(0.078, 0.145)
        plt.xlim(0.5, 30.5)
        #plt.plot(np.linspace(1, 30, 30), REE80W1)
        #plt.plot(np.linspace(1, 30, 30), REE80W2)
        #plt.plot(np.linspace(1, 30, 30), REE80W3)
        #plt.plot(np.linspace(1, 30, 30), REE80W4)

        plt.xticks([])
        plt.yticks(fontsize=15)
        plt.text(1.5, 0.074, 'CCD1 - CCD5', fontsize=15)
        plt.text(6.5, 0.074, 'CCD6 - CCD10', fontsize=15)
        plt.text(11.5, 0.074, 'CCD11 - CCD15', fontsize=15)
        plt.text(16.5, 0.074, 'CCD16 - CCD20', fontsize=15)
        plt.text(21.5, 0.074, 'CCD21 - CCD25', fontsize=15)
        plt.text(26.5, 0.074, 'CCD26 - CCD30', fontsize=15)

        plt.plot([27], [0.183], 'ko')
        plt.text(27.5, 0.182, 'wave-1',fontsize=15)
        plt.plot([27], [0.180], 'ro')
        plt.text(27.5, 0.179, 'wave-2',fontsize=15)
        plt.plot([27], [0.177], 'go')
        plt.text(27.5, 0.176, 'wave-3',fontsize=15)
        plt.plot([27], [0.174], 'bo')
        plt.text(27.5, 0.173, 'wave-4',fontsize=15)


        #overplot stackedPSF
        xccd = np.linspace(1, 30, 30)
        plt.plot(xccd,ttREE80, 'm*', ms = 20, markerfacecolor='None', markeredgewidth=2)
        plt.plot([27], [0.168], 'm*', ms = 20, markerfacecolor='None', markeredgewidth=2)
        plt.text(27.5, 0.1665, 'stacked',fontsize=20)
        plt.savefig('figs/psfStackedREE80.pdf')    



##############################
##############################
##############################
if __name__=='__main__':
    comm = MPI.COMM_WORLD
    ThisTask = comm.Get_rank()
    NTasks   = comm.Get_size()

    psfPath = '/data/simudata/CSSOSDataProductsSims/data/csstPSFdata/CSSOS_psf_20210108/CSST_psf_ciomp_2p5um_cycle3'
    test_psfREE80(psfPath, ThisTask, NTasks)