psfConfigTest.py

"""
CSST image simulation module (in python3): Point Spread Function (PSF)
author:: Wei Chengliang <chengliangwei@pmo.ac.cn>
"""

import sys
from itertools import islice

import numpy as np
import matplotlib.pyplot as plt

import scipy.io
from scipy.io import loadmat
#import xlrd

from scipy import ndimage
from scipy.interpolate import RectBivariateSpline

#from astropy.modeling.models import Ellipse2D
#from astropy.coordinates import Angle
#import matplotlib.patches as mpatches

import ctypes
import galsim



def setupPSFimg(iccd, iwave, psfPath="/data/simudata/CSSOSDataProductsSims/data/csstPSFdata/CSSOS_psf_ciomp"):
    """
    psf model setup for csst-sim
   
    Parameters:
        iccd, iwave (int, int): psf model on iccd & iwave
        psfPath (string, optional): path to psf matrix

    Returns:
        psf_model (psf_class): psf model
  
    Methods:
        psf_model.PSFinplace(self, px, py, interpScheme=1): psf interpolation
        psf_model.PSFspin(self, psf, sigSpin, sigGauss, dx, dy): psf rotation (from Yudong)
    """
    psf_model = PSFimg(iccd, iwave, psfPath)
    return psf_model


##################################################
#       A. psf matrix loading & checking         #
##################################################
def psfPixelLayout(nrows, ncols, cenPosRow, cenPosCol, pixSizeInMicrons=5.0):
    """
    convert psf pixels to physical position
    
    Parameters:
        nrows, ncols (int, int): psf sampling with [nrows, ncols].
        cenPosRow, cenPosCol (float, float): A physical position of the chief ray for a given psf.
        pixSizeInMicrons (float-optional): The pixel size in microns from the psf sampling.
        
    Returns:
        psfPixelPos (numpy.array-float): [posx, posy] in mm for [irow, icol]
 
    Notes:
        1. show positions on ccd, but not position on image only [+/- dy]
    """
    psfPixelPos = np.zeros([2, nrows, ncols])
    if nrows % 2 != 0:
        sys.exit()
    if ncols % 2 != 0:
        sys.exit()
        
    cenPix_row = nrows/2 + 1 #中心主光线对应pixle [由长光定义]
    cenPix_col = ncols/2 + 1

    for irow in range(nrows):
        for icol in range(ncols):
            delta_row = ((irow + 1) - cenPix_row)*pixSizeInMicrons*1e-3
            delta_col = ((icol + 1) - cenPix_col)*pixSizeInMicrons*1e-3
            psfPixelPos[0, irow, icol] = cenPosCol + delta_col
            psfPixelPos[1, irow, icol] = cenPosRow - delta_row  #note-1
            
    return psfPixelPos


def imSigRange(img, fraction=0.80):
    """
    extract the image within x-percent (DISCARD)
    
    Parameters:
        img (numpy.array-float): image
        fraction (float-optional): a percentage

    Returns:
        im1 (numpy.array-float): image
    """
    im1 = img.copy()
    im1size = im1.shape
    im2 = np.sort(im1.reshape(im1size[0]*im1size[1]))
    im2 = im2[::-1]
    im3 = np.cumsum(im2)/np.sum(im2)
    loc = np.where(im3 > fraction)
    #print(im3[loc[0][0]], im2[loc[0][0]])
    im1[np.where(im1 <= im2[loc[0][0]])]=0

    return im1


def imPlotInRange(img):
    """
    plot image within a selected range
    
    Parameters:
        img (numpy.array-float): image

    Returns:
    """
    im1 = img.copy()
    im1size = im1.shape
    X,Y = np.meshgrid(range(im1size[1]),range(im1size[0]))
    Z = im1

    resolution = 25

    f = lambda x,y: Z[int(y),int(x) ]
    g = np.vectorize(f)

    x = np.linspace(0,Z.shape[1], Z.shape[1]*resolution)
    y = np.linspace(0,Z.shape[0], Z.shape[0]*resolution)
    X2, Y2= np.meshgrid(x[:-1],y[:-1])
    Z2 = g(X2,Y2)

    #plt.pcolormesh(X,Y, Z)
    #plt.imshow(img, origin='lower')
    plt.contour(X2-0.5,Y2-0.5,Z2, [0.], colors='w', linestyles='--',  linewidths=[1])

    return 


def findMaxPix(img):
    """
    get the pixel position of the maximum-value
    
    Parameters:
        img (numpy.array-float): image
    
    Returns:
        imgMaxPix_x, imgMaxPix_y (int, int): pixel position in columns & rows
    """
    maxIndx = np.argmax(img)
    maxIndx = np.unravel_index(maxIndx, np.array(img).shape)
    imgMaxPix_x = maxIndx[1]
    imgMaxPix_y = maxIndx[0]

    return imgMaxPix_x, imgMaxPix_y


def psfTailor(img, apSizeInArcsec=0.5, psfSampleSizeInMicrons=5, focalLengthInMeters=28):
    """
    psf tailor within a given aperture size
 
    Parameters:
        img (numpy.array-float): image
        apSizeInArcsec (float-optional): aperture size in arcseconds.
        psfSampleSizeInMicrons (float-optional): psf pixel size in microns.
        focalLengthInMeters (float-optional): csst focal length im meters.
    Returns:
        imgT (numpy.array-float): image
    """
    imgMaxPix_x, imgMaxPix_y = findMaxPix(img)
    apSizeInMicrons = np.deg2rad(apSizeInArcsec/3600.)*focalLengthInMeters*1e6
    apSizeInPix = apSizeInMicrons/psfSampleSizeInMicrons
    apSizeInPix = np.int(np.ceil(apSizeInPix))
    imgT = np.zeros_like(img)
    imgT[imgMaxPix_y-apSizeInPix:imgMaxPix_y+apSizeInPix+1, 
         imgMaxPix_x-apSizeInPix:imgMaxPix_x+apSizeInPix+1] = \
    img[imgMaxPix_y-apSizeInPix:imgMaxPix_y+apSizeInPix+1, 
        imgMaxPix_x-apSizeInPix:imgMaxPix_x+apSizeInPix+1]    
    return imgT


def psfEncircle(img, fraction=0.8, psfSampleSizeInMicrons=5, focalLengthInMeters=28):
    """
    psf tailor within a given percentage.
 
    Parameters:
        img (numpy.array-float): image
        fraction (float-optional): a percentage for psf tailor.
        psfSampleSizeInMicrons (float-optional): psf pixel size in microns.
        focalLengthInMeters (float-optional): csst focal length im meters.
    Returns:
        img*wgt (numpy.array-float): image
        REE80 (float): radius of REE80 in arcseconds.
    """
    imgMaxPix_x, imgMaxPix_y = findMaxPix(img)
    im1 = img.copy()
    im1size = im1.shape
    
    dis = np.zeros_like(img)
    for irow in range(im1size[0]):
        for icol in range(im1size[1]):
            dx = icol - imgMaxPix_x