Commit f6f9a6b8 authored by FETICK Romain's avatar FETICK Romain
Browse files

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

parent e543d47d
......@@ -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)
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment