Simplification and bug correction on the PSD null frequency for undersampled cases

maoppy/psfmodel.py

@@ -343,13 +343,20 @@ class Moffat(ParametricPSF):
         super(Moffat,self).tofits(param,filename,*args,keys=keys,**kwargs)
 
 
-class Psfao(ParametricPSF):
-    """PSF model based on a parametrization of the phase PSD
-    See documentation of methods "__init__" and "__call__"
-    
-    """
-    
-    def __init__(self,Npix,system=None,Lext=10.,samp=None,symmetric=False,diffotf=True):
+
+
+def oversample(samp):
+    """Oversample with an integer"""
+    if samp>2:
+        return (samp,1)
+    else:
+        k = int(np.ceil(2.0/samp))
+        return (k*samp,k)
+        
+
+
+class Psfao(ParametricPSF):    
+    def __init__(self,Npix,system=None,Lext=10.,samp=None,symmetric=False):
         """
         Parameters
         ----------
@@ -362,27 +369,18 @@ class Psfao(ParametricPSF):
         Lext : float
             Von-Karman external scale (default = 10 m)
             Useless if Fao >> 1/Lext
-        diffotf : bool
-            Enable/disable diffraction OTF for PSF computation
-            (default=True)
         """
         #super(ParametricPSF,self).__init__(Npix)
-        if not (type(Npix) in [tuple,list,np.ndarray]):
-            raise ValueError("Npix must be a tuple, list or numpy.ndarray")
-        if len(Npix)!=2:
-            raise ValueError("Npix must be of length = 2")
-        if (Npix[0]%2) or (Npix[1]%2):
-            raise ValueError("Each Npix component must be even")
+        if not (type(Npix) in [tuple,list,np.ndarray]): raise ValueError("Npix must be a tuple, list or numpy.ndarray")
+        if len(Npix)!=2: raise ValueError("Npix must be of length = 2")
+        if (Npix[0]%2) or (Npix[1]%2): raise ValueError("Each Npix component must be even")
+        if system is None: raise ValueError("Keyword `system` must be defined")
+        if samp is None: raise ValueError("Keyword `samp` must be defined")
         self.Npix = Npix
-        if system is None:
-            raise ValueError("Keyword `system` must be defined")
-        if samp is None:
-            raise ValueError("Keyword `samp` must be defined")
         self.system = system
         self.Lext = Lext
         self.samp = samp
         self.symmetric = symmetric
-        self.diffotf = diffotf
     
     @property
     def symmetric(self):
@@ -407,49 +405,7 @@ class Psfao(ParametricPSF):
     def samp(self,value):
         # Manage cases of undersampling
         self._samp = value
-        if value >=2:
-            self._samp_num = value
-            self._k = 1
-        else:
-            self._k = int(np.ceil(2.0/value))
-            self._samp_num = self._k * value
-    
-    @lru_cache(maxsize=2)
-    def _freq_array(self,Nx,Ny,samp,D):
-        """
-        Returns
-        -------
-        tab - numpy.array (dim=3)
-            3D array of frequencies [1/m]
-        """
-        pix2freq = 1.0/(D*samp)
-        f2D = np.mgrid[0:Nx, 0:Ny].astype(float)
-        #null frequency at [Nx//2,Ny//2] according to numpy fft convention
-        f2D[0] -= Nx//2
-        f2D[1] -= Ny//2
-        return f2D * pix2freq
-    
-    @lru_cache(maxsize=2)
-    def _shift_array(self,Nx,Ny):
-        Y, X = np.mgrid[0:Nx,0:Ny].astype(float)
-        X = (X-Nx/2) * 2*np.pi*1j/Nx
-        Y = (Y-Ny/2) * 2*np.pi*1j/Ny
-        return X, Y
-    
-    @lru_cache(maxsize=5)
-    def _dlFTO(self,Nx,Ny,pupfct,samp):
-        # samp as a tuple is not ready yet, don't use it
-        if type(samp)==tuple:
-            dlFTO = np.zeros((Nx,Ny,len(samp)))
-            for i in range(len(samp)):
-                dlFTO[...,i] = self._dlFTO(Nx,Ny,pupfct, samp[i])
-            return np.mean(dlFTO,axis=2)
-        
-        NpupX = np.ceil(Nx/samp)
-        NpupY = np.ceil(Ny/samp)
-        tab = np.zeros((Nx, Ny), dtype=np.complex)
-        tab[0:int(NpupX), 0:int(NpupY)] = pupfct((NpupX,NpupY),samp=samp)
-        return fftshift(abs(ifft2(abs(fft2(tab)) ** 2)) / np.sum(tab))
+        self._samp_num, self._k = oversample(value)
     
     def psd(self,x0):
         """Compute the PSD model from parameters
@@ -472,8 +428,16 @@ class Psfao(ParametricPSF):
         else:
             raise ValueError("Wrong size of x0")
         
-        f2D = self._freq_array(self.Npix[0]*self._k,self.Npix[1]*self._k,self._samp_num,self.system.D)
-        F2 = f2D[0] ** 2. + f2D[1] ** 2.
+        Nx_over = self.Npix[0]*self._k
+        Ny_over = self.Npix[1]*self._k
+        
+        pix2freq = 1.0/(self.system.D*self._samp_num)
+        f2D = np.mgrid[0:Nx_over, 0:Ny_over].astype(float)
+        f2D[0] -= Nx_over//2
+        f2D[1] -= Ny_over//2
+        f2D *= pix2freq
+        F2 = f2D[0]**2. + f2D[1]**2.
+        
         Fao = self.system.Nact/(2.0*self.system.D)
         
         PSD = 0.0229* x[0]**(-5./3.) * ((1. / self.Lext**2.) + F2)**(-11./6.)
@@ -481,8 +445,8 @@ class Psfao(ParametricPSF):
         
         param = np.concatenate((x[2:],[0,0]))
         PSD += (F2 < Fao**2.) * np.abs(x[1] + moffat(f2D,param,norm=Fao))
-        # Set PSD = 0 at null frequency (according to numpy fft convention)
-        PSD[self.Npix[0]//2,self.Npix[1]//2] = 0.0
+        # Set PSD = 0 at null frequency
+        PSD[Nx_over//2,Ny_over//2] = 0.0
         return PSD
     
     def otf(self,x0,dx=0,dy=0,_caller='user'):
@@ -496,19 +460,37 @@ class Psfao(ParametricPSF):
         if (self._k > 1) and (_caller != 'self'):
             raise ValueError("Cannot call `Psfao.otf(...)` when undersampled (functionality not implemented yet)")
         
-        PSD = self.psd(x0)
+        OTF_TURBULENT = self._otf_turbulent(x0)
+        OTF_DIFFRACTION = self._otf_diffraction()
+        OTF_SHIFT = self._otf_shift(dx,dy)
         
+        return OTF_TURBULENT * OTF_DIFFRACTION * OTF_SHIFT
+    
+    def _otf_turbulent(self,x0):
+        PSD = self.psd(x0)
         L = self.system.D * self._samp_num
-        Bg = fft2(fftshift(PSD)) / L ** 2
+        Bg = fft2(fftshift(PSD)) / L**2
         Dphi = fftshift(np.real(2 * (Bg[0, 0] - Bg)))
+        return np.exp(-Dphi/2.) 
+    
+    def _otf_diffraction(self):  
+        Nx_over = self.Npix[0]*self._k
+        Ny_over = self.Npix[1]*self._k
         
-        if self.diffotf:
-            dlFTO = self._dlFTO(self.Npix[0]*self._k,self.Npix[1]*self._k, 
-                            self.system.pupil, self._samp_num)      
-        else:
-            dlFTO = 1.
-        X, Y = self._shift_array(self.Npix[0]*self._k,self.Npix[1]*self._k)
-        return np.exp(-Dphi/2.)*dlFTO*np.exp(-X*dx*self._k - Y*dy*self._k)
+        NpupX = np.ceil(Nx_over/self._samp_num)
+        NpupY = np.ceil(Ny_over/self._samp_num)
+        tab = np.zeros((Nx_over, Ny_over), dtype=np.complex)
+        tab[0:int(NpupX), 0:int(NpupY)] = self.system.pupil((NpupX,NpupY),samp=self._samp_num)
+        return fftshift(abs(ifft2(abs(fft2(tab)) ** 2)) / np.sum(tab))
+    
+    def _otf_shift(self,dx,dy):
+        Nx_over = self.Npix[0]*self._k
+        Ny_over = self.Npix[1]*self._k
+        
+        Y, X = np.mgrid[0:Nx_over,0:Ny_over].astype(float)
+        X = (X-Nx_over/2) * 2*np.pi*1j/Nx_over
+        Y = (Y-Ny_over/2) * 2*np.pi*1j/Ny_over
+        return np.exp(-X*dx*self._k - Y*dy*self._k)
     
     def __call__(self,x0,dx=0,dy=0):
         """
@@ -568,9 +550,10 @@ class Psfao(ParametricPSF):
         hdr = self._getfitshdr(param,keys=keys,keys_comment=keys_comment)
         
         hdr["HIERARCH SYSTEM"] = (self.system.name,"System name")
+        hdr["HIERARCH SYSTEM D"] = (self.system.D,"Primary mirror diameter")
+        hdr["HIERARCH SYSTEM NACT"] = (self.system.Nact,"Linear number of AO actuators")
         hdr["HIERARCH SAMP"] = (self.samp,"Sampling (eg. 2 for Shannon)")
         hdr["HIERARCH LEXT"] = (self.Lext,"Von-Karman outer scale")
-        hdr["HIERARCH DIFFOTF"] = (self.diffotf,"Is diffraction OTF enabled")
         
         hdu = fits.PrimaryHDU(psf, hdr)
         hdu.writeto(filename, overwrite=True)