Commit 81857514 authored by rfetick's avatar rfetick
Browse files

create master class ParametricPSFfromPSD

parent 75908ba0
......@@ -397,25 +397,10 @@ 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)
#%% 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
"""
#%% MASTER CLASS
class ParametricPSFfromPSD(ParametricPSF):
"""This class is NOT to be instantiated"""
def __init__(self,nparam,Npix,system=None,samp=None):
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")
......@@ -424,11 +409,8 @@ class Turbulent(ParametricPSF):
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)
self._nparam = nparam
@property
def samp(self):
return self._samp_over/self._k
......@@ -437,52 +419,16 @@ class Turbulent(ParametricPSF):
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 len(x0)!=self._nparam: 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
......@@ -492,7 +438,7 @@ class Turbulent(ParametricPSF):
"""
if (self._k > 1) and (_caller != 'self'):
raise ValueError("Cannot call `Psfao.otf(...)` when undersampled (functionality not implemented yet)")
raise ValueError("Cannot call `otf(...)` when undersampled (functionality not implemented yet)")
OTF_TURBULENT = self._otf_turbulent(x0)
OTF_DIFFRACTION = self._otf_diffraction()
......@@ -500,13 +446,43 @@ class Turbulent(ParametricPSF):
return OTF_TURBULENT * OTF_DIFFRACTION * OTF_SHIFT
def getpix2freq(self):
"""
pixel to frequency conversion in the PSD plane
"""
return 1.0/(self.system.D*self._samp_over)
def getnxny_over(self):
"""
Return the number of pixels for the correctly sampled array (at least at Shannon-Nyquist)
"""
return self.Npix[0]*self._k, self.Npix[1]*self._k
def get_null_freq_index(self):
nx,ny = self.getnxny_over()
return nx//2, ny//2
def getfxfy(self):
"""
Return the array of frequencies (correctly sampled above Shannon-Nyquist)
"""
Nx_over, Ny_over = self.getnxny_over()
f2D = _np.mgrid[0:Nx_over, 0:Ny_over].astype(float)
f2D[0] -= Nx_over//2
f2D[1] -= Ny_over//2
f2D *= self.getpix2freq()
return f2D
def psd(self,x0):
raise ValueError("ParametricPSFfromPSD is not to be instantiated. the `psd` method must be override in the subclasses")
def _otf_turbulent(self,x0):
PSD = self.psd(x0)
PSD,integral = self.psd(x0)
L = self.system.D * self._samp_over
Bg = _fft.fft2(_fft.fftshift(PSD)) / L**2
Dphi = _fft.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.)
#Dphi = _fft.fftshift(_np.real(2 * (Bg[0, 0] - Bg))) # normalized on the numerical FoV
Dphi = _fft.fftshift(_np.real(2 * (integral - Bg))) # normalized up to infinity
return _np.exp(-Dphi/2.)
@_lru_cache(maxsize=2)
def _otf_diffraction(self):
......@@ -533,8 +509,7 @@ class Turbulent(ParametricPSF):
Parameters
----------
x0 : numpy.array (dim=1), tuple, list
x[0] - Fried parameter r0 [m]
x[1] - Von Karman external length [m]
Array of parameters for the PSD (see __doc__)
dx : float
PSF X shifting [pix] (default = 0)
dy : float
......@@ -546,15 +521,98 @@ class Turbulent(ParametricPSF):
if self._k==1:
return out
else:
return _binning(out,int(self._k))
return _binning(out,int(self._k)) # undersample PSF if needed (if it was oversampled for computations)
#%% TURBULENT PSF
class Turbulent(ParametricPSFfromPSD):
"""
Summary
-------
PSF model dedicated to long-exposure imaging with turbulence
p = Turbulent((npix,npix),system=system,samp=samp)
Description
-----------
the PSF is parametrised through the PSD of the electromagnetic phase
x[0] - Fried parameter r0 [m]
x[1] - Von Karman external length [m]
"""
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
"""
super().__init__(2,Npix,system=system,samp=samp)
# r0,L0
bounds_down = [_EPSILON, _EPSILON]
bounds_up = [_np.inf, _np.inf]
self.bounds = (bounds_down,bounds_up)
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 __doc__ for more details
Returns
-------
psd : numpy.array (dim=2)
integral : float : the integral of the `psd` up to infinity
"""
self.check_parameters(x0)
nx0,ny0 = self.get_null_freq_index()
f2D = self.getfxfy()
pix2freq = self.getpix2freq()
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[nx0,ny0] = 0.0
# Compute PSD integral up to infinity
fmax = _np.min([nx0,ny0])*pix2freq
integral_in = _np.sum(PSD*(F2<(fmax**2))) * pix2freq**2 # numerical sum
integral_out = 0.0229*6*_np.pi/5 * (r0*fmax)**(-5./3.) # analytical sum
integral = integral_in + integral_out
return PSD, integral
#%% PSFAO MODEL
class Psfao(ParametricPSF):
class Psfao(ParametricPSFfromPSD):
"""
Name: PSFAO
Description: a long-exposure PSF model dedicated to AO correction
Reference: Fétick et al., August 2019, A&A, Vol. 628 [Fétick2019b]
Note: the PSF is parametrised through the PSD of the electromagnetic phase
Summary
-------
PSF model dedicated to long-exposure imaging with adaptive optics
p = Psfao((npix,npix),system=system,samp=samp)
Description
-----------
the PSF is parametrised through the PSD of the electromagnetic phase
x[0] - Fried parameter : r0 [m]
x[1] - Corrected area background : C [rad² m²]
x[2] - Moffat phase variance : A [rad²]
x[3] - Moffat width : alpha [1/m]
x[4] - Moffat elongation ratio : sqrt(ax/ay)
x[5] - Moffat angle : theta [rad]
x[6] - Moffat power law : beta
Reference
---------
Fétick et al., August 2019, A&A, Vol.628
"""
def __init__(self,Npix,system=None,Lext=10.,samp=None):
"""
......@@ -570,40 +628,14 @@ class Psfao(ParametricPSF):
Von-Karman external scale (default = 10 m)
Useless if Fao >> 1/Lext
"""
super().__init__(7,Npix,system=system,samp=samp)
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.Lext = Lext
self.samp = samp
# r0,C,A,alpha,ratio,theta,beta
bounds_down = [_EPSILON,0,0,_EPSILON,_EPSILON,-_np.inf,1+_EPSILON]
bounds_up = [_np.inf for i in range(7)]
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)!=7: 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²]
......@@ -611,7 +643,7 @@ class Psfao(ParametricPSF):
Parameters
----------
x0 : numpy.array (dim=1), tuple, list
See __call__ for more details
See __doc__ for more details
Returns
-------
......@@ -620,15 +652,10 @@ class Psfao(ParametricPSF):
"""
self.check_parameters(x0)
nx0,ny0 = self.get_null_freq_index()
f2D = self.getfxfy()
pix2freq = self.getpix2freq()
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.
Fao = self.system.Nact/(2.0*self.system.D)
......@@ -643,85 +670,15 @@ class Psfao(ParametricPSF):
param = (ax,ay,theta,beta,0,0)
PSD += (F2 < Fao**2.) * _np.abs(C + A*moffat(f2D,param,norm=Fao))
# Set PSD = 0 at null frequency
PSD[Nx_over//2,Ny_over//2] = 0.0
PSD[nx0,ny0] = 0.0
# Compute PSD integral up to infinity
fmax = _np.min([Nx_over//2,Ny_over//2])*pix2freq
fmax = _np.min([nx0,ny0])*pix2freq
integral_in = _np.sum(PSD*(F2<(fmax**2))) * pix2freq**2 # numerical sum
integral_out = 0.0229*6*_np.pi/5 * (r0*fmax)**(-5./3.) # analytical sum
integral = integral_in + integral_out
return PSD, integral
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,integral = self.psd(x0)
L = self.system.D * self._samp_over
Bg = _fft.fft2(_fft.fftshift(PSD)) / L**2
#Dphi = _fft.fftshift(_np.real(2 * (Bg[0, 0] - Bg))) # normalized on the numerical FoV
Dphi = _fft.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 _fft.fftshift(_np.abs(_fft.ifft2(_np.abs(_fft.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] - PSD corrected area background C [rad² m²]
x[2] - PSD corrected area phase variance A [rad²]
x[3] - PSD alpha [1/m]
x[4] - PSD ax/ay ratio
x[5] - PSD theta [rad]
x[6] - PSD beta power law
dx : float
PSF X shifting [pix] (default = 0)
dy : float
PSF Y shifting [pix] (default = 0)
"""
out = _np.real(_fft.fftshift(_fft.ifft2(_fft.fftshift(self.otf(x0,dx=dx,dy=dy,_caller='self')))))
if self._k==1:
return out
else:
return _binning(out,int(self._k))
def tofits(self,param,filename,*args,keys=None,overwrite=False,**kwargs):
if keys is None:
......
Markdown is supported
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