normalize gradient r0 up to infinity

maoppy/example/psfao_derivative.py

+# -*- 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 = 1 # 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/psfmodel.py

@@ -514,7 +514,7 @@ class ParametricPSFfromPSD(ParametricPSF):
     def _otfTurbulent(self,x0,grad=False):
         L = self.system.D * self._samp_over
         if grad:
-            PSD,integral,g = self.psd(x0,grad=True)
+            PSD,integral,g,integral_g = self.psd(x0,grad=True)
         else:
             PSD,integral = self.psd(x0,grad=False)
             
@@ -527,8 +527,8 @@ class ParametricPSFfromPSD(ParametricPSF):
             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
-                # BUG: 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
+                #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
@@ -577,7 +577,7 @@ class ParametricPSFfromPSD(ParametricPSF):
             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)
+                    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)
 
@@ -646,14 +646,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
         
-        # TODO: I can also compute the integral gradient
-        
         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
+        if grad: return PSD,integral,g,integral_g
         return PSD, integral
 
 #%% PSFAO MODEL        
@@ -765,9 +768,12 @@ class Psfao(ParametricPSFfromPSD):
             # Remove central freq from all derivative
             g[:,nx0,ny0] = 0
         
-        #TODO: I can also compute the integral gradient towards the parameters!
-        
-        if grad: return PSD,integral,g
+            # 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):