psfMap.py

 
 
import numpy as np
import astropy.io.fits as fits
import matplotlib.pyplot as plt
 
 
class PSF_map:
 
    def __init__(self, SR=np.zeros((0, 0)), radius=0, wl=0, NGS=None, LGS=None, TS=None,
                 AOtype="AO", sup=None, filename=None):
        """ SR_map class
 
        SR      : np.ndarray        : array of strehl (organised by position)
        radius  : float             : radius of the map (in arcsec)
        wl      : float             : wavelength of observation (in nm)
        NGS     : np.ndarray        : position (arcsec) of the NGS (1st line:x axis, 2nd:y axis)
        LGS     : np.ndarray        : position (arcsec) of the LGS (1st line:x axis, 2nd:y axis)
        TS      : np.ndarray        : position (arcsec) of the TS (1st line:x axis, 2nd:y axis)
        sup     : shesha.supervisor : shesha supervisor to retrieve all data from
        filename: str               : filename to read the psf map from
 
        warning:
        using 'sup' will take only the supervisor into consideration
        using 'filename' will overide previous values of the psf map (unless called with sup)
 
        """
        self.mapmap = SR
 
        self._Rtar_Rtar = radius
        self._Rngs_Rngs = 0
        self._Rlgs_Rlgs = 0
        self._Rlgs_Rlgs = 0
        if (NGS is not None):
            self.NGSxNGSx = NGS[0]
            self.NGSyNGSy = NGS[1]
            self._Rngs_Rngs = max((self.NGSxNGSx.max(), self.NGSyNGSy.max()))
        else:
            self.NGSxNGSx = np.zeros((0))
            self.NGSyNGSy = np.zeros((0))
 
        if (LGS is not None):
            self.LGSxLGSx = LGS[0]
            self.LGSyLGSy = LGS[1]
            self._Rlgs_Rlgs = max(self.LGSxLGSx.max(), self.LGSyLGSy.max())
        else:
            self.LGSxLGSx = np.zeros((0))
            self.LGSyLGSy = np.zeros((0))
 
        if (TS is not None):
            self.TSxTSx = TS[0]
            self.TSyTSy = TS[1]
            self._Rts_Rts = max(self.TSxTSx.max(), self.TSyTSy.max())
        else:
            self.TSxTSx = np.zeros((0))
            self.TSyTSy = np.zeros((0))
 
        self.wavelengthwavelength = wl
        self.typetype = AOtype
 
        if (sup is not None):
            self.NGSxNGSx = np.array([t.xpos for t in sup.config.p_wfs_ngs[:-1]])
            self.NGSyNGSy = np.array([t.ypos for t in sup.config.p_wfs_ngs[:-1]])
            self._Rngs_Rngs = max((self.NGSxNGSx.max(), self.NGSyNGSy.max()))
            self.LGSxLGSx = np.array([t.xpos for t in sup.config.p_wfs_lgs])
            self.LGSyLGSy = np.array([t.ypos for t in sup.config.p_wfs_lgs])
            self._Rlgs_Rlgs = max(self.LGSxLGSx.max(), self.LGSyLGSy.max())
            self.TSxTSx = np.array([t.xpos for t in sup.config.p_wfs_ts])
            self.TSyTSy = np.array([t.ypos for t in sup.config.p_wfs_ts])
            self._Rts_Rts = max(self.TSxTSx.max(), self.TSyTSy.max())
            self.wavelengthwavelength = sup.config.p_targets[0].Lambda
            NTAR = len(sup.config.p_targets)
            NTAR_side = int(np.sqrt(NTAR))
            if (NTAR != NTAR_side**2):
                raise ValueError("not a square nb of targets")
            self.mapmap = np.zeros((NTAR_side, NTAR_side))
            for i in range(NTAR):
                #self.map.itemset(i,sup._sim.get_strehl(i)[1])
                self.mapmap.itemset(i, sup._sim.get_strehl(i)[0])
                tar = sup._sim.tar.d_targets[i]
                self._Rtar_Rtar = max(self._Rtar_Rtar, tar.posx, tar.posy)
 
        elif (filename is not None):
            self.readread(filename)
 
    def setNGS(self, NGS, NGSy=None):
        if (NGSy is None):
            self.NGSxNGSx = NGS[0]
            self.NGSyNGSy = NGS[1]
        else:
            self.NGSNGS = NGS
            self.NGSyNGSy = NGS
        self._Rngs_Rngs = max((self.NGSxNGSx.max(), self.NGSyNGSy.max()))
 
    def setLGS(self, LGS, LGSy=None):
        if (LGSy is None):
            self.LGSxLGSx = LGS[0]
            self.LGSyLGSy = LGS[1]
        else:
            self.LGSLGS = LGS
            self.LGSyLGSy = LGS
        self._Rlgs_Rlgs = max(self.LGSxLGSx.max(), self.LGSyLGSy.max())
 
    def setTS(self, TS, TSy=None):
        if (TSy is None):
            self.TSxTSx = TS[0]
            self.TSyTSy = TS[1]
        else:
            self.TSxTSx = TS
            self.TSyTSy = TSy
        self._Rts_Rts = max(self.TSxTSx.max(), self.TSyTSy.max())
 
    def setWaveLength(self, wl):
        self.wavelengthwavelength = wl
 
    def plot(self, title=False, GS=False, WFS=False, LGS=False, NGS=False, TS=False):
        if (self.mapmap.shape[0] > 0 and self.mapmap.shape[1] > 0):
            plt.matshow(self.mapmap,
                        extent=[-self._Rtar_Rtar, self._Rtar_Rtar, -self._Rtar_Rtar, self._Rtar_Rtar])
            plt.colorbar()
            if (GS or WFS or LGS):
                plt.scatter(self.LGSyLGSy, self.LGSxLGSx, color="red")
            if (GS or WFS or NGS):
                plt.scatter(self.NGSyNGSy, self.NGSxNGSx, color="blue")
            if (GS or TS):
                plt.scatter(self.TSyTSy, self.TSxTSx, color="yellow", s=1)
 
            t = self.typetype + " Strehl"
            if (self.wavelengthwavelength > 0):
                t += " @ {:.3f} nm".format(self.wavelengthwavelength)
 
            if (self._Rts_Rts > 0):
                t += ", {:.2f}".format(self._Rts_Rts) + " arcsec ring optim"
            if (title):
                plt.title(t)
 
    def save(self, name=""):
        hdu_map = fits.PrimaryHDU(self.mapmap)
        hdu_map.header["TYPE"] = self.typetype
        hdu_map.header["LAMBDA"] = self.wavelengthwavelength
        hdu_map.header["RTAR"] = self._Rtar_Rtar
        hdu_map.header["RLGS"] = self._Rlgs_Rlgs
        hdu_map.header["RNGS"] = self._Rngs_Rngs
        hdu_map.header["RTS"] = self._Rts_Rts
        hdu_LGSX = fits.ImageHDU(self.LGSxLGSx, name="LGSX")
        hdu_LGSY = fits.ImageHDU(self.LGSyLGSy, name="LGSY")
        hdu_NGSX = fits.ImageHDU(self.NGSxNGSx, name="NGSX")
        hdu_NGSY = fits.ImageHDU(self.NGSyNGSy, name="NGSY")
        hdu_TSX = fits.ImageHDU(self.TSxTSx, name="TSX")
        hdu_TSY = fits.ImageHDU(self.TSyTSy, name="TSY")
 
        hdul = fits.HDUList([
                hdu_map, hdu_LGSX, hdu_LGSY, hdu_NGSX, hdu_NGSY, hdu_TSX, hdu_TSY
        ])
        t = self.typetype + "_StrehlMap"
        if (self.wavelengthwavelength > 0):
            t += "_{:.3f}nm".format(self.wavelengthwavelength)
        if (self._Rts_Rts > 0):
            t += "_{:.2f}arcsec".format(self._Rts_Rts)
        t += ".fits"
        if (name == ""):
            name = t
        hdul.writeto(name, overwrite=1)
 
    def read(self, name):
        hdu_map = fits.open(name)
        self.typetype = hdu_map[0].header["TYPE"]
        self.wavelengthwavelength = hdu_map[0].header["LAMBDA"]
        self._Rtar_Rtar = hdu_map[0].header["RTAR"]
        self._Rlgs_Rlgs = hdu_map[0].header["RLGS"]
        self._Rngs_Rngs = hdu_map[0].header["RNGS"]
        self._Rts_Rts = hdu_map[0].header["RTS"]
        self.mapmap = hdu_map[0].data
        self.LGSxLGSx = hdu_map["LGSx"].data
        self.LGSyLGSy = hdu_map["LGSy"].data
        self.NGSxNGSx = hdu_map["NGSx"].data
        self.NGSyNGSy = hdu_map["NGSy"].data
        self.TSxTSx = hdu_map["TSx"].data
        self.TSyTSy = hdu_map["TSy"].data