Commit b8fe833d authored by rfetick's avatar rfetick
Browse files

syntax improvements

parent e0c4f3ec
......@@ -6,20 +6,21 @@ Created on Mon May 27 17:30:51 2019
@author: rfetick
"""
from maoppy.utils import circarr as _circarr
from maoppy.utils import RAD2ARCSEC as _RAD2ARCSEC
import sys
import os
import sys as _sys
import yaml as _yaml
import os as _os
from astropy.io import fits as _fits
import yaml
from astropy.io import fits
import numpy as _np
from scipy.interpolate import interp2d as _interp2d
from scipy.interpolate import interp2d
from maoppy.utils import circarr as _circarr
from maoppy.utils import RAD2ARCSEC as _RAD2ARCSEC
def _get_data_folder():
folder = _os.path.abspath(__file__)
folder = _os.sep.join(folder.split(_os.sep)[0:-1])+_os.sep+'data'+_os.sep
folder = os.path.abspath(__file__)
folder = os.sep.join(folder.split(os.sep)[0:-1])+os.sep+'data'+os.sep
return folder
#def load(name):
......@@ -61,10 +62,14 @@ class Instrument(object):
def __init__(self,D=None,occ=0.,res=None,Nact=0,gain=1.,ron=0.):
if D is None: raise ValueError("Please enter keyword `D` to set Instrument's aperture diameter")
if res is None: raise ValueError("Please enter keyword `res` to set instrument resolution in rad")
if D <= 0: raise ValueError("Keyword `D` must be strictly positive")
if res <= 0: raise ValueError("Keyword `res` must be strictly positive")
if D is None:
raise ValueError("Please enter keyword `D` to set Instrument's aperture diameter")
if res is None:
raise ValueError("Please enter keyword `res` to set instrument resolution in rad")
if D <= 0:
raise ValueError("Keyword `D` must be strictly positive")
if res <= 0:
raise ValueError("Keyword `res` must be strictly positive")
self.D = D
self.occ = occ # occultation diameter ratio
......@@ -92,9 +97,9 @@ class Instrument(object):
s += "Diameter : %.2f m (occ=%u%%)\n" % (self.D,self.occ*100)
s += "Resolution : %.1f mas (binning=%u)\n" % (self.resolution_mas,self.binning)
s += "Nact : %u\n" % self.Nact
K = tuple(self.filters.keys())
keys = tuple(self.filters.keys())
s += "Filters : " # %u" % len(self.filters)
for k in K:
for k in keys:
s += "%s " % k
s += "\n"
s += "Detector : gain=%.1f e-/ADU\n"%self.gain
......@@ -111,23 +116,24 @@ class Instrument(object):
def resolution_mas(self):
return self.resolution_rad * _RAD2ARCSEC * 1e3
def pupil(self,Npix,wvl=None,samp=None):
def pupil(self,shape,wvl=None,samp=None):
"""Returns the 2D array of the pupil transmission function (complex data)"""
Dpix = min(Npix)/2
pup = _circarr(Npix)
Dpix = min(shape)/2
pup = _circarr(shape)
if self.phasemask_enable:
if self._phasemask is None:
if self.phasemask_path is None: raise ValueError('phasemask_path must be defined')
p = _fits.open(self.phasemask_path)[0].data * 1e-9 # fits data in nm, converted here to meter
if self.phasemask_path is None:
raise ValueError('phasemask_path must be defined')
p = fits.open(self.phasemask_path)[0].data * 1e-9 # fits data in nm, converted here to meter
x = _np.arange(p.shape[0])/p.shape[0]
y = _np.arange(p.shape[1])/p.shape[1]
self._phasemask = _interp2d(x,y,p)
self._phasemask = interp2d(x,y,p)
cx,cy = self.phasemask_shift
x = _np.arange(Npix[0])/Npix[0] - cx/Npix[0]
y = _np.arange(Npix[1])/Npix[1] - cy/Npix[1]
wf = self._phasemask(x,y)
if wvl is None: wvl = self.wvl(samp) # samp must be defined if wvl is None
wf = _np.exp(2j*_np.pi/wvl*wf)
x = _np.arange(shape[0])/shape[0] - cx/shape[0]
y = _np.arange(shape[1])/shape[1] - cy/shape[1]
if wvl is None:
wvl = self.wvl(samp) # samp must be defined if wvl is None
wf = _np.exp(2j*_np.pi/wvl*self._phasemask(x,y))
else:
wf = 1.0 + 0j # complex type for output array, even if real data
return (pup < Dpix) * (pup >= Dpix*self.occ) * wf
......@@ -154,15 +160,16 @@ class Instrument(object):
data['filters'] = self.filters
data['phasemask_path'] = self.phasemask_path
data['phasemask_shift'] = self.phasemask_shift
with open(filename,'w') as f: _yaml.dump(data,f)
with open(filename,'w') as f:
yaml.dump(data,f)
#%% LOAD INSTRUMENT INSTANCES (make them attributes of this module)
_this_module = _sys.modules[__name__]
_this_module = sys.modules[__name__]
_d = _get_data_folder()
_l = [_f for _f in _os.listdir(_d) if _f.endswith('.yml')]
_l = [_f for _f in os.listdir(_d) if _f.endswith('.yml')]
for _f in _l:
with open(_d+_f,'r') as _fi:
_i = _yaml.full_load(_fi)
_i = yaml.full_load(_fi)
_instru = Instrument(D=_i['d'],occ=_i['occ'],res=_i['res'],gain=_i['gain'],ron=_i['ron'],Nact=_i['nact'])
_instru.name = _i['name']
_instru.fullname = _i['fullname']
......
......@@ -218,21 +218,21 @@ def psffit(psf,Model,x0,weights=None,dxdy=(0.,0.),flux_bck=(True,True),
#%% BASIC FUNCTIONS
def reduced_coord(XY,ax,ay,theta,cx,cy,grad=False):
def reduced_coord(xxyy,ax,ay,theta,cx,cy,grad=False):
"""Create an array of reduced coordinated with elongation and rotation"""
c = _np.cos(theta)
s = _np.sin(theta)
s2 = _np.sin(2.0 * theta)
Rxx = (c/ax)**2 + (s/ay)**2
Rxy = s2/ay**2 - s2/ax**2
Ryy = (c/ay)**2 + (s/ax)**2
rxx = (c/ax)**2 + (s/ay)**2
rxy = s2/ay**2 - s2/ax**2
ryy = (c/ay)**2 + (s/ax)**2
uxx = (XY[0]-cx)**2
uxy = (XY[0]-cx)*(XY[1]-cy)
uyy = (XY[1]-cy)**2
uxx = (xxyy[0]-cx)**2
uxy = (xxyy[0]-cx)*(xxyy[1]-cy)
uyy = (xxyy[1]-cy)**2
uu = Rxx*uxx + Rxy*uxy + Ryy*uyy
uu = rxx*uxx + rxy*uxy + ryy*uyy
if grad:
gg = _np.zeros((3,)+uu.shape)
......@@ -287,19 +287,20 @@ def moffat(xxyy, param, norm=None, removeInside=0, grad=False):
ax,ay,theta,beta,cx,cy = param
if grad:
u,du = reduced_coord(xxyy,ax,ay,theta,cx,cy,grad=True)
uu,du = reduced_coord(xxyy,ax,ay,theta,cx,cy,grad=True)
else:
u = reduced_coord(xxyy,ax,ay,theta,cx,cy,grad=False)
uu = reduced_coord(xxyy,ax,ay,theta,cx,cy,grad=False)
V = (1.0+u)**(-beta) # Moffat shape
V = (1.0+uu)**(-beta) # Moffat shape
E = 1.0 # normalization factor (eventually up to infinity)
F = 1.0 # normalization factor (eventually up to a limited radius)
if grad:
dVdu = -beta*V/(1.0+u)
dV = _np.zeros((4,)+u.shape)
for i in range(3): dV[i,...] = dVdu*du[i,...]
dV[3,...] = -V*_np.log(1.0+u)
dVdu = -beta*V/(1.0+uu)
dV = _np.zeros((4,)+uu.shape)
for i in range(3):
dV[i,...] = dVdu*du[i,...]
dV[3,...] = -V*_np.log(1.0+uu)
if norm is None:
if grad:
......@@ -348,7 +349,7 @@ def gauss(xxyy,param):
ax = _np.sqrt(2)*param[0]
ay = _np.sqrt(2)*param[1]
uu = reduced_coord(xxyy,ax,ay,param[2],param[3],param[4])
return 1.0/(2*_np.pi*param[0]*param[1])*_np.exp(-uu)
return _np.exp(-uu) / (2*_np.pi*param[0]*param[1])
#%% CLASS PARAMETRIC PSF AND ITS SUBCLASSES
......@@ -433,22 +434,22 @@ class Moffat(ParametricPSF):
bounds_up = [_np.inf for i in range(4)]
self.bounds = (bounds_down,bounds_up)
self._XY = _np.mgrid[0:npix[0],0:npix[1]]
self._XY[0] -= npix[0]//2
self._XY[1] -= npix[1]//2
self._xxyy = _np.mgrid[0:npix[0],0:npix[1]]
self._xxyy[0] -= npix[0]//2
self._xxyy[1] -= npix[1]//2
def __call__(self,x,dx=0,dy=0):
def __call__(self, param, dx=0, dy=0):
"""
Parameters
----------
x : list, tuple, numpy.ndarray (len=4)
x[0] - Alpha X
x[1] - Alpha Y
x[2] - Theta
x[3] - Beta
param : list, tuple, numpy.ndarray (len=4)
param[0] - Alpha X
param[1] - Alpha Y
param[2] - Theta
param[3] - Beta
"""
param = _np.concatenate((x,[dx,dy]))
return moffat(self._XY, param, norm=self.norm)
paramdxdy = _np.concatenate((param,[dx,dy]))
return moffat(self._xxyy, paramdxdy, norm=self.norm)
def tofits(self,param,filename,*args,keys=["ALPHA_X","ALPHA_Y","THETA","BETA"],**kwargs):
super(Moffat,self).tofits(param,filename,*args,keys=keys,**kwargs)
......@@ -462,21 +463,21 @@ class Gaussian(ParametricPSF):
bounds_up = [_np.inf for i in range(3)]
self.bounds = (bounds_down,bounds_up)
self._XY = _np.mgrid[0:npix[0],0:npix[1]]
self._XY[1] -= npix[0]/2
self._XY[0] -= npix[1]/2
self._xxyy = _np.mgrid[0:npix[0],0:npix[1]]
self._xxyy[1] -= npix[0]/2
self._xxyy[0] -= npix[1]/2
def __call__(self,x,dx=0,dy=0):
def __call__(self, param, dx=0, dy=0):
"""
Parameters
----------
x : list, tuple, numpy.ndarray (len=4)
x[0] - Sigma X
x[1] - Sigma Y
x[2] - Theta
param : list, tuple, numpy.ndarray (len=4)
param[0] - Sigma X
param[1] - Sigma Y
param[2] - Theta
"""
param = _np.concatenate((x,[dx,dy]))
return gauss(self._XY, param)
paramdxdy = _np.concatenate((param,[dx,dy]))
return gauss(self._xxyy, paramdxdy)
def tofits(self,param,filename,*args,keys=["SIGMA_X","SIGMA_Y","THETA"],**kwargs):
super(Gaussian,self).tofits(param,filename,*args,keys=keys,**kwargs)
......@@ -500,12 +501,12 @@ class ParametricPSFfromPSD(ParametricPSF):
Methods
-------
strehlOTF : compute Strehl ratio from OTF ratio (recommanded)
strehlMarechal : compute Strehl ratio using Marechal approximation (only in high Strehl case)
psd : return the phase PSD
otf : return the OTF
__call__ : return the PSF
tofits : write a FITS file of the PSF
strehlOTF : compute Strehl ratio from OTF ratio (recommanded).
strehlMarechal : compute Strehl ratio using Marechal approximation (only in high Strehl case).
psd : return the phase PSD.
otf : return the OTF.
__call__ : return the PSF.
tofits : write a FITS file of the PSF.
"""
def __init__(self, nparam, npix, system=None, samp=None, fixed_k=None):
......@@ -538,11 +539,17 @@ class ParametricPSFfromPSD(ParametricPSF):
self._npix = value
self._computeXYarray()
self._computeOtfDiffraction()
@property
def npixOver(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
@property
def samp(self):
"""Sampling to compute arrays"""
"""Sampling of the output PSF"""
return self._samp_over/self._k
......@@ -555,29 +562,23 @@ class ParametricPSFfromPSD(ParametricPSF):
def _computeXYarray(self):
nx,ny = self._nxnyOver
nx,ny = self.npixOver
xyarray = _np.mgrid[0:nx, 0:ny].astype(float)
xyarray[0] -= nx//2
xyarray[1] -= ny//2
self._fxfy = xyarray * self._pix2freq
self._fxfy = xyarray * self.pix2freq
self._f2 = self._fxfy[0]**2. + self._fxfy[1]**2.
@property
def _pix2freq(self):
def pix2freq(self):
"""Pixel to frequency conversion in the PSD plane"""
return 1.0/(self.system.D*self._samp_over)
@property
def _nxnyOver(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
@property
def _nullFreqIndex(self):
nx,ny = self._nxnyOver
nx,ny = self.npixOver
return nx//2, ny//2
......@@ -589,55 +590,56 @@ class ParametricPSFfromPSD(ParametricPSF):
return fftshift(self._otfDiffraction)
def strehlOTF(self, x0):
def strehlOTF(self, parampsd):
"""Compute the Strehl ratio based on the sum of the OTF"""
return _np.real(_np.sum(self._otf(x0))/_np.sum(self._otfDiffraction))
return _np.real(_np.sum(self._otf(parampsd))/_np.sum(self._otfDiffraction))
def strehlMarechal(self, x0):
def strehlMarechal(self, parampsd):
"""
Compute the Strehl ratio based on the Maréchal approximation.
Note that `strehlOTF` might provide more accurate results.
"""
_,sig2 = self.psd(x0)
_,sig2 = self.psd(parampsd)
return _np.exp(-sig2)
def psd(self, x0, grad=False):
def psd(self, parampsd, grad=False):
"""Compute the PSD. This method should be overriden by the dedicated method from subclasses."""
raise ValueError("ParametricPSFfromPSD is not to be instantiated. the `psd` method must be override in the subclasses")
def check_parameters(self, x0):
bd, bu = self.bounds
x0 = _np.array(x0)
bd = _np.array(bd)
bu = _np.array(bu)
if len(x0)!=self._nparam:
raise ValueError('len(x0) is different from length of bounds')
if _np.any(x0<bd):
def check_parameters(self, parampsd):
"""Check whether input parameters comply with bounds"""
bdn, bup = self.bounds
parampsd = _np.array(parampsd)
bdn = _np.array(bdn)
bup = _np.array(bup)
if len(parampsd)!=self._nparam:
raise ValueError('len(parampsd) is different from length of bounds')
if _np.any(parampsd<bdn):
raise ValueError('Lower bounds are not respected')
if _np.any(x0>bu):
if _np.any(parampsd>bup):
raise ValueError('Upper bounds are not respected')
def otf(self, x0, dx=0, dy=0, grad=False):
def otf(self, parampsd, dx=0, dy=0, grad=False):
"""
Public implemetation of the OTF.
Only available if samp>=2.
Null frequency is at the center of the array.
"""
if (self._k > 1):
if self._k > 1:
raise ValueError("Cannot call `otf(...)` when undersampled")
if grad:
otf, gg = self._otf(x0, dx=dx, dy=dy, grad=True)
otf, gg = self._otf(parampsd, dx=dx, dy=dy, grad=True)
return fftshift(otf), fftshift(gg,axes=(1,2))
else:
otf = self._otf(x0, dx=dx, dy=dy, grad=False)
return fftshift(otf)
otf = self._otf(parampsd, dx=dx, dy=dy, grad=False)
return fftshift(otf)
def _otf(self, x0, dx=0, dy=0, grad=False):
def _otf(self, parampsd, dx=0, dy=0, grad=False):
"""
Private implementation of the OTF.
It is correctly oversampled.
......@@ -646,11 +648,11 @@ class ParametricPSFfromPSD(ParametricPSF):
"""
otfShift = self._otfShift(dx,dy)
if grad:
otfTurb, gg = self._otfTurbulent(x0,grad=True)
otfTurb, gg = self._otfTurbulent(parampsd,grad=True)
for i in range(gg.shape[0]):
gg[i,:] *= self._otfDiffraction*otfShift
else:
otfTurb = self._otfTurbulent(x0,grad=False)
otfTurb = self._otfTurbulent(parampsd,grad=False)
otf = otfTurb * self._otfDiffraction * otfShift
......@@ -659,33 +661,33 @@ class ParametricPSFfromPSD(ParametricPSF):
return otf
def _otfTurbulent(self, x0, grad=False):
def _otfTurbulent(self, parampsd, grad=False):
"""Atmospheric part of the OTF"""
L = self.system.D * self._samp_over
df2 = self.pix2freq**2
if grad:
psd, integral, g, integral_g = self.psd(x0,grad=True)
psd, integral, gg, integral_g = self.psd(parampsd,grad=True)
else:
psd, integral = self.psd(x0,grad=False)
psd, integral = self.psd(parampsd,grad=False)
Bphi = _np.real(fft2(fftshift(psd))) / L**2
Bphi = _np.real(fft2(fftshift(psd))) * df2
#Bphi = Bphi[0, 0] - Bphi # normalized on the numerical FoV
Bphi = integral - Bphi # normalized up to infinity
otf = _np.exp(-Bphi)
if grad:
g2 = _np.zeros(g.shape,dtype=complex) # I cannot override 'g' here due to float to complex type
for i in range(len(x0)):
Bphi = _np.real(fft2(fftshift(g[i,...]))) / L**2
gg2 = _np.zeros(gg.shape,dtype=complex) # I cannot override 'g' here due to float to complex type
for i in range(len(parampsd)):
Bphi = _np.real(fft2(fftshift(gg[i,...]))) * df2
#Bphi = Bphi[0, 0] - Bphi # normalized on the numerical FoV
Bphi = integral_g[i] - Bphi # normalized up to infinity
g2[i,...] = -Bphi*otf
return otf, g2
gg2[i,...] = -Bphi*otf
return otf, gg2
return otf
def _computeOtfDiffraction(self):
"""Precompute the diffraction part of the OTF"""
nx, ny = self._nxnyOver
nx, ny = self.npixOver
npupx = _np.ceil(nx/self._samp_over)
npupy = _np.ceil(ny/self._samp_over)
tab = _np.zeros((nx, ny), dtype=complex)
......@@ -695,7 +697,7 @@ class ParametricPSFfromPSD(ParametricPSF):
def _otfShift(self, dx, dy):
"""Shifting part of the OTF"""
yy, xx = self._fxfy/self._pix2freq # so it is just an array of pixels
yy, xx = self._fxfy/self.pix2freq # so it is just an array of pixels
ny,nx = xx.shape
# Compensate oversampling shift
dx += (self._k-1)/(2*self._k)
......@@ -706,11 +708,11 @@ class ParametricPSFfromPSD(ParametricPSF):
return ifftshift(_np.exp(-2j*_np.pi*self._k*(dx*xx/nx + dy*yy/ny)))
def __call__(self, x0, dx=0, dy=0, grad=False):
def __call__(self, parampsd, dx=0, dy=0, grad=False):
"""
Parameters
----------
x0 : numpy.array (dim=1), tuple, list
parampsd : numpy.array (dim=1), tuple, list
Array of parameters for the PSD (see __doc__)
Keywords
......@@ -724,11 +726,11 @@ class ParametricPSFfromPSD(ParametricPSF):
"""
if grad:
otf, gg = self._otf(x0, dx=dx, dy=dy, grad=True)
for i in range(len(x0)):
otf, gg = self._otf(parampsd, dx=dx, dy=dy, grad=True)
for i in range(len(parampsd)):
gg[i,...] = fftshift(_np.real(ifft2(gg[i,...])))
else:
otf = self._otf(x0, dx=dx, dy=dy, grad=False)
otf = self._otf(parampsd, dx=dx, dy=dy, grad=False)
psf = fftshift(_np.real(ifft2(otf)))
......@@ -736,13 +738,14 @@ class ParametricPSFfromPSD(ParametricPSF):
if k==1:
if grad: return psf, gg
return psf
else:
if grad:
g2 = _np.zeros((len(x0),psf.shape[0]//k,psf.shape[1]//k))
for i in range(len(x0)):
g2[i,...] = binning(gg[i,...].astype(float),k)
return binning(psf,k), g2
return binning(psf,k) # undersample PSF if needed (if it was oversampled for computations)
if grad:
gg2 = _np.zeros((len(parampsd),psf.shape[0]//k,psf.shape[1]//k))
for i in range(len(parampsd)):
gg2[i,...] = binning(gg[i,...].astype(float),k)
return binning(psf,k), gg2
return binning(psf,k) # undersample PSF if needed (if it was oversampled for computations)
#%% TURBULENT PSF
......@@ -777,13 +780,13 @@ class Turbulent(ParametricPSFfromPSD):
bounds_up = [_np.inf, _np.inf]
self.bounds = (bounds_down,bounds_up)
def psd(self, x0, grad=False):
def psd(self, parampsd, grad=False):
"""Compute the PSD model from parameters
PSD is given in [rad²/f²] = [rad² m²]
Parameters
----------
x0 : numpy.array (dim=1), tuple, list
parampsd : numpy.array (dim=1), tuple, list
See __doc__ for more details
Returns
......@@ -792,30 +795,30 @@ class Turbulent(ParametricPSFfromPSD):
integral : float : the integral of the `psd` up to infinity
"""
self.check_parameters(x0)
self.check_parameters(parampsd)
nx0,ny0 = self._nullFreqIndex
r0,Lext = x0
r0,Lext = parampsd
psd = 0.0229* r0**(-5./3.) * ((1./Lext**2.) + self._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])*self._pix2freq
integral_in = _np.sum(psd*(self._f2<(fmax**2))) * self._pix2freq**2 # numerical sum (in the array's tangent circle)
fmax = _np.min([nx0,ny0])*self.pix2freq
integral_in = _np.sum(psd*(self._f2<(fmax**2))) * self.pix2freq**2 # numerical sum (in the array's tangent circle)
integral_out = 0.0229*6*_np.pi/5 * (r0*fmax)**(-5./3.) # analytical sum (outside the array's tangent circle)
integral = integral_in + integral_out
if grad:
gg = _np.zeros((len(x0),)+self._f2.shape)
gg = _np.zeros((len(parampsd),)+self._f2.shape)
gg[0,...] = psd*(-5./3)/r0
gg[1,...] = psd*(-11./6)/((1./Lext**2.) + self._f2)*(-2/Lext**3)
# compute integral gradient
integral_g = [0,0]
for i in range(2):
integral_g[i] = _np.sum(gg[i,...]*(self._f2<(fmax**2))) * self._pix2freq**2 # numerical sum (in the array's tangent circle)
integral_g[i] = _np.sum(gg[i,...]*(self._f2<(fmax**2))) * self.pix2freq**2 # numerical sum (in the array's tangent circle)
integral_g[0] += integral_out*(-5./3)/r0
if grad:
......@@ -867,6 +870,8 @@ class Psfao(ParametricPSFfromPSD):
Define a fixed oversampling factor
"""
self.Lext = Lext
self._maskin = None
self._vk = None
super().__init__(7,npix,system=system,samp=samp,fixed_k=fixed_k)
# r0,C,A,alpha,ratio,theta,beta
......@@ -881,6 +886,7 @@ class Psfao(ParametricPSFfromPSD):
@property
def cutoffAOfreq(self):
"""AO cutoff frequency due to limited number of corrected modes"""
return self.system.Nact/(2.0*self.system.D)
......@@ -890,13 +896,13 @@ class Psfao(ParametricPSFfromPSD):
self._vk = (1-self._maskin) * 0.0229 * ((1. / self.Lext**2.) + self._f2)**(-11./6.)
def psd(self, x0, grad=False):
def psd(self, parampsd, grad=False):
"""Compute the PSD model from parameters
PSD is given in [rad²/f²] = [rad² m²]
Parameters
----------
x0 : numpy.array (dim=1), tuple, list
parampsd : numpy.array (dim=1), tuple, list