PSFMatsREE80.py

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)