edited plotting tools

Author: kahil_dell

PD.py

@@ -5,8 +5,6 @@
 from functools import partial
 
 
-import pyfits
-
 
 
 ## function to compute the FT of focused and defocused image

README.md

@@ -32,7 +32,7 @@ This class allows the user to:
 4. Restore blurred images and correct for the PSF choosing between the Wiener filter and the Richardson-Lucy filter
 
 The parameters of this class:
-  1. the pair of focused and defocused images, 
+  1. the pair of focused and defocused images, the input image should have a format of `2xsize_xxsize_y`.
   2. the parameters of the `Telescope` class
   3. `cutoff` frequency for the noise filtering
   4.  `reg`: regularization parameter for the Wiener filter
@@ -56,7 +56,7 @@ The `minimization` class returns the best-fit zernike polynomials, a visualisati
 python3 minimization.py -i 'path/input.fits' -s 150 -w 617.3e-6 -a 140 -f 4125.3 -p 0.5 -c 0.5 -r 1e-10 -ap 10 -x1 500 -x2 650 -y1 500 -y2 650 -z 10 -del 0.5 -o path/reduced.fits -fl 'Wiener'
 ```
 
-The specific description of the parsers can be found inside [the main code](https://github.com/fakahil/PyPD/blob/master/minimization.py). The values given above are for an example PD dataset taken by the PHI/HRT telescope. You can change the values according to your telescope.
+The specific description of the parsers and input to the class can be found inside [the main code](https://github.com/fakahil/PyPD/blob/master/minimization.py). The values given above are for an example PD dataset taken by the PHI/HRT telescope. You can change the values according to your telescope.
 
 To use the `patch_pd` class:
 

deconvolution.py

@@ -5,7 +5,7 @@
 from scipy.fftpack import fftshift, ifftshift, fft2, ifft2
 from astropy.io import fits
 import imreg_dft
-import pyfits
+
 import tools
 import aperture
 import PD

minimization.py

@@ -34,15 +34,15 @@
 
 class minimization(object):
 
-   def __init__(self,foc_defoc,size,lam, diameter,focal_length,platescale,cut_off,reg,ap,x1,x2,y1,y2,co_num,del_z,output,filterr):
+   def __init__(self,foc_defoc,lam, diameter,focal_length,platescale,cut_off,reg,ap,x1,x2,y1,y2,co_num,del_z,output,filterr):
       self.data = fits.getdata(foc_defoc)
 
 
       self.foc =  self.data[0,:,:]
       self.defoc = self.data[1,:,:]
       self.cut_off = cut_off
       self.reg = reg
-      self.size = size
+      self.size =self.y2-self.y1
       self.ap = ap
       self.co_num = co_num
       self.x1 = x1
@@ -112,7 +112,7 @@ def plot_results(self,Z):
       ph =wavefront.phase(Z, self.telescope.pupil_size(),self.co_num)
 
 
-      fig = plt.figure(figsize=(10,10))
+      fig = plt.figure(figsize=(20,20))
       ax1 = fig.add_subplot(1,3,1)
       im1 = ax1.imshow(ph/(2*np.pi), origin='lower',cmap='gray')
       ax1.set_xlabel('[Pixels]',fontsize=18)
@@ -187,7 +187,6 @@ def restored_scene(self,Z,iterations_RL=10):
  parser = argparse.ArgumentParser(description='Retrieving wavefront error')
  parser.add_argument('-i','--input', help='input')
  parser.add_argument('-o','--out', help='out')
- parser.add_argument('-s','--size', help='size',default=150)
  parser.add_argument('-w','--wavelength', help='wavelength',default=617.3e-6)
  parser.add_argument('-a','--aperture', help='aperture', default=140)
  parser.add_argument('-f','--focal_length', help='focal_length',default=4125.3)
@@ -207,12 +206,13 @@ def restored_scene(self,Z,iterations_RL=10):
 
 
 
- res = minimization(foc_defoc='{0}'.format(parsed['input']), size=int(parsed['size']), lam=float(parsed['wavelength']),diameter=float(parsed['aperture']),focal_length=float(parsed['focal_length']),
+ res = minimization(foc_defoc='{0}'.format(parsed['input']), lam=float(parsed['wavelength']),diameter=float(parsed['aperture']),focal_length=float(parsed['focal_length']),
  platescale=float(parsed['plate_scale']),cut_off=float(parsed['cut_off']),reg=float(parsed['reg']),ap=int(parsed['apod']),x1=int(parsed['x1']),x2=int(parsed['x2']),y1=int(parsed['y1']),y2=int(parsed['y2']),
  co_num=int(parsed['Z']),del_z=float(parsed['del']),output='{0}'.format(parsed['out']),filterr=parsed['filter'])
 
  Z = res.fit() 
  print(Z)
+ tools.plot_zernike(Z)
  res.plot_results(Z)
  res.restored_scene(Z,10)
 

patch_pd.py

@@ -1,6 +1,5 @@
 import time
-import imreg_dft
-import pyfits
+
 import tools
 import aperture
 import PD
@@ -14,8 +13,7 @@
 import argparse
 import telescope
 from telescope import *
-
-from tools import imreg, apo2d
+from tools import *
 from aperture import *
 from PD import *
 from noise import *
@@ -144,56 +142,30 @@ def Minimise(coefficients):
             hdu = fits.PrimaryHDU(self.output_mtf)
             hdu.writeto(self.output_mtf,overwrite=True)
      else:   
+            print('Initialising parallel computation')
+            t0 = time.time()
             self.patch = tools.prepare_patches(self.data,self.Del,self.Im0,self.Imk)
             n_workers = min(6, os.cpu_count())
+            print(f'number of workers is {n_workers}')
             self.args_list = [i for i in range(len(self.patch))]
             self.results_parallel = list(processing.MP.simultaneous(self.run_pd, self.args_list, workers=n_workers))
+            dt = (time.time() - t0)/60.
+            print(f'Time spent in fitting the wavefront error is: {dt: .3f}min')
 
-    def plot_results(self,output):
+    def plot_results(self):
 
-         # change here the format of the output
          if not self.parallel:
             data_mtf = self.output_MTF
             data_wfe = self.output_WF
 
          if self.parallel:
-            ## call here the stitching function
             data_mtf,data_wfe = tools.stitch_patches(self.results_parallel,self.Del)
+            hdu = fits.PrimaryHDU(data_mtf)
+            hdu.writeto(self.output_mtf,overwrite=True)
+            hdu = fits.PrimaryHDU(data_wfe)
+            hdu.writeto(self.output_wf,overwrite=True)
+         tools.plot_mtf_wf(data_wfe,data_mtf)
 
-
-         fig, ax = plt.subplots(1,2,figsize=(10,10))
-        
-         im=ax[1].imshow(data_mtf,vmin=0,vmax=1,origin='lower',cmap='gray')
-         ax[1].set_title('MTF')
-         divider = make_axes_locatable(ax[1])
-         cax = divider.append_axes('right',pad=0.05,size=0.03)
-         cbar1 = plt.colorbar(im,cax=cax)
-         cbar1.ax.tick_params(labelsize=16)
-         cbar1.set_label('MTF',fontsize=16)
-         ax[1].set_xlabel('[Pixels]')
-         major_ticks = np.arange(0, 2048,350)
-         major_ticks_y = np.arange(0, 2048,350)
-         ax[1].set_xticks(major_ticks)
-         ax[1].set_yticks(major_ticks_y)
-         ax[1].tick_params(labelsize=4)
-
-
-
-         im2=ax[0].imshow(data_wfe/(2*np.pi),vmin=-0.5,vmax=1.4,origin='lower',cmap='gray')
-         ax[0].set_xlabel('[Pixels]')
-         ax[0].set_title('WF error[$\lambda$]')
-         divider2 = make_axes_locatable(ax[0])
-         cax2 = divider2.append_axes('right',pad=0.05,size=0.03) #size is the width of the color bar and pad is the fraction of the new axis
-         cbar2=plt.colorbar(im2,cax=cax2)
-         cbar2.ax.tick_params(labelsize=16)
-         ax[0].tick_params(labelbottom=False,labelsize=16)
-         plt.subplots_adjust(wspace=None, hspace=0.1)
-         ax[0].set_xticks(major_ticks)
-         ax[0].set_yticks(major_ticks_y)
-         ax[0].set_ylabel('[Pixels]')
-         #plt.subplots_adjust(wspace=None, hspace=-0.1)
-         #plt.axis('off')
-         plt.savefig(output,dpi=300)
 
 if (__name__ == '__main__'):
  parser = argparse.ArgumentParser(description='PD on sub-fields')
@@ -202,10 +174,10 @@ def plot_results(self,output):
  parser.add_argument('-d','--Del', help='Del',default=265)
  parser.add_argument('-ow','--ow', help='output_WFE')
  parser.add_argument('-om','--om', help='output MTF')
- parser.add_argument('-r','--res', help='results')
+
  parser.add_argument('-p','--parallel',choices=['True','False'],default=True)
  parsed = vars(parser.parse_args())
  st = patch_pd(pd_data='{0}'.format(parsed['input']),Del=int(parsed['Del']),co_num=int(parsed['Z']),output_wf='{0}'.format(parsed['ow']),output_mtf='{0}'.format(parsed['om']),parallel=bool(parsed['parallel']))
  st.fit_patch()
- st.plot_results(output='{0}'.format(parsed['res']))
+ st.plot_results()
 

tools.py

@@ -1,13 +1,15 @@
 import numpy as np
+import pylab
 import matplotlib.pyplot as plt
 import scipy
 from scipy.fftpack import fftshift, ifftshift, fft2, ifft2
 from mpl_toolkits.axes_grid1 import make_axes_locatable, axes_size
 from scipy.signal import correlate2d as correlate
 from scipy.signal import general_gaussian
 from scipy import ndimage
-
-
+import imreg_dft as ird
+from image_registration import chi2_shift
+from image_registration.fft_tools import shift
 
 
 def GetPSD1D(psd2D):
@@ -72,9 +74,7 @@ def strehl(rms):
     return np.exp(-2*(np.pi*rms**2))
 
 def imreg(im0,imk):
-    import imreg_dft as ird
-    from image_registration import chi2_shift
-    from image_registration.fft_tools import shift
+
     xoff, yoff, exoff, eyoff = chi2_shift(im0,imk)
     timg = ird.transform_img(imk, tvec=np.array([-yoff,-xoff]))
     return timg
@@ -87,7 +87,6 @@ def noise(im):
     s = estimate_sigma(im)
     return s
 
-
 def plot_zernike(coeff): 
     n = coeff.shape[0] 
     index = np.arange(n) 
@@ -100,10 +99,11 @@ def plot_zernike(coeff):
         plt.xlabel('Zernike Polynomials',fontsize=18) 
         plt.ylabel('Coefficient [$\lambda$]',fontsize=18) 
         plt.title('Zernike Polynomials Coefficients',fontsize=18)
+        plt.savefig('Zernikes.png',dpi=300)
 
 def prepare_patches(d,Del,Im0,Imk):
     n = d.shape[0]
-    upper = 1700-Del
+    upper = 1700
     lower = 300
     Nx = np.arange(lower,upper,Del)
     Ny = np.arange(lower,upper,Del)
@@ -141,10 +141,58 @@ def stitch_patches(results,Del):
              st_wf[n2:n2+Del,n1:n1+Del] = data1[k]
              st_mtf[n2:n2+Del,n1:n1+Del] = data2[k]
              k=k+1
-    return st_wf,st_mtf
+    return st_mtf,st_wf
 
 
 
+def plot_mtf_wf(ph,mtf):
+   
+    fig=plt.figure(figsize=(20,8))
+    aspect = 5
+    pad_fraction = 0.5
+    ax = fig.add_subplot(1,2,1)
+    im=ax.imshow(ph/(2*np.pi), cmap=pylab.gray(),origin='lower',vmin=-1.2,vmax=1.2)
+  
+    ax.set_xlabel('[Pixels]',fontsize=18)
+    ax.set_ylabel('[Pixels]',fontsize=18)
+    divider = make_axes_locatable(ax)
+    cax = divider.append_axes("right", size=0.15, pad=0.05)
+    cbar = plt.colorbar(im, cax=cax,orientation='vertical')
+    cbar.set_label('WF error HRT [$\lambda$]',fontsize=20)
+    cax.tick_params(labelsize=14)
+    ax2 = fig.add_subplot(1,2,2)
+
+    im2=ax2.imshow(mtf,cmap=pylab.gray(),origin='lower',vmin=0,vmax=1)
+    ax2.set_xlabel('[Pixels]',fontsize=18)
+    divider = make_axes_locatable(ax2)
+    cax2 = divider.append_axes("right", size=0.15, pad=0.05)
+    cbar2 = plt.colorbar(im2, cax=cax2,orientation='vertical')
+    cax2.tick_params(labelsize=14)
+    cbar2.set_label('MTF',fontsize=16)
+
+    plt.subplots_adjust(wspace=.2, hspace=None)
+    plt.savefig('WFE+MTF.png',dpi=300)
+    
 
+def compute_residual_shifts(pd_pair,Del):
 
+    d = fits.getdata(pd_pair)
+    xoff, yoff, exoff, eyoff = chi2_shift(d[0,500:1000,500:1000],d[1,500:1000,500:1000])
+    Imk = ird.transform_img(d[1,:,:], tvec=np.array([-yoff,-xoff]))
+    Nx = np.arange(200,1800,Del)
+    Ny = np.arange(200,1800,Del)
+    shifts_x = np.zeros((2048,2048))
+    shifts_y = np.zeros((2048,2048))
+    S_x = []
+    S_y = []
 
+    for n1 in Nx :
+        for n2 in Ny:
+            
+            im0 = Im0[n2:n2+Del,n1:n1+Del]
+            imk = Imk[n2:n2+Del,n1:n1+Del]
+            xoff, yoff, exoff, eyoff = chi2_shift(im0,imk)
+            print(xoff, yoff)
+            shifts_x[n2:n2+Del,n1:n1+Del] = xoff
+            shifts_y[n2:n2+Del,n1:n1+Del] = yoff
+            return shifts_x, shifts_y