Skip to content
GitLab
Commit 7df5aa50 authored by rfetick's avatar rfetick
Browse files

Compute gradients

parent 31dc34f4
Hide whitespace changes
Inline Side-by-side
maoppy/psfmodel.py
View file @ 7df5aa50
......@@ -476,7 +476,7 @@ class ParametricPSFfromPSD(ParametricPSF):
_,sig2 = self.psd(x0)
return _np.exp(-sig2)
def psd(self,x0):
def psd(self,x0,grad=False):
raise ValueError("ParametricPSFfromPSD is not to be instantiated. the `psd` method must be override in the subclasses")
def check_parameters(self,x0):
......@@ -488,7 +488,7 @@ class ParametricPSFfromPSD(ParametricPSF):
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 otf(self,x0,dx=0,dy=0,_caller='user'):
def otf(self,x0,dx=0,dy=0,_caller='user',grad=False):
"""
See __call__ for input arguments
Warning: result of otf will be unconsistent if undersampled!!!
......@@ -499,19 +499,39 @@ class ParametricPSFfromPSD(ParametricPSF):
if (self._k > 1) and (_caller != 'self'):
raise ValueError("Cannot call `otf(...)` when undersampled (functionality not implemented yet)")
otfTurbulent = self._otfTurbulent(x0)
otfDiffraction = self._otfDiffraction()
otfDiffr = self._otfDiffraction()
otfShift = self._otfShift(dx,dy)
if grad:
otfTurb,g = self._otfTurbulent(x0,grad=True)
for i in range(g.shape[0]): g[i,:] *= otfDiffr*otfShift
else:
otfTurb = self._otfTurbulent(x0,grad=False)
return otfTurbulent * otfDiffraction * otfShift
otf = otfTurb * otfDiffr * otfShift
if grad: return otf,g
return otf
def _otfTurbulent(self,x0):
PSD,integral = self.psd(x0)
def _otfTurbulent(self,x0,grad=False):
L = self.system.D * self._samp_over
if grad:
PSD,integral,g = self.psd(x0,grad=True)
else:
PSD,integral = self.psd(x0,grad=False)
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.)
#Bphi = _fft.fftshift(_np.real(Bg[0, 0] - Bg)) # normalized on the numerical FoV
Bphi = _fft.fftshift(_np.real(integral - Bg)) # normalized up to infinity
otf = _np.exp(-Bphi)
if grad:
g2 = _np.zeros(g.shape,dtype=complex)
for i in range(len(x0)):
Bg = _fft.fft2(_fft.fftshift(g[i,...])) / L**2
# Warning: grad is normalized on the FoV, that is different of the otf!
Bphi = _fft.fftshift(_np.real(Bg[0, 0] - Bg)) # normalized on the numerical FoV
g2[i,...] = -Bphi*otf
return otf, g2
return otf
#@_lru_cache(maxsize=2)
def _otfDiffraction(self):
......@@ -528,7 +548,7 @@ class ParametricPSFfromPSD(ParametricPSF):
ny,nx = X.shape
return _np.exp(-2j*_np.pi*self._k*(dx*X/nx + dy*Y/ny))
def __call__(self,x0,dx=0,dy=0):
def __call__(self,x0,dx=0,dy=0,grad=False):
"""
Parameters
----------
......@@ -540,12 +560,26 @@ class ParametricPSFfromPSD(ParametricPSF):
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 grad:
otf,g = self.otf(x0,dx=dx,dy=dy,_caller='self',grad=True)
for i in range(len(x0)):
g[i,...] = _np.real(_fft.fftshift(_fft.ifft2(_fft.fftshift(g[i,...]))))
else:
otf = self.otf(x0,dx=dx,dy=dy,_caller='self',grad=False)
if self._k==1:
return out
psf = _np.real(_fft.fftshift(_fft.ifft2(_fft.fftshift(otf))))
k = int(self._k)
if k==1:
if grad: return psf,g
return psf
else:
return _binning(out,int(self._k)) # undersample PSF if needed (if it was oversampled for computations)
if grad:
g2 = _np.zeros((len(x0),psf.shape[0]//k,psf.shape[1]//k))
for i in range(len(x0)):
g2[i,...] = _binning(g[i,...],k)
return _binning(psf,k), g2
return _binning(psf,k) # undersample PSF if needed (if it was oversampled for computations)
#%% TURBULENT PSF
class Turbulent(ParametricPSFfromPSD):
......@@ -579,7 +613,7 @@ class Turbulent(ParametricPSFfromPSD):
bounds_up = [_np.inf, _np.inf]
self.bounds = (bounds_down,bounds_up)
def psd(self,x0):
def psd(self,x0,grad=False):
"""Compute the PSD model from parameters
PSD is given in [rad²/f²] = [rad² m²]
......@@ -612,6 +646,12 @@ class Turbulent(ParametricPSFfromPSD):
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:
g = _np.zeros((len(x0),)+F2.shape)
g[0,...] = PSD*(-5./3)/r0
g[1,...] = PSD*(-11./6)/((1./Lext**2.) + F2)*(-2/Lext**3)
if grad: return PSD,integral,g
return PSD, integral
#%% PSFAO MODEL
......@@ -659,7 +699,7 @@ class Psfao(ParametricPSFfromPSD):
bounds_up = [_np.inf for i in range(7)]
self.bounds = (bounds_down,bounds_up)
def psd(self,x0):
def psd(self,x0,grad=False):
"""Compute the PSD model from parameters
PSD is given in [rad²/f²] = [rad² m²]
......@@ -667,6 +707,9 @@ class Psfao(ParametricPSFfromPSD):
----------
x0 : numpy.array (dim=1), tuple, list
See __doc__ for more details
grad : bool (default=False)
Return both (psd,integral,gradient) if set to True
Warning: not finished yet, do not use!
Returns
-------
......@@ -677,27 +720,26 @@ class Psfao(ParametricPSFfromPSD):
self.check_parameters(x0)
nx0,ny0 = self._nullFreqIndex
f2D = self._fxfy
F2 = f2D[0]**2. + f2D[1]**2.
Fao = self.system.Nact/(2.0*self.system.D)
maskin = (F2 < Fao**2.)
r0,C,A,alpha,ratio,theta,beta = x0
PSD = 0.0229* r0**(-5./3.) * ((1. / self.Lext**2.) + F2)**(-11./6.)
PSD *= (F2 >= Fao**2.)
PSD *= (1-maskin)
ax = alpha*ratio
ay = alpha/ratio
param = (ax,ay,theta,beta,0,0)
removeInside = (1+_np.sqrt(2))/2 * self._pix2freq/2 # remove central pixel in energy computation
moff = moffat(f2D,param,norm=Fao,removeInside=removeInside) * (F2 < Fao**2.)
moff = moffat(f2D,param,norm=Fao,removeInside=removeInside) * maskin
moff[nx0,ny0] = 0.0 # Set Moffat PSD = 0 at null frequency
moff = moff / (_np.sum(moff)*self._pix2freq**2) # normalize moffat numerically to get correct A=sigma² in the AO zone
# Warning: Moffat numerical normalization generates strehlOTF jump when "_k" is changed
PSD += (F2 < Fao**2.) * _np.abs(C + A*moff) # before I wrote "moffat(f2D,param,norm=Fao)" instead of "moff", so no numerical normalization
PSD += maskin * (C + A*moff) # before I wrote "moffat(f2D,param,norm=Fao)" instead of "moff", so no numerical normalization
PSD[nx0,ny0] = 0.0 # Set PSD = 0 at null frequency
# Compute PSD integral up to infinity
......@@ -706,6 +748,22 @@ class Psfao(ParametricPSFfromPSD):
integral_out = 0.0229*6*_np.pi/5 * (r0*fmax)**(-5./3.) # analytical sum
integral = integral_in + integral_out
if grad:
g = _np.zeros((len(x0),)+F2.shape)
# derivative towards r0
g[0,...] = PSD*(1-maskin)*(-5./3)/r0
# derivative towards C
g[1,...] = maskin
# derivative towards A
g[2,...] = maskin*moff
# derivative towards alpha
# derivative towards ratio
# derivative towards theta
# derivative towards beta
# Remove central freq from all derivative
g[:,nx0,ny0] = 0
if grad: return PSD,integral,g
return PSD, integral
def tofits(self,param,filename,*args,keys=None,overwrite=False,**kwargs):
......
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