added parallel computation

Author: kahil_dell

README.md

@@ -48,6 +48,8 @@ The modules of this class:
 
 ## The `patch_pd` class:
 For fitting the wavefront error is sub-regions of the full FOV of the PD dataset. The user can enter the size of the subregion (it has to be quadratic) plus the number of Zernike polynomials to be fit and the names of the output files (one for the wavefront error and one for the 2D MTF).
+
+The class offers the option to run parallel computation by setting the parser -p to True (see below).
 # How to use the code?
 The `minimization` class returns the best-fit zernike polynomials, a visualisation of the results (wavefront error+MTF), and the restored scene (in this case the focused image of the PD dataset). To get these specific results, type in the shell terminal:
 ```
@@ -59,7 +61,7 @@ The specific description of the parsers can be found inside [the main code](http
 To use the `patch_pd` class:
 
 ```
-python3 patch_pd.py -i 'path/input.fits' -z 10 -d 265 -ow 'path/output_wf.fits' -om 'path/output_wf.fits'
+python3 patch_pd.py -i 'path/input.fits' -z 10 -d 265 -ow 'path/output_wf.fits' -om 'path/output_wf.fits -p True
 
 ```
 The parsers description can be found in the [main code](https://github.com/fakahil/PyPD/blob/master/patch_pd.py)

__pycache__/cost_func.cpython-38.pyc

@@ -51,4 +51,6 @@ def Minimize_res(self,coefficients):
 
 
        return t0,ph
+
+
 

__pycache__/logging.cpython-38.pyc

@@ -22,10 +22,12 @@
 from wavefront import *
 from deconvolution import *
 from cost_func import *
+import os
+import processing
 
 
 class patch_pd(object):
-    def __init__(self,pd_data,Del,co_num,output_wf,output_mtf):
+    def __init__(self,pd_data,Del,co_num,output_wf,output_mtf,parallel=True):
 
 
      self.data= fits.getdata(pd_data)
@@ -42,15 +44,56 @@ def __init__(self,pd_data,Del,co_num,output_wf,output_mtf):
      self.Imk = self.data[1,:,:]/self.mean_imk
 
      self.output_WF = np.zeros((2048,2048))
-     self.output_mtf = np.zeros((2048,2048))
-
-
+     self.output_MTF = np.zeros((2048,2048))
+     self.parallel = parallel
+    #@staticmethod
+    def run_pd(self,k):
+       
+     im0 = self.patch[k,:,:,0]
+     imk  =self.patch[k,:,:,1]
+     im0 = im0/im0.mean()
+     imk = imk/imk.mean()
+     imk = imreg(im0,imk)
+
+     
+     im0 = apo2d(im0,10)
+     imk = apo2d(imk,10)
+ 
+     d0,dk = FT(im0,imk)
+     gam =1# Gamma(d0,dk,M_gamma)
+     p0 =   np.zeros(self.co_num)
+     fit = cost_func(self.Del,0.5,1e-10,10,self.co_num,0.5,617.3e-6, 140,4125.3,0.5)
+     Mask = aperture.mask_pupil(fit.telescope.pupil_size(),fit.size)
+     noise_temp = noise_mask_high(fit.size,fit.cut_off)
+     noise_filter = fftshift(noise_temp)
+     def Minimise(coefficients):
+         A_f = wavefront.pupil_foc(coefficients,fit.size,fit.telescope.pupil_size(),self.co_num)
+         A_def = wavefront.pupil_defocus(coefficients,fit.size,fit.del_z,fit.telescope.pupil_size(),self.co_num)
+         psf_foc = wavefront.PSF(Mask,A_f,False) 
+         psf_defoc = wavefront.PSF(Mask,A_def,False)
+         t0 = wavefront.OTF(psf_foc)
+         tk = wavefront.OTF(psf_defoc) 
+         q,q2 = PD.Q_matrix(t0,tk,fit.reg,gam)
+         F_m = PD.F_M(q2,d0, dk,t0,tk,noise_filter,gam)
+         E_metric = PD.Error_metric(t0,tk,d0,dk,q,noise_filter)
+         L_m = PD.L_M(E_metric,fit.size)
+         return L_m
+
+
+
+     Minimize_partial = partial(Minimise)
+     mini = minimize(Minimize_partial,p0,method='L-BFGS-B')
+     result = fit.Minimize_res(mini.x)
+     patch_wfe = result[1]
+     patch_mtf = MTF(fftshift(result[0]))
+     return patch_wfe,patch_mtf
 
     def fit_patch(self):
 
-      upper = 1700
-      Nx = np.arange(300,upper,self.Del)
-      Ny = np.arange(300,upper,self.Del)
+     upper = 1700
+     Nx = np.arange(300,upper,self.Del)
+     Ny = np.arange(300,upper,self.Del)
+     if not self.parallel:
       for n1 in Nx :
         for n2 in Ny:
 
@@ -60,9 +103,7 @@ def fit_patch(self):
             im0 = im0/im0.mean()
             imk = imk/imk.mean()
             imk = imreg(im0,imk)
-
-
-           
+         
             im0 = apo2d(im0,10)
             imk = apo2d(imk,10)
 
@@ -91,16 +132,68 @@ def Minimise(coefficients):
                   L_m = PD.L_M(E_metric,fit.size)
                   return L_m
 
-       
+
             Minimise_partial = partial(Minimise)
             mini = scipy.optimize.minimize(Minimise_partial,p0,method= 'L-BFGS-B')
 
             result = fit.Minimize_res(mini.x)
             self.output_WF[n2:n2+self.Del,n1:n1+self.Del] = result[1]
-            self.output_mtf[n2:n2+self.Del,n1:n1+self.Del] = MTF(fftshift(result[0]))
+            self.output_MTF[n2:n2+self.Del,n1:n1+self.Del] = MTF(fftshift(result[0]))
             hdu = fits.PrimaryHDU(self.output_WF)
             hdu.writeto(self.output_wf,overwrite=True)
+            hdu = fits.PrimaryHDU(self.output_mtf)
+            hdu.writeto(self.output_mtf,overwrite=True)
+     else:   
+            self.patch = tools.prepare_patches(self.data,self.Del,self.Im0,self.Imk)
+            n_workers = min(6, os.cpu_count())
+            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))
+
+    def plot_results(self,output):
+
+         # 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)
+
+
+         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')
@@ -109,7 +202,10 @@ def Minimise(coefficients):
  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']))
+ 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']))
 

__pycache__/processing.cpython-38.pyc

@@ -0,0 +1,124 @@
+import os
+from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor, Executor
+class Thread:
+    """
+    # Thread
+    The threads holds the information on the function to execute in a thread or process.
+    Provides an interface to the `future` object once submitted to an executer.
+    """
+
+    def __init__(self, func, args):
+        self.function = func
+        self.arguments = args
+        self.future = None
+
+    def submit(self, executor: Executor):
+        """Start execution via executor"""
+        if not self.is_submitted():
+            self.future = executor.submit(self.function, self.arguments)
+        return self
+
+    def is_submitted(self) -> bool:
+        return self.future is not None
+
+    def is_done(self):
+        return self.is_submitted() and self.future.done()
+
+    def exception(self):
+        if not self.is_done():
+            return None
+        return self.future.exception()
+
+    def result(self):
+        if not self.is_submitted():
+            return None
+        return self.future.result()
+
+class MP:
+    """
+    ## MP Multi-Processing
+    Class provides housekeeping / setup methods to reduce the programming overhead of
+    spawning threads or processes.
+    """
+
+    #: Number of CPUs of the current machine
+    NUM_CPUs = round(os.cpu_count() * 0.8)
+
+    @staticmethod
+    def threaded(func, args, workers=10, raise_exception=True):
+        """
+        Calls the given function in multiple threads for the set of given arguments
+            Note that this does not spawn processes, but threads. Use this for non CPU
+            CPU dependent tasks, i.e. I/O
+        Method returns once all calls are done.
+
+        ### Params
+        - func: [Function] the function to call
+        - args: [Iterable] the 'list' of arguments for each call
+        - workers: [Integer] the number of concurrent threads to use
+        - raise_exception: [Bool] Flag if an exception in a thread shall be raised or just logged
+
+        ### Returns
+        Results from all `Threads` as list
+        """
+        if len(args) == 1:
+            return list(func(arg) for arg in args)
+
+        with ThreadPoolExecutor(workers) as ex:
+            threads = [Thread(func, arg).submit(ex) for arg in args]
+        return MP.collect_results(threads, raise_exception)
+
+    @staticmethod
+    def simultaneous(func, args, workers=None, raise_exception=True):
+        """
+        Calls the given function in multiple processes for the set of given arguments
+            Note that this does spawn processes, not threads. Use this for task that
+            depend heavily on CPU and can be done in parallel.
+        Method returns once all calls are done.
+
+        ### Params
+        - func: [Function] the function to call
+        - args: [Iterable] the 'list' of arguments for each call
+        - workers: [Integer] the number of concurrent threads to use (Default: NUM_CPUs)
+        - raise_exception: [Bool] Flag if an exception in a thread shall be raised or just logged
+
+        ### Returns
+        Results from all `Threads` as list
+        """
+        if len(args) == 1:
+            return list(func(arg) for arg in args)
+
+        if workers is None:
+            workers = MP.NUM_CPUs
+        with ProcessPoolExecutor(workers) as ex:
+            threads = [Thread(func, arg).submit(ex) for arg in args]
+        return MP.collect_results(threads, raise_exception)
+
+    @staticmethod
+    def collect_results(threads: list, raise_exception: bool = True) -> list:
+        """
+        Takes a list of threads and waits for them to be executed. Collects results.
+
+        ### Params
+        - threads: [List<Thread>] a list of submitted threads
+        - raise_exception: [Bool] Flag if an exception in a thread shall be raised or just logged
+
+        ### Returns
+        Results from all `Threads` as list
+        """
+        result = []
+        while len(threads) > 0:
+            for thread in threads:
+                if not thread.is_submitted():
+                    threads.remove(thread)
+                if not thread.is_done():
+                    continue
+
+                if thread.exception() is not None:
+                    MP.__exception_handling(threads, thread, raise_exception)
+                else:
+                    result.append(thread.result())
+                threads.remove(thread)
+        return result
+
+   

__pycache__/tools.cpython-38.pyc

@@ -101,5 +101,50 @@ def plot_zernike(coeff):
         plt.ylabel('Coefficient [$\lambda$]',fontsize=18) 
         plt.title('Zernike Polynomials Coefficients',fontsize=18)
 
+def prepare_patches(d,Del,Im0,Imk):
+    n = d.shape[0]
+    upper = 1700-Del
+    lower = 300
+    Nx = np.arange(lower,upper,Del)
+    Ny = np.arange(lower,upper,Del)
+    i_max = np.floor((upper-lower)/Del)+1
+    patches = np.zeros((int(i_max**2),Del,Del,n))
+    #output_WF =np.zeros((int(i_max**2),Del,Del))
+    #output_mtf = np.zeros((int(i_max**2),Del,Del))
+    k=0
+    for n1 in Nx :
+       for n2 in Ny:
+        patches[k,:,:,0]=Im0[n2:n2+Del,n1:n1+Del]
+        patches[k,:,:,1]=Imk[n2:n2+Del,n1:n1+Del]
+        k = k+1
+    return patches
+
+
+def stitch_patches(results,Del):
+    data1 = [r[0] for r in results]
+    data2 = [r[1] for r in results]
+    upper = 1700-Del
+    lower = 300
+    Nx = np.arange(lower,upper,Del)
+    Ny = np.arange(lower,upper,Del)
+    i_max = np.floor((upper-lower)/Del)+1
+    k = 0
+    if len(data1)==(np.floor((upper-lower)/Del)+1)**2:
+
+      st_wf = np.zeros((2048,2048))
+      st_mtf = np.zeros((2048,2048))
+    else:
+        raise TypeError('Check dimensions!')
+
+    for n1 in Nx :
+         for n2 in Ny:
+             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
+    
+
+   
+