diff --git a/maoppy/psfmodel.py b/maoppy/psfmodel.py
index 2f535e8ee71dae6ad02aae5ee383bc769a89515a..85900c15cc2890aa493202829ec39c2b190e2308 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)