Add Turbulent PSF

maoppy/psfmodel.py

@@ -18,7 +18,7 @@ from functools import lru_cache
 # Value to compute finite differences
 _EPSILON = np.sqrt(sys.float_info.epsilon)
 
-#%% FITTING FUNCTION
+#%% FUNCTIONS
 def rmserror(model,image,weights=None,background=True, positive_bck=False):
     """Compute the RMS error between a PSF model and an image"""
     if weights is None: weights = np.ones_like(image)
@@ -29,6 +29,12 @@ def rmserror(model,image,weights=None,background=True, positive_bck=False):
     err = err/np.sum(image)
     return err
 
+def oversample(samp):
+    """Find the minimal integer that allows oversampling"""
+    k = int(np.ceil(2.0/samp))
+    return (k*samp,k)
+
+#%% FITTING
 def lsq_flux_bck(model, data, weights, background=True, positive_bck=False):
     """Compute the analytical least-square solution for flux and background
     LS = SUM_pixels { weights*(flux*model + bck - data)² }
@@ -401,14 +407,158 @@ class Gaussian(ParametricPSF):
     def tofits(self,param,filename,*args,keys=["SIGMA_X","SIGMA_Y","THETA"],**kwargs):
         super(Gaussian,self).tofits(param,filename,*args,keys=keys,**kwargs)
 
-
-def oversample(samp):
-    """Find the minimal integer that allows oversampling"""
-    k = int(np.ceil(2.0/samp))
-    return (k*samp,k)
+#%% TURBULENT PSF
+class Turbulent(ParametricPSF):
+    """
+    Name:        PSFAO
+    Description: long-exposure turbulent PSF
+    Note:        the PSF is parametrised through the PSD of the electromagnetic phase
+    """
+    def __init__(self,Npix,system=None,samp=None):
+        """
+        Parameters
+        ----------
+        Npix : tuple
+            Size of output PSF
+        system : OpticalSystem
+            Optical system for this PSF
+        samp : float
+            Sampling at the observation wavelength
+        """
+        
+        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
+        self.system = system
+        self.samp = samp
+        # r0,L0
+        bounds_down = [_EPSILON, _EPSILON]
+        bounds_up = [np.inf, np.inf]
+        self.bounds = (bounds_down,bounds_up)
+    
+    @property
+    def samp(self):
+        return self._samp_over/self._k
+    
+    @samp.setter
+    def samp(self,value):
+        # Manage cases of undersampling
+        self._samp_over, self._k = oversample(value)
+    
+    def check_parameters(self,x0):
+        bd, bu = self.bounds
+        x0 = np.array(x0)
+        bd = np.array(bd)
+        bu = np.array(bu)
+        if len(x0)!=2: raise ValueError('len(x0) is different from length of bounds')
+        if np.any(x0<bd): raise ValueError('Lower bounds are not respected')
+        if np.any(x0>bu): raise ValueError('Upper bounds are not respected')
+        
+    
+    def psd(self,x0):
+        """Compute the PSD model from parameters
+        PSD is given in [rad²/f²] = [rad² m²]
+        
+        Parameters
+        ----------
+        x0 : numpy.array (dim=1), tuple, list
+            See __call__ for more details
+            
+        Returns
+        -------
+        psd : numpy.array (dim=2)
+        integral : float : the integral of the `psd` up to infinity
+            
+        """
+        self.check_parameters(x0)
+        
+        Nx_over = self.Npix[0]*self._k
+        Ny_over = self.Npix[1]*self._k
         
+        pix2freq = 1.0/(self.system.D*self._samp_over)
+        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.
+        
+        r0,Lext = x0
+        PSD = 0.0229* r0**(-5./3.) * ((1./Lext**2.) + F2)**(-11./6.)
 
+        # 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'):
+        """
+        See __call__ for input arguments
+        Warning: result of otf will be unconsistent if undersampled!!!
+        This issue is solved with oversampling + binning in __call__ but not here
+        For the moment, the `_caller` keyword prevents user to misuse otf
+        """
+        
+        if (self._k > 1) and (_caller != 'self'):
+            raise ValueError("Cannot call `Psfao.otf(...)` when undersampled (functionality not implemented yet)")
+        
+        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_over
+        Bg = fft2(fftshift(PSD)) / L**2
+        Dphi = fftshift(np.real(2 * (Bg[0, 0] - Bg))) # normalized on the numerical FoV 
+        #Dphi = fftshift(np.real(2 * (integral - Bg))) # normalized up to infinity
+        return np.exp(-Dphi/2.) 
+    
+    @lru_cache(maxsize=2)
+    def _otf_diffraction(self):
+        Nx_over = self.Npix[0]*self._k
+        Ny_over = self.Npix[1]*self._k
+        
+        NpupX = np.ceil(Nx_over/self._samp_over)
+        NpupY = np.ceil(Ny_over/self._samp_over)
+        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_over)
+        return fftshift(np.abs(ifft2(np.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):
+        """
+        Parameters
+        ----------
+        x0 : numpy.array (dim=1), tuple, list
+            x[0] - Fried parameter r0 [m]
+            x[1] - Von Karman external length [m]
+        dx : float
+            PSF X shifting [pix] (default = 0)
+        dy : float
+            PSF Y shifting [pix] (default = 0)
+        """
+        
+        out = np.real(fftshift(ifft2(fftshift(self.otf(x0,dx=dx,dy=dy,_caller='self')))))
+        
+        if self._k==1:
+            return out
+        else:
+            return binning(out,int(self._k))
 
+#%% PSFAO MODEL
 class Psfao(ParametricPSF):
     """
     Name:        PSFAO