Skip to content

GitLab

Commit 9d11e206 authored by rfetick's avatar rfetick
Browse files

Merge branch 'dev-rfetick'

parents ede2e2e9 b3de27f8
Hide whitespace changes
Inline Side-by-side
maoppy/example/psfao_derivative.py 0 → 100644
View file @ 9d11e206
# -*- coding: utf-8 -*-
"""
Created on Tue Jun 1 10:20:53 2021
Show the analytical and numerical derivatives for Psfao
@author: rfetick
"""
import matplotlib.pyplot as plt
import numpy as np
import copy
from maoppy.psfmodel import Psfao
from maoppy.instrument import muse_nfm
Npix = 64 # pixel size of PSF
wvl = 600*1e-9 # wavelength [m]
#%% Initialize PSF model
samp = muse_nfm.samp(wvl) # sampling (2.0 for Shannon-Nyquist)
Pmodel = Psfao((Npix,Npix),system=muse_nfm,samp=samp)
#%% Choose parameters
r0 = 0.15 # Fried parameter [m]
bck = 1e-5 # Phase PSD background [rad² m²]
amp = 5.0 # Phase PSD Moffat amplitude [rad²]
alpha = 0.8 # Phase PSD Moffat alpha [1/m]
ratio = 1.2
theta = np.pi/4
beta = 1.6 # Phase PSD Moffat beta power law
param = [r0,bck,amp,alpha,ratio,theta,beta]
#%% Compute PSD derivative
psd,_,g,_ = Pmodel.psd(param,grad=True)
gnum = 0*g
eps = 1e-5
for i in range(len(param)):
dp = copy.copy(param)
dp[i] += eps
psd2,_ = Pmodel.psd(dp)
gnum[i,...] = (psd2-psd)/eps
#%% Plot results
names = ['r0','bck','amp','alpha','ratio','theta','beta']
plt.figure(1)
plt.clf()
for i in range(len(param)):
plt.subplot(2,len(param),i+1)
plt.pcolormesh(g[i,...])
plt.axis('image')
plt.axis('off')
plt.title(names[i])
plt.colorbar(orientation='horizontal')
plt.subplot(2,len(param),i+1+len(param))
plt.pcolormesh(gnum[i,...])
plt.axis('image')
plt.axis('off')
#plt.title(names[i])
plt.colorbar(orientation='horizontal')
#%% Compute PSF derivative
psf,g = Pmodel(param,grad=True)
gnum = 0*g
eps = 1e-5
for i in range(len(param)):
dp = copy.copy(param)
dp[i] += eps
psf2 = Pmodel(dp)
gnum[i,...] = (psf2-psf)/eps
#%% Plot results
names = ['r0','bck','amp','alpha','ratio','theta','beta']
plt.figure(2)
plt.clf()
for i in range(len(param)):
plt.subplot(2,len(param),i+1)
plt.pcolormesh(g[i,...])
plt.axis('image')
plt.axis('off')
plt.title(names[i])
plt.colorbar(orientation='horizontal')
plt.subplot(2,len(param),i+1+len(param))
plt.pcolormesh(gnum[i,...])
plt.axis('image')
plt.axis('off')
#plt.title(names[i])
plt.colorbar(orientation='horizontal')
\ No newline at end of file
maoppy/example/turbulent_derivative.py 0 → 100644
View file @ 9d11e206
# -*- coding: utf-8 -*-
"""
Created on Tue Jun 1 10:56:24 2021
@author: rfetick
"""
import matplotlib.pyplot as plt
import numpy as np
import copy
from maoppy.psfmodel import Turbulent
from maoppy.instrument import muse_nfm
Npix = 128 # pixel size of PSF
wvl = 600*1e-9 # wavelength [m]
#%% Initialize PSF model
samp = muse_nfm.samp(wvl) # sampling (2.0 for Shannon-Nyquist)
Pmodel = Turbulent((Npix,Npix),system=muse_nfm,samp=samp)
#%% Choose parameters
r0 = 0.15 # Fried parameter [m]
L0 = 20 # external length [m]
param = [r0,L0]
#%% Compute PSD derivative
psd,_,g,_ = Pmodel.psd(param,grad=True)
gnum = 0*g
eps = 1e-5
for i in range(len(param)):
dp = copy.copy(param)
dp[i] += eps
psd2,_ = Pmodel.psd(dp)
gnum[i,...] = (psd2-psd)/eps
#%% Plot results
names = ['r0','L0']
plt.figure(1)
plt.clf()
for i in range(len(param)):
plt.subplot(2,len(param),i+1)
plt.pcolormesh(np.arcsinh(g[i,...]))
plt.axis('image')
plt.axis('off')
plt.title(names[i])
plt.colorbar(orientation='horizontal')
plt.subplot(2,len(param),i+1+len(param))
plt.pcolormesh(np.arcsinh(gnum[i,...]))
plt.axis('image')
plt.axis('off')
#plt.title(names[i])
plt.colorbar(orientation='horizontal')
#%% Compute PSF derivative
psf,g = Pmodel(param,grad=True)
gnum = 0*g
eps = 1e-5
for i in range(len(param)):
dp = copy.copy(param)
dp[i] += eps
psf2 = Pmodel(dp)
gnum[i,...] = (psf2-psf)/eps
#%% Plot results
names = ['r0','L0']
plt.figure(2)
plt.clf()
for i in range(len(param)):
plt.subplot(2,len(param),i+1)
plt.pcolormesh(g[i,...])
plt.axis('image')
plt.axis('off')
plt.title(names[i])
plt.colorbar(orientation='horizontal')
plt.subplot(2,len(param),i+1+len(param))
plt.pcolormesh(gnum[i,...])
plt.axis('image')
plt.axis('off')
#plt.title(names[i])
plt.colorbar(orientation='horizontal')
\ No newline at end of file
maoppy/psfmodel.py
View file @ 9d11e206
......@@ -203,23 +203,42 @@ def psffit(psf,Model,x0,weights=None,dxdy=(0.,0.),flux_bck=(True,True),
return result
#%% BASIC FUNCTIONS
def reduced_coord(XY,ax,ay,theta,cx,cy):
def reduced_coord(XY,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
Ryy = (c/ay)**2 + (s/ax)**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
u = Rxx*uxx + Rxy*uxy + Ryy*uyy
if grad:
g = _np.zeros((3,)+u.shape)
# du/dax
g[0,...] = -2/ax * (uxx*(c/ax)**2 - uxy*s2/ax**2 + uyy*(s/ax)**2)
# du/day
g[1,...] = -2/ay * (uxx*(s/ay)**2 + uxy*s2/ay**2 + uyy*(c/ay)**2)
# du/dtheta
drxx = -2*c*s/ax**2 + 2*s*c/ay**2
drxy = 2*_np.cos(2*theta)*(1/ay**2 - 1/ax**2)
dryy = -2*c*s/ay**2 + 2*s*c/ax**2
g[2,...] = drxx*uxx + drxy*uxy + dryy*uyy
return u,g
u = Rxx*(XY[0]-cx)**2 + Rxy*(XY[0]-cx)*(XY[1]-cy) + Ryy*(XY[1]-cy)**2
return u
def moffat(XY, param, norm=None, removeInside=0):
"""
Compute a Moffat function on a meshgrid
moff = Enorm * (1+u)^(-beta)
def moffat(XY, param, norm=None, removeInside=0, grad=False):
"""Compute a Moffat function on a meshgrid
```moff = E * (1+u)^(-beta)```
with `u` the quadratic coordinates in the shifted and rotated frame
and `E` the energy normalization factor
Parameters
----------
......@@ -235,14 +254,14 @@ def moffat(XY, param, norm=None, removeInside=0):
Keywords
--------
norm : None, _np.inf, float (>0), int (>0)
norm : None, np.inf, float (>0), int (>0)
Radius for energy normalization
None - No energy normalization (maximum=1.0)
Enorm = 1.0
_np.inf - Total energy normalization (on the whole X-Y plane)
Enorm = (beta-1)/(pi*ax*ay)
E = 1.0
np.inf - Total energy normalization (on the whole X-Y plane)
E = (beta-1)/(pi*ax*ay)
float,int - Energy normalization up to the radius defined by this value
Enorm = (beta-1)/(pi*ax*ay)*(1-(1+(R**2)/(ax*ay))**(1-beta))
E = (beta-1)/(pi*ax*ay)*(1-(1+(R**2)/(ax*ay))**(1-beta))
removeInside: float (default=0)
Used to remove the central pixel in energy computation
......@@ -250,19 +269,43 @@ def moffat(XY, param, norm=None, removeInside=0):
if len(param)!=6: raise ValueError("Parameter `param` must contain exactly 6 elements, but input has %u elements"%len(param))
ax,ay,theta,beta,cx,cy = param
u = reduced_coord(XY,ax,ay,theta,cx,cy)
if grad:
u,du = reduced_coord(XY,ax,ay,theta,cx,cy,grad=True)
else:
u = reduced_coord(XY,ax,ay,theta,cx,cy,grad=False)
V = (1.0+u)**(-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)
if norm is None:
Enorm = 1
if grad:
return V,dV
else: # norm can be float or np.inf
if (beta<=1) and (norm==_np.inf): raise ValueError("Cannot compute Moffat energy for beta<=1")
if beta==1: raise ValueError("Energy computation for beta=1.0 not implemented yet. Sorry!")
Enorm = (beta-1) / (_np.pi*ax*ay)
Fout = (1 - (1 + (norm**2) / (ax*ay))**(1-beta))
Fin = (1 - (1 + (removeInside**2) / (ax*ay))**(1-beta))
Enorm = Enorm / (Fout-Fin)
return Enorm * (1.0+u)**(-beta)
E = (beta-1) / (_np.pi*ax*ay)
Fout = (1 + (norm**2)/(ax*ay))**(1-beta)
Fin = (1 + (removeInside**2)/(ax*ay))**(1-beta)
F = 1/(Fin-Fout)
if grad:
dE = [-E/ax,-E/ay,0,E/(beta-1)]
k = (1-beta)*Fout/(1 + (norm**2)/(ax*ay))*(norm**2)/(ax*ay)
dFout = _np.array([-k/ax,-k/ay,0,-Fout*_np.log(1 + (norm**2)/(ax*ay))]) if norm<_np.inf else _np.zeros(4)
k = (1-beta)*Fin/(1 + (removeInside**2)/(ax*ay))*(removeInside**2)/(ax*ay)
dFin = _np.array([-k/ax,-k/ay,0,-Fin*_np.log(1 + (removeInside**2)/(ax*ay))])
dF = -(dFin-dFout)/(Fin-Fout)**2
dm = 0*dV
for i in range(4): dm[i,...] = V*E*dF[i] + V*dE[i]*F + dV[i,...]*E*F
return V*E*F, dm
return V*E*F
def gauss(XY,param):
"""
......@@ -491,7 +534,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):
......@@ -503,7 +546,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!!!
......@@ -514,26 +557,46 @@ 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,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) # I cannot override 'g' here due to float to complex type
for i in range(len(x0)):
Bg = _fft.fft2(_fft.fftshift(g[i,...])) / L**2
#Bphi = _fft.fftshift(_np.real(Bg[0, 0] - Bg)) # normalized on the numerical FoV
Bphi = _fft.fftshift(_np.real(integral_g[i] - Bg)) # normalized up to infinity
g2[i,...] = -Bphi*otf
return otf, g2
return otf
#@_lru_cache(maxsize=2)
def _otfDiffraction(self):
nx, ny = self._nxnyOver
NpupX = _np.ceil(nx/self._samp_over)
NpupY = _np.ceil(ny/self._samp_over)
tab = _np.zeros((nx, ny), dtype=_np.complex)
tab = _np.zeros((nx, ny), dtype=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))
......@@ -543,7 +606,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
----------
......@@ -555,12 +618,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,...].astype(float),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):
......@@ -594,7 +671,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²]
......@@ -627,6 +704,17 @@ 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)
# compute integral gradient
integral_g = [0,0]
for i in range(2): integral_g[i] = _np.sum(g[i,...]*(F2<(fmax**2))) * self._pix2freq**2 # numerical sum (in the array's tangent circle)
integral_g[0] += integral_out*(-5./3)/r0
if grad: return PSD,integral,g,integral_g
return PSD, integral
#%% PSFAO MODEL
......@@ -669,12 +757,16 @@ class Psfao(ParametricPSFfromPSD):
super().__init__(7,npix,system=system,samp=samp,fixed_k=fixed_k)
self.Lext = Lext
# 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)]
# r0,C,A,alpha,ratio,theta,beta
### Mathematical bounds
#bounds_down = [_EPSILON,0,0,_EPSILON,_EPSILON,-_np.inf,1+_EPSILON]
#bounds_up = [_np.inf for i in range(7)]
### Physical bounds
bounds_down = [1e-3,0,0,1e-3,1e-2,-_np.inf,1.01]
bounds_up = [_np.inf]*4 + [1e2,_np.inf,5]
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²]
......@@ -682,6 +774,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
-------
......@@ -692,36 +787,68 @@ class Psfao(ParametricPSFfromPSD):
self.check_parameters(x0)
nx0,ny0 = self._nullFreqIndex
f2D = self._fxfy
pix = self._pix2freq
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[nx0,ny0] = 0.0 # Set Moffat PSD = 0 at null frequency
# Newly added for the PSFAO19 model
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
removeInside = (1+_np.sqrt(2))/2 * pix/2 # remove central pixel in energy computation
if grad:
moff,dm = moffat(f2D,param,norm=Fao,removeInside=removeInside,grad=True)
else:
moff = moffat(f2D,param,norm=Fao,removeInside=removeInside)
moff *= maskin
numericalNorm = False # set to true in order to activate code below
if numericalNorm:
moff[nx0,ny0] = 0.0 # Set Moffat PSD = 0 at null frequency
moff = moff / (_np.sum(moff)*pix**2) # normalize moffat numerically to get correct A=sigma² in the AO zone
if grad: raise ValueError("PSFAO analytical gradient computation is not compatible with numerical Moffat normalization")
# 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
fmax = _np.min([nx0,ny0])*self._pix2freq
integral_in = _np.sum(PSD*(F2<(fmax**2))) * self._pix2freq**2 # numerical sum
fmax = _np.min([nx0,ny0])*pix
integral_in = _np.sum(PSD*(F2<(fmax**2))) * pix**2 # numerical sum
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
g[3,...] = A*maskin*(dm[0,...]*ratio + dm[1,...]/ratio)
# derivative towards ratio
g[4,...] = A*maskin*(dm[0,...]*alpha - dm[1,...]*ay/ratio)
# derivative towards theta
g[5,...] = A*maskin*dm[2,...]
# derivative towards beta
g[6,...] = A*maskin*dm[3,...]
# Remove central freq from all derivative
g[:,nx0,ny0] = 0
# Compute integral gradient
integral_g = _np.zeros(len(x0))
for i in range(len(x0)): integral_g[i] = _np.sum(g[i,...]*(F2<(fmax**2))) * self._pix2freq**2 # numerical sum
integral_g[0] += integral_out*(-5./3)/r0
if grad: return PSD,integral,g,integral_g
return PSD, integral
def tofits(self,param,filename,*args,keys=None,overwrite=False,**kwargs):
......
maoppy/test/test_psfmodel.py
View file @ 9d11e206
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri May 8 20:21:20 2020
Created on Tue Jun 1 16:30:27 2021
@author: rfetick
"""
import unittest
import pytest
import numpy as np
from maoppy.psfmodel import lsq_flux_bck, moffat, gauss, Psfao
from maoppy.instrument import zimpol, muse_nfm
import copy
from maoppy.psfmodel import (reduced_coord, lsq_flux_bck, moffat, gauss,
ParametricPSFfromPSD, Psfao)
from maoppy.instrument import muse_nfm, zimpol
class TestLsqFluxBck(unittest.TestCase):
def test_fluxbck(self):
bck = -5.0
flux = 8.0
model = np.zeros((10,10))
model[3,3] = 1.0
image = flux*model+bck
w = np.ones_like(image)
f,b = lsq_flux_bck(model,image,w)
self.assertAlmostEqual(flux,f,delta=1e-8)
self.assertAlmostEqual(bck,b,delta=1e-8)
def test_positivebck(self):
bck = -2.0
flux = 8.0