From f6f9a6b854693ccc2b89920f3a758107a36a899f Mon Sep 17 00:00:00 2001 From: rfetick <r.fetick@gmail.com> Date: Thu, 27 Feb 2020 16:49:49 +0100 Subject: [PATCH] Simplification and bug correction on the PSD null frequency for undersampled cases --- maoppy/psfmodel.py | 139 ++++++++++++++++++++------------------------- 1 file changed, 61 insertions(+), 78 deletions(-) diff --git a/maoppy/psfmodel.py b/maoppy/psfmodel.py index 2f535e8..85900c1 100644 --- a/maoppy/psfmodel.py +++ b/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) -- GitLab