initial commit

ptycho/python/ptycho_recon/__init__.py

@@ -0,0 +1,2 @@
+#!usr/bin/python
+#import reconstruction functions

ptycho/python/ptycho_recon/fft_shift.py

@@ -0,0 +1,26 @@
+from numpy import *
+from numpy.fft import *
+import pyfftw
+
+class shift_fftw:
+      """Shift function via FFT"""
+      def __init__(self, dk_x, dk_y, N):
+          kx = linspace(-floor(N/2.0),ceil(N/2.0)-1,N)
+          kx = ifftshift(kx)
+          [self.kX, self.kY] = meshgrid(kx,kx)
+          self.kX = self.kX*dk_x
+          self.kY = self.kY*dk_y
+          self.f = pyfftw.empty_aligned((N,N),dtype='complex128',n=16) 
+          self.r = pyfftw.empty_aligned((N,N),dtype='complex128',n=16)
+          self.N_tot = N*N
+          self.fft_forward = pyfftw.FFTW(self.r, self.f, axes=(0,1))
+          self.fft_inverse = pyfftw.FFTW(self.f, self.r, direction='FFTW_BACKWARD', axes=(0,1))
+
+      def shift(self, func, px, py):
+          self.r[:,:] = ifftshift(func)
+          self.fft_forward.update_arrays(self.r, self.f)
+          self.fft_forward.execute()
+          self.f = self.f*exp(-2*pi*1j*px*self.kX)*exp(-2*pi*1j*py*self.kY)
+          self.fft_inverse.update_arrays(self.f, self.r)
+          self.fft_inverse.execute();
+          return fftshift(self.r) / self.N_tot #fix normalization

ptycho/python/ptycho_recon/filters.py

@@ -0,0 +1,52 @@
+from numpy import *
+from numpy.fft import *
+
+def disk(N,radius):
+    output = ones((N,N))
+    x = linspace(-floor(N/2.0),ceil(N/2.0)-1,N)
+    [X,Y] = meshgrid(x,x)
+    S = sqrt(X**2+Y**2)
+    output[S>radius] = 0
+    return output
+
+def gaussian(N,innerRadius,outerRadius,sigma):
+    x = linspace(-floor(N/2.0),ceil(N/2.0)-1,N)
+    [X,Y] = meshgrid(x,x)
+    S = sqrt(X**2+Y**2)
+
+    output = exp(-(S-innerRadius)**2/(2*sigma**2))
+    output[S<=innerRadius] = 1
+    if outerRadius>0:
+        output[S>outerRadius] = 0
+    return output
+
+def cosine(N,innerRadius,outerRadius):
+    x = linspace(-floor(N/2.0),ceil(N/2.0)-1,N)
+    [X,Y] = meshgrid(x,x)
+    S = sqrt(X**2+Y**2)
+
+    T = (outerRadius - innerRadius)*2.0
+    output = (cos(2.0*pi/T*(S-innerRadius)) + 1 ) / 2
+    output[S<=innerRadius] = 1
+    output[S>outerRadius] = 0
+
+    return output
+
+def cbed(N_roi,N_dp):
+    center_index_roi = floor(N_roi/2.0)
+    index_dp_lb = int(-floor(N_dp/2.0) + center_index_roi)
+    index_dp_ub = int(ceil(N_dp/2.0) + center_index_roi)
+    output = zeros((N_roi,N_roi))
+    output[index_dp_lb:index_dp_ub,index_dp_lb:index_dp_ub] = 1
+    return output
+
+def square(N, halfSideLength):
+    x = linspace(-floor(N/2.0),ceil(N/2.0)-1,N)    
+    [X,Y] = meshgrid(x,x)
+    print(halfSideLength)
+    output = ones((N,N))
+    output[X>halfSideLength] = 0
+    output[X<-halfSideLength] = 0
+    output[Y>halfSideLength] = 0
+    output[Y<-halfSideLength] = 0
+    return output

ptycho/python/ptycho_recon/pie_mixed_states.py

@@ -0,0 +1,180 @@
+import scipy.io as sio #for read/write matlab file
+from scipy import ndimage
+import utility_function as u
+import utility_function_recon as ur
+import pyfftw
+from numpy import *
+from numpy.fft import *
+import numpy as np
+import time
+
+########## mixed states ##########
+def reconPIE_mixed_state(dp, paraDict):
+    N_probe = paraDict['N_probe']; N_object = paraDict['N_object']
+    N_roi = paraDict['N_roi']
+    Niter = paraDict['Niter'];
+    N_tot = N_roi*N_roi
+    if 'randomSeed' in paraDict: random.seed(mod(int(paraDict['randomSeed']),4294967295))
+
+    auxiFunc = ur.auxiliary_function(paraDict)
+
+    ########## initialize exit wave ########## 
+    psi = zeros((N_probe, N_roi, N_roi), dtype=np.complex128)
+    psi_old = zeros((N_probe, N_roi, N_roi), dtype=np.complex128)
+    delta_psi = zeros((N_probe, N_roi, N_roi), dtype=np.complex128)
+    
+    ########## initialize CBED ##########
+    dp_tot = sum(dp)
+    dp_avg = np.sum(dp * dp, axis = 0)/paraDict['N_scan']
+    cbed_region_mag = zeros((N_probe, paraDict['N_dp'], paraDict['N_dp']))
+
+    ########## initialize transmission function ##########
+    O = auxiFunc.initializeObject()
+
+    ########## initialize probe function ##########
+    if 'previous_probe' in paraDict:
+        probes = paraDict['previous_probe'].copy()
+    elif 'probes0' in paraDict:
+        probes = paraDict['probes0'].copy()
+    else:
+        probes = np.zeros((N_probe, N_roi, N_roi), dtype=np.complex128)
+        if paraDict['normalizeInitialProbe']:
+            #print('normalize initial probe intensity to match CBED')
+            probe = paraDict['probe0'] * sqrt(np.sum(dp_avg) / np.sum(abs(paraDict['probe0'])**2) / N_tot)
+        else:
+            probe = paraDict['probe0']
+        for i in range(N_probe):
+            probes[i,:,:] = probe / (i+1)
+
+    probes_shifted = np.zeros((N_probe,N_roi,N_roi), dtype=np.complex128)
+    probes_old = np.zeros((N_probe,N_roi,N_roi), dtype=np.complex128)
+
+    ########## miscellaneous ##########
+    startNiter = 0
+    timeLeft = 0
+    start_time = time.time()
+    time_counter = 1
+
+    #Use single object and probe before mixed states update
+    N_object_recon = 1
+    N_probe_recon = 1
+
+    dp_error = zeros(paraDict['N_scan']) #difference between data and recon wave
+    s, dp_error_old = auxiFunc.initializeDataError()
+
+    if 'previousIteration' in paraDict: 
+        startNiter = paraDict['previousIteration']
+        if startNiter >= paraDict['Niter_update_states']:
+            N_object_recon = N_object; N_probe_recon = N_probe
+
+    ################################# prepare result dictionary ################################# 
+    resultDir = {'object':O}
+    resultDir['psi'] = psi
+    resultDir['probes'] = probes
+    resultDir['probe0'] = paraDict['probe0'].copy()
+    if 'probes0' in paraDict: resultDir['probes0'] = paraDict['probes0'].copy()
+    resultDir['dx_x'] = 1.0/(paraDict['dk_x']*N_roi); 
+    resultDir['dx_y'] = 1.0/(paraDict['dk_y']*N_roi);
+    resultDir['dk_x'] = paraDict['dk_x']; resultDir['dk_y'] = paraDict['dk_y']
+    resultDir['ppX'] = paraDict['ppX']; resultDir['ppY'] = paraDict['ppY']
+    resultDir['dp_avg'] = dp_avg
+    resultDir['s'] = s
+    resultDir['badPixels'] = paraDict['badPixels']
+    
+    resultDir['dp_error'] = dp_error
+    if 'filter_r_probe' in paraDict:
+        resultDir['filter_r_probe'] = paraDict['filter_r_probe']
+    if 'filter_r_psi' in paraDict:
+        resultDir['filter_r_psi'] = paraDict['filter_r_psi']
+    if 'filter_f_probe' in paraDict:
+        resultDir['filter_f_probe'] = paraDict['filter_f_probe']
+    if 'filter_f_psi' in paraDict:
+        resultDir['filter_f_psi'] = paraDict['filter_f_psi']
+    if 'probe0_info' in paraDict:
+        resultDir['probe0_info'] = paraDict['probe0_info']
+
+    ################################## main recon loop #################################    
+    for k in range(startNiter, Niter):
+        if mod(k, paraDict['Niter_print'])==0: auxiFunc.printStatus(timeLeft, k)
+
+        if k == paraDict['Niter_update_probe']: 
+            print('start probe update')
+            start_time = time.time()
+            time_counter = 1
+        if k == paraDict['Niter_update_position']: 
+            print('start position correction')
+            start_time = time.time()
+            time_counter = 1
+        if k == paraDict['Niter_update_states']:
+            print('start mixed states update')
+            N_object_recon = N_object; N_probe_recon = N_probe
+            start_time = time.time()
+            time_counter = 1
+            
+            probes_temp = gramschmidt(probes.reshape(N_probe, N_roi*N_roi))
+            probes[:,:,:] = probes_temp.reshape(N_probe, N_roi, N_roi)
+           
+        update_order = random.permutation(paraDict['N_scan']) #random order
+ 
+        for i in update_order:
+            O_old = auxiFunc.getObjectROI(O, i)
+            for p in range(N_probe_recon):
+                probes_shifted[p,:,:] = auxiFunc.shiftProb(probes[p,:,:], i, 'toScanPosition')
+
+                #overlap projection
+                psi[p,:,:] = O_old * probes_shifted[p,:,:]
+
+                #Fourier projection
+                psi_old[p,:,:] = psi[p,:,:]
+                psi[p,:,:], cbed_region_mag[p,:,:] = auxiFunc.FFTpsi(psi[p,:,:])
+
+            psi_f_mag_tot = sqrt(np.sum(cbed_region_mag**2, axis = 0))
+            dp_error[i] = np.sum((psi_f_mag_tot - dp[i,:,:])**2)
+
+            #Fourier projection
+            for p in range(N_probe_recon):
+                psi[p,:,:] = auxiFunc.updateFourierIntensity(psi[p,:,:], dp[i,:,:], psi_f_mag_tot)
+                if 'filter_r_psi' in paraDict: psi[p,:,:] = psi[p,:,:] * paraDict['filter_r_psi']
+
+            delta_psi = psi - psi_old
+            probe_old = probes_shifted[:,:,:]
+            O_update = auxiFunc.calculateMixedStatesUpdate(probe_old, delta_psi, 'o')
+            auxiFunc.updateObj(O, O_update, i)
+
+            if k>=paraDict['Niter_update_probe']:
+                O_tot_max = amax(abs(O_old)**2)
+                for p in range(N_probe_recon):
+                    probe_update = auxiFunc.calculateUpdate(O_old, delta_psi[p,:,:], 'p')
+                    probes_shifted[p,:,:] += probe_update
+                    probes[p,:,:] = auxiFunc.shiftProb(probes_shifted[p,:,:], i, 'toOrigin', checkFilter=True)
+
+            if k>=paraDict['Niter_update_position']:
+                auxiFunc.gradPositionCorrection(probe_old[0,:,:], O_old, i, delta_psi[0,:,:])
+
+        ############################### orthogonalise probe #######################################
+        if k >= paraDict['Niter_update_states']:
+            probes_temp = gramschmidt(probes.reshape(N_probe, N_roi*N_roi))
+            probes[:,:,:] = probes_temp.reshape(N_probe, N_roi, N_roi)
+           
+        ############################### calcuate data error #######################################
+        s[k] = np.sum(dp_error)/dp_tot
+        dp_error_old = dp_error.copy()
+
+        ############################### save results #######################################
+        if mod(k+1, paraDict['Niter_save'])==0: #save results
+            saveName = paraDict['saveName'] + '_Niter'+str(k+1) + '.mat'
+            sio.savemat(saveName,resultDir)
+    
+        timeLeft = (time.time()-start_time)/time_counter * (Niter-k-1)
+        time_counter += 1
+
+def proj(u, v):
+    return u * np.vdot(u,v) / np.vdot(u,u)  
+
+def gramschmidt(V):
+    U = np.copy(V)
+    for i in range(1, V.shape[0]):
+        for j in range(i):
+            U[i,:] -= proj(U[j,:], V[i,:])
+    return U
+

ptycho/python/ptycho_recon/probe.py

@@ -0,0 +1,76 @@
+import numpy as np
+import scipy.io as sio #for read/write matlab file
+from scipy import ndimage
+
+class STEMprobe:
+      """Probe function for STEM"""
+      def __init__(self):
+          self.dx = 1.0
+          self.Nside = 256
+          self.px = 0
+          self.py = 0
+          self.voltage = 300 #keV
+          self.alpha_max = 30 #mrad
+          self.df = 0 #angstrom
+          self.cs = 0 #mm
+          self.f_a2 = 0 #angstrom
+          self.theta_a2 = 0
+          self.f_a3 = 0 #angstrom
+          self.theta_a3 = 0
+          self.f_c3 = 0 #angstrom
+          self.theta_c3 = 0
+          self.Fourier_mag = 1
+
+      def printParameters(self):
+          self.wavelength = 12.398/np.sqrt((2*511.0+self.voltage)*self.voltage)  #angstrom
+          #print out all the parameters in the dm reconstruction
+          print("probe size:",self.Nside,"x",self.Nside)
+          print("distance between adjacent pixels: dx =", self.dx)
+          print("distance between adjacent pixels in Fourier space: dk =", 1.0/(self.dx*self.Nside))
+
+          print("probe position: (px,py) =(",self.px,",",self.py,")")
+          print("beam voltage =", self.voltage, "keV")
+          print("beam wavelength =",self.wavelength, "angstrom")
+          print("semi-convergence angle =", self.alpha_max, "mrad")
+          print("defocus=", self.df, "angstrom")
+          print("spherical aberration =", self.cs, "angstrom")
+          print("two-fold astigmatism =", self.f_a2, "angstrom.", "azimuthal orientation=",self.theta_a2, "rad")
+          print("three-fold astigmatism =", self.f_a3, "angstrom.", "azimuthal orientation=",self.theta_a3, "rad")
+          print("coma =", self.f_c3, "angstrom.", "azimuthal orientation=",self.theta_c3, "rad")
+
+      def generateProbe(self):
+          print("generating probe function...")
+          self.wavelength = 12.398/np.sqrt((2*511.0+self.voltage)*self.voltage)  #angstrom
+          amax = self.alpha_max*1e-3 # in rad
+          amin = 0.0
+
+          k_max = amax/self.wavelength
+          k_min = amin/self.wavelength
+
+          dk= 1.0/(self.dx*self.Nside)
+          kx = np.linspace(-np.floor(self.Nside/2.0),np.ceil(self.Nside/2.0)-1,self.Nside)
+          [kY,kX] = np.meshgrid(kx,kx)
+          kX = kX*dk; kY = kY*dk;
+          kR = np.sqrt(kX**2+kY**2)
+          theta = np.arctan2(kY,kX)
+          
+          chi = -np.pi*self.wavelength*kR**2*self.df + np.pi/2*self.cs*1e7*self.wavelength**3*kR**4+np.pi*self.f_a2*self.wavelength*kR**2*np.sin(2*(theta-self.theta_a2))+2*np.pi/3*self.f_a3*self.wavelength**2*kR**3*np.sin(3*(theta-self.theta_a3))+2*np.pi/3*self.f_c3*self.wavelength**2*kR**3*np.sin(theta-self.theta_c3)
+
+          probe = np.exp(-1j*chi)*np.exp(-2*np.pi*1j*self.px*kX)*np.exp(-2*np.pi*1j*self.py*kY)
+          probe[kR>k_max] = 0
+          probe[kR<k_min] = 0
+          
+          if self.Fourier_mag != 1:
+              probe = probe/abs(probe) * self.Fourier_mag
+
+          #probe = probe/np.sum((np.abs(probe)**2)) #normalize probe
+          
+          probe = np.fft.fftshift(np.fft.ifft2(np.fft.ifftshift(probe)))
+          probe = probe/np.sqrt(np.sum((np.abs(probe)**2)*self.dx*self.dx)) #normalize probe
+          #probe = probe/np.sqrt(np.sum((np.abs(probe)**2))) #normalize probe
+          
+          mask = np.ones(kR.shape)
+          mask[kR>k_max] = 0
+          mask[kR<k_min] = 0
+
+          return probe

ptycho/python/ptycho_recon/ptycho.py

@@ -0,0 +1,143 @@
+from numpy import *
+from numpy.fft import *
+import numpy as np
+import filters as filters
+import scipy.io as sio #for read/write matlab file
+from scipy import ndimage
+import os #for change directory
+
+import utility_function as u
+import pie_mixed_states
+from probe import STEMprobe
+
+##############################################################################################
+class ptycho:
+      """Ptychography reconstruction"""
+      def __init__(self, dp, dk, initialProbe, ppX, ppY):
+          self.initialProbe = initialProbe
+          self.paraDict = {'dk':dk}
+
+          ########## reshape scan positions and diffraction patterns ##########
+          N_scan_y = dp.shape[2]
+          N_scan_x = dp.shape[3]
+          N_scan_tot = N_scan_y * N_scan_x
+
+          self.dp = zeros((N_scan_tot,dp.shape[0],dp.shape[1]))
+          self.paraDict['ppX'] = ppX.reshape(N_scan_tot)
+          self.paraDict['ppY'] = ppY.reshape(N_scan_tot)
+
+          for i in range(N_scan_y):
+              for j in range(N_scan_x):
+                  index = i*N_scan_x + j
+                  self.dp[index,:,:] = sqrt(dp[:,:,i,j])
+
+          self.paraDict['dk_y'] = dk
+          self.paraDict['dk_x'] = dk
+          
+          self.paraDict['N_dp'] = self.dp.shape[1]
+          self.paraDict['N_scan'] = self.dp.shape[0]
+          self.paraDict['badPixels'] = zeros((self.dp.shape[1],self.dp.shape[1]))
+          self.paraDict['Niter'] = 200
+          self.paraDict['Niter_save'] = 50
+          self.paraDict['Niter_print'] = 1
+          self.paraDict['beta'] = 1.0
+          self.paraDict['alpha'] = 0.1
+          
+          self.paraDict['Niter_update_probe'] = 10
+          self.paraDict['uniformInitialObject'] = True
+          self.paraDict['normalizeInitialProbe'] = True
+
+          self.paraDict['filter_r_type_psi'] = 'none'
+          self.paraDict['filter_r_type_probe'] = 'none'
+          self.paraDict['filter_f_type_psi'] = 'cbed'
+          self.paraDict['filter_f_type_probe'] = 'none'
+         
+          self.paraDict['saveData'] = False
+          self.paraDict['loadData'] = False
+          
+          self.paraDict['reconID'] = 0
+          self.paraDict['printID'] = ''
+          
+          #mixed-states
+          self.paraDict['N_probe'] = 1
+          self.paraDict['N_object'] = 1
+      
+      def recon(self):
+          print("begin ptychographic reconstruction")
+          pie_mixed_states.reconPIE_mixed_state(self.dp, self.paraDict)
+
+      def initialize(self, result_dir):
+           ########## initial probe ##########
+           #create initial probe
+           self.initialProbe.dx = 1.0/(self.paraDict['dk']*self.paraDict['N_roi'])
+           #print self.initialProbe.dx
+           self.initialProbe.Nside = self.paraDict['N_roi']
+
+           if not 'probe0' in self.paraDict: self.paraDict['probe0'] = self.initialProbe.generateProbe()
+           #self.initialProbe.printParameters()
+
+           ########## save data ##########
+           if self.paraDict['saveData']:
+               print('saving dp_recon...')
+               if 'dataDir' in self.paraDict:
+                   if not os.path.exists(self.paraDict['dataDir']): os.makedirs(self.paraDict['dataDir'])
+                   os.chdir(self.paraDict['dataDir'])
+               else:
+                   if not os.path.exists(result_dir): os.makedirs(result_dir)
+                   os.chdir(result_dir)
+               sio.savemat('dp_recon',{'dp_recon':self.dp,'dk':self.paraDict['dk']})
+
+           ########## filters ##########
+           a = self.generateFilters('', self.paraDict['filter_r_type_psi'], 'r', 'psi')
+           a = self.generateFilters('', self.paraDict['filter_r_type_probe'], 'r', 'probe')
+           a = self.generateFilters('', self.paraDict['filter_f_type_psi'], 'f', 'psi')
+           a = self.generateFilters('', self.paraDict['filter_f_type_probe'], 'f', 'probe')
+
+           self.paraDict['saveName'] =  "recon"
+
+           ########## create result dir ##########
+           if not os.path.exists(result_dir): os.makedirs(result_dir)
+           os.chdir(result_dir)
+           return result_dir
+
+      def generateFilters(self, result_dir, filterType, space, waveFunction):
+          if self.paraDict['filter_' + space + '_type_' + waveFunction] == 'none':
+              return result_dir
+          #print("Generating " + filterType + " filter in " + space + " space for " + waveFunction)
+          N_roi = self.paraDict['N_roi']
+          result_dir = result_dir + "/filter_" + space + "_"
+          filerKey = 'filter_' + space + '_' + waveFunction
+          if filterType == "cbed":
+               self.paraDict[filerKey]  = filters.cbed(N_roi, self.paraDict['N_dp'] )
+               result_dir = result_dir + "cbed" + str(self.paraDict['N_dp'])
+
+          elif filterType == "square":
+               cutoff = self.paraDict['filter_' + space + '_inner_cutoff_' + waveFunction]
+               self.paraDict[filerKey]  = filters.square(N_roi, cutoff)
+               result_dir = result_dir + "square_cutoff"+str(cutoff)
+
+          elif filterType == "gaussian":
+               inner_cutoff = self.paraDict['filter_' + space + '_inner_cutoff_' + waveFunction]
+               outer_cutoff = self.paraDict['filter_' + space + '_outer_cutoff_' + waveFunction]
+               sigma = self.paraDict['filter_' + space +'_gaussian_sigma_' + waveFunction]
+               self.paraDict[filerKey] = filters.gaussian(N_roi, inner_cutoff, outer_cutoff, sigma)
+               result_dir = result_dir + "gaussian_cutoff"+str(inner_cutoff)+"_outer_cutoff"+str(outer_cutoff)+"_sigma"+str(sigma)
+
+          elif filterType  == "cosine":
+               inner_cutoff = self.paraDict['filter_' + space + '_inner_cutoff_' + waveFunction]
+               outer_cutoff = self.paraDict['filter_' + space + '_outer_cutoff_' + waveFunction]
+               self.paraDict[filerKey] = filters.cosine(N_roi, inner_cutoff, outer_cutoff)
+               result_dir = result_dir + "cosine_inner_cutoff"+str(inner_cutoff)+"_outer_cutoff"+str(ff_outer_cutoff)
+
+          elif filterType == "disk":
+               cutoff = self.paraDict['filter_' + space + '_disk_cutoff_' + waveFunction]
+               self.paraDict[filerKey] = filters.disk(N_roi, cutoff)
+               result_dir = result_dir + "disk_cutoff"+str(cutoff)
+          else:
+              raise RuntimeError('Unknown filter type!')
+          if space == 'f':
+              self.paraDict[filerKey] = ifftshift( self.paraDict[filerKey] )
+
+          result_dir = result_dir + "_" + waveFunction
+          return result_dir
+

ptycho/python/ptycho_recon/utility_function.py

@@ -0,0 +1,466 @@
+import numpy as np
+from numpy import *
+from scipy import ndimage
+import scipy.ndimage
+import pyfftw
+from numpy.fft import *
+import pyfftw
+
+import zipfile as zp
+import os
+import warnings
+
+def readraw(filename):
+    scanx = int(filename.rstrip('.raw').split('_')[-1].lstrip('abcdefghijklmnopqrstuvwxyz'))
+    scany = int(filename.rstrip('.raw').split('_')[-2].lstrip('abcdefghijklmnopqrstuvwxyz'))
+    contents = np.fromfile(filename, dtype = 'float32')
+    data_arr = np.reshape(contents, (125, 125, scany, scanx), order = 'C')
+    return data_arr
+    
+##############################################################################################
+def shift(input,dx,px,py):
+    N_image = input.shape[0]
+
+    dk= 1.0/(dx*N_image)
+    kx = np.linspace(-np.floor(N_image/2.0),np.ceil(N_image/2.0)-1,N_image)
+    [kX,kY] = np.meshgrid(kx,kx)
+    kX = kX*dk; kY = kY*dk;
+
+    f = np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(input)))
+    f = f*np.exp(-2*np.pi*1j*px*kX)*np.exp(-2*np.pi*1j*py*kY)
+    f = np.fft.fftshift(np.fft.ifft2(np.fft.ifftshift(f)))
+
+    return f
+
+##############################################################################################
+def convertADUtoElectronCount(input, ADU_electronCount_ratio, directory = ""):
+    print("converting ADU to # of electrons:", ADU_electronCount_ratio)
+    output = input/ADU_electronCount_ratio
+    directory = directory + "_ADUtoElectron" + str(ADU_electronCount_ratio)
+    return output, directory
+
+##############################################################################################
+def partition_scan(N_scan_x, N_scan_y, partion_x, partion_y):
+    scan_partition = np.zeros((N_scan_y, N_scan_x), dtype=np.int)
+    delta_x = int(np.ceil(N_scan_x / partion_x))
+    delta_y = int(np.ceil(N_scan_y / partion_y))
+    value = 0
+    for i in xrange(partion_y):
+        for j in xrange(partion_x):
+             index_x_lb = delta_x*j
+             index_x_ub = min(delta_x*(j+1), N_scan_x)
+             index_y_lb = delta_y*i
+             index_y_ub = min(delta_y*(i+1), N_scan_y)
+             scan_partition[index_y_lb:index_y_ub,index_x_lb:index_x_ub] = value
+             value = value + 1
+    return scan_partition
+
+##############################################################################################
+def shift_cbed(input, px, py, directory = ""):
+    N_cbed = input.shape[0]
+    N_tot = N_cbed*N_cbed
+    dx = 1.0
+    dk= 1.0/(dx*N_cbed)
+    kx = np.linspace(-np.floor(N_cbed/2.0),np.ceil(N_cbed/2.0)-1,N_cbed)
+    kx = ifftshift(kx)
+    [kX,kY] = np.meshgrid(kx,kx)
+    kX = kX*dk; kY = kY*dk;
+    output = zeros(input.shape)
+    for i in range(input.shape[2]):
+        for j in range(input.shape[3]):
+            f = np.fft.fft2(input[:,:,i,j]);
+            f = f*exp(-2*pi*1j*px*kX)*exp(-2*pi*1j*py*kY)
+            output[:,:,i,j] = abs(np.fft.ifft2(f)); #fix normalization
+    directory = directory + "_shift_sx" + str(px)+"_sy" + str(py)
+    return output, directory
+
+##############################################################################################
+def lowpassfilter(input, dx,cutoff):
+    N = input.shape[0]
+    dk= 1.0/(dx*N)
+    kx = np.linspace(-np.floor(N/2.0),np.ceil(N/2.0)-1,N)
+    [kX,kY] = np.meshgrid(kx,kx)
+    kR = np.sqrt(kX**2 + kY**2)
+
+    a = np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(input)))
+    a[kR>(N/2*cutoff)] = 0
+    output = np.fft.fftshift(np.fft.ifft2(np.fft.ifftshift(a)))
+
+    return output
+
+##############################################################################################
+def lowpassfilter_alpha(input, cutoff, dk_x, dk_y, alpha_max, voltage):
+    N = input.shape[0]
+    print("applying low pass filter to image")
+    print("mask cutoff =",cutoff ,'alpha')
+    kx = linspace(-floor(N/2.0),ceil(N/2.0)-1, N)
+    [kX,kY] = meshgrid(kx,kx)
+    kX = kX*dk_x; kY = kY*dk_y;
+    kR = np.sqrt(kX**2+ kY**2)
+    
+    wavelength = 12.398/np.sqrt((2*511.0 + voltage) * voltage)  #angstrom
+
+    k_cutoff = cutoff * alpha_max *1e-3 / wavelength
+    
+    f = fftshift(fft2(ifftshift(input)))
+    f[kR > k_cutoff] = 0
+    output = real(fftshift(ifft2(ifftshift(f))))
+
+    return output
+##############################################################################################
+def propagtor_function(N, dk_x, dk_y,dz, wavelength):
+    kx = np.linspace(-np.floor(N/2.0),np.ceil(N/2.0)-1,N)
+    [kX,kY] = np.meshgrid(kx,kx)
+    kX = kX*dk_x; kY = kY*dk_y;
+
+    kR = np.sqrt(kX**2 + kY**2)
+    cutoff = (np.ceil(N/2.0)-1)*min(dk_x,dk_y)*2/3
+
+    P = zeros((dz.size,N,N), dtype=np.complex128)
+    for i in range(dz.size):
+        temp = np.exp(-1j*np.pi*wavelength*kR**2*dz[i])
+        temp[kR>cutoff] = 0  #apply a low pass filter to keep 2/3 of maximum spatial frequency
+        P[i,:,:] = ifftshift(temp)
+
+    '''
+    x = np.linspace(-np.floor(N/2.0),np.ceil(N/2.0)-1,N)*dx
+    [X,Y] = np.meshgrid(x,x)
+    R = np.sqrt(X**2 + Y**2)
+    p = 1/(1j*wavelength*dz)*np.exp(1j*np.pi/(wavelength*dz)*R**2)
+    '''
+    return P
+
+##############################################################################################
+def recenter(input):
+    N = input.shape[0]
+    center_index = N // 2
+    [yy,xx] = where(abs(input) == np.max(abs(input)))
+
+    output = roll(input,int(-(yy[0]-center_index)), axis = 0)
+    output = roll(output,int(-(xx[0]-center_index)), axis = 1)
+
+    return output
+
+##############################################################################################
+def upsample_cbed_ff(dp, dk_x, dk_y, resizeFactor = 2, directory = ""):
+    print("upsample cbed using free float ptychography")
+    print("old dk_x=", dk_x, "old dk_y=", dk_y)
+
+    #resize data
+    output = np.zeros((int(dp.shape[0]*resizeFactor), int(dp.shape[1]*resizeFactor), dp.shape[2], dp.shape[3]))
+    output[0:-1:resizeFactor,0:-1:resizeFactor,:,:] = dp
+    mask = np.ones((int(dp.shape[0]*resizeFactor), int(dp.shape[1]*resizeFactor)))
+    print(mask.shape)
+    mask[0:-1:2,0:-1:2] = 0
+    dk_x_r = dk_x/resizeFactor
+    dk_y_r = dk_y/resizeFactor
+    print("new dk_x=", dk_x_r, "new dk_y=", dk_y_r)
+    directory = directory + "_upsampleCBED" + str(resizeFactor)
+    #output[output<0] = 0
+    return output, mask, dk_x_r, dk_y_r, directory
+
+
+##############################################################################################
+def resize_cbed(dp, resizeFactor, dk_x, dk_y, directory = "", order = 1):
+    print("resize cbed")
+    print("old dk_x=", dk_x, "old dk_y=", dk_y)
+
+    #resize data
+    output = np.zeros((int(dp.shape[0]*resizeFactor),int(dp.shape[1]*resizeFactor),dp.shape[2],dp.shape[3]))
+    for i in range(0,dp.shape[2]):
+        for j in range(0,dp.shape[3]):
+            scipy.ndimage.interpolation.zoom(dp[:,:,i,j],[resizeFactor,resizeFactor],output[:,:,i,j], order)
+    dk_x_r = dk_x/resizeFactor
+    dk_y_r = dk_y/resizeFactor
+    print("new dk_x=", dk_x_r, "new dk_y=", dk_y_r)
+    directory = directory + "_resizeCBED" + str(resizeFactor)
+    #output[output<0] = 0
+    return output, dk_x_r, dk_y_r, directory
+
+##############################################################################################
+def resample_cbed(dp, Npix, dk_x, dk_y, directory = ""):
+    print("resample cbed using every ", str(Npix), 'pixels...')
+    print("old dk_x=", dk_x, "old dk_y=", dk_y)
+    output = dp[0:-1:Npix,0:-1:Npix,:,:]
+    if Npix>1: directory = directory + "_resampleCBED" + str(Npix)
+    dk_x_new = dk_x*Npix; dk_y_new = dk_y*Npix
+    print("new dk_x=", dk_x_new, "new dk_y=", dk_y_new)
+    return output, dk_x_new, dk_y_new, directory
+    
+##############################################################################################
+def crop_cbed(dp, N_dp_x_new, N_dp_y_new, directory = ""):
+    print("crop cbed to ", str(N_dp_y_new), 'x', str(N_dp_x_new))
+    cen_x = floor(dp.shape[1]/2.0)
+    cen_y = floor(dp.shape[0]/2.0)
+    index_x_lb = (cen_x - floor(N_dp_x_new/2.0)).astype(np.int)
+    index_x_ub = (cen_x + ceil(N_dp_x_new/2.0)).astype(np.int)
+    index_y_lb = (cen_y - floor(N_dp_y_new/2.0)).astype(np.int)
+    index_y_ub = (cen_y + ceil(N_dp_y_new/2.0)).astype(np.int)
+    
+    #crop data
+    output = np.zeros((N_dp_x_new, N_dp_y_new,dp.shape[2],dp.shape[3]))
+    for i in range(0,dp.shape[2]):
+        for j in range(0,dp.shape[3]):
+            output[:,:,i,j] = dp[index_y_lb:index_y_ub,index_x_lb:index_x_ub,i,j]
+    directory = directory + "_crop_Ndpx" + str(N_dp_x_new) + "_Ndpy" + str(N_dp_y_new)
+    return output, directory
+
+##############################################################################################
+def pad_cbed(input, N_pad_x, N_pad_y, value = 0, directory = ""):
+    print("pad " + str(value) + "s to cbed patterns...")
+    Ny = input.shape[0]; Nx = input.shape[1];
+    pad_pre_y = int(np.ceil((N_pad_y - Ny) / 2.0))
+    pad_post_y = int(np.floor((N_pad_y - Ny) / 2.0))
+    pad_pre_x = int(np.ceil((N_pad_x - Nx) / 2.0))
+    pad_post_x = int(np.floor((N_pad_x - Nx) / 2.0))
+  
+    output = np.pad(input, ((pad_pre_y, pad_post_y), (pad_pre_x, pad_post_x), (0,0), (0,0)), 'constant', constant_values=value)
+    directory = directory + "_padCBED" + str(value) + "_" + str(N_pad_x)
+    return output, directory
+
+##############################################################################################
+def apply_circular_mask(input, radius, offset_x=0, offset_y=0, directory = ""):
+    N_dp = input.shape[0]
+    print("applying circullar mask to cbed")
+    print("mask radius =",radius)
+    print("mask offset_x =",offset_x, "mask offset_y =",offset_y)
+
+    x = np.linspace(-np.floor(N_dp/2.0),np.ceil(N_dp/2.0)-1,N_dp)
+    [X,Y] = np.meshgrid(x,x)
+    mask_disk = np.sqrt(X**2+ Y**2)
+    
+    output = input.copy()
+
+    for i in range(input.shape[2]):
+        for j in range(input.shape[3]):
+            temp = input[:,:,i,j].copy()
+            temp = np.roll(temp, offset_y, axis=0)
+            temp = np.roll(temp, offset_x, axis=1)
+            if radius>0:
+               temp[mask_disk>radius] = 0
+            output[:,:,i,j] = temp.copy()
+    if radius>0: directory = directory + "_lowPassFilter" + str(radius)
+    if offset_x!= 0: directory = directory + "_sx" + str(offset_x)
+    if offset_y!= 0: directory = directory + "_sy" + str(offset_y)
+    
+    return output, directory, mask_disk
+
+##############################################################################################
+def apply_circular_mask_alpha(dp, cutoff, dk_x, dk_y, alpha_max, voltage, offset_x=0, offset_y=0, directory = ""):
+    if cutoff>0:
+        N_dp = dp.shape[0]
+        print("applying circullar mask to cbed")
+        print("mask cutoff =",cutoff ,'alpha')
+        print("mask offset_x =",offset_x, "mask offset_y =",offset_y)
+        kx = linspace(-floor(N_dp/2.0),ceil(N_dp/2.0)-1, N_dp)
+
+        [kX,kY] = meshgrid(kx,kx)
+        kX = kX*dk_x; kY = kY*dk_y;
+        kR = np.sqrt(kX**2+ kY**2)
+        
+        wavelength = 12.398/np.sqrt((2*511.0 + voltage) * voltage)  #angstrom
+
+        k_cutoff = cutoff * alpha_max *1e-3 / wavelength
+    output = dp.copy()
+    for i in range(dp.shape[2]):
+        for j in range(dp.shape[3]):
+            temp = dp[:,:,i,j].copy()
+            temp = np.roll(temp, offset_y, axis=0)
+            temp = np.roll(temp, offset_x, axis=1)
+            #np.roll(dp[:,:,i,j], offset_y, axis=0)
+            #np.roll(dp[:,:,i,j], offset_x, axis=1)
+            if cutoff > 0:
+               #dp[kR > k_cutoff,i,j]= 0
+               temp[kR > k_cutoff] = 0
+            output[:,:,i,j] = temp.copy()
+    directory = directory + "_cutoff" + str(cutoff)+"alpha"
+    if offset_x!= 0: directory = directory + "_sx" + str(offset_x)
+    if offset_y!= 0: directory = directory + "_sy" + str(offset_y)
+    
+    return output, directory
+##############################################################################################
+def add_poisson_noise(input, current, readOutTime, directory = ""):
+    N_dp_tot = input.shape[0] * input.shape[1]
+    print("applying poisson noise to cbed")
+    print("beam current =", current, 'pA')
+    Nc_avg = current*1e-12*readOutTime/(1.6e-19)/N_dp_tot
+    print("average count per pixel = ", Nc_avg) 
+    print("snr = ", sqrt(Nc_avg)) 
+    
+    output = input.copy()
+    snr = zeros((input.shape[2],input.shape[3]))
+    for i in range(input.shape[2]):
+        for j in range(input.shape[3]):
+            cbed_noise = input[:,:,i,j].copy()
+            cbed_noise = cbed_noise/np.sum(input[:,:,i,j])*(N_dp_tot*Nc_avg*1.0)
+            cbed_noise = random.poisson(cbed_noise)
+            cbed_noise = cbed_noise*np.sum(input[:,:,i,j])/(N_dp_tot*Nc_avg*1.0)
+            output[:,:,i,j] = cbed_noise
+            snr[i,j] = np.mean(input[:,:,i,j])/np.std(input[:,:,i,j] - cbed_noise)
+    directory = directory + "_poissonNoise" + str(current) + "pA"
+    return output, snr, directory
+
+##############################################################################################
+def average_cbed(input, windowSize, scanStepSize_x, scanStepSize_y, directory = ""):
+    print("average cbed patterns... window size =",windowSize)
+    
+    if windowSize==1:
+        output = input
+    else:
+        output = zeros((input.shape[0],input.shape[1],input.shape[2]//windowSize,input.shape[3]//windowSize))
+        for i in range(output.shape[2]):
+            for j in range(output.shape[3]):
+                temp = np.sum(input[:,:,i*windowSize:(i+1)*windowSize,j*windowSize:(j+1)*windowSize], axis=(2,3))/windowSize**2
+                output[:,:,i,j] = temp.copy()
+        directory = directory + "_averageCBED"+str(windowSize)
+        scanStepSize_x = scanStepSize_x * windowSize
+        scanStepSize_y = scanStepSize_y * windowSize
+    
+    return output, scanStepSize_x, scanStepSize_y, directory
+
+##############################################################################################
+def resample_scan(input, windowSize, scanStepSize_x, scanStepSize_y, directory = ""):
+    print("resample scans ... window size =",windowSize)
+     
+    if windowSize==1:
+        output = input
+    else:
+        output = input[:,:,::windowSize,::windowSize]
+        
+        scanStepSize_x = scanStepSize_x * windowSize
+        scanStepSize_y = scanStepSize_y * windowSize
+
+        directory = directory + "_resampleScan"+str(windowSize)
+
+    return output, scanStepSize_x, scanStepSize_y, directory
+
+##############################################################################################
+def transpose_cbed(input, directory = ""):
+    print("transpose cbed patterns")
+    output = zeros((input.shape[1],input.shape[0],input.shape[2],input.shape[3]))
+
+    for i in range(input.shape[2]):
+        for j in range(input.shape[3]):
+            output[:,:,i,j] = input[:,:,i,j].T
+    directory = directory + "_transpose" 
+    return output, directory
+
+##############################################################################################
+def rot90_cbed(input, k, directory = ""):
+    print("rotate cbed patterns by", 90*k, "degrees...")
+    output = zeros((input.shape[1],input.shape[0],input.shape[2],input.shape[3]))
+    for i in range(input.shape[2]):
+        for j in range(input.shape[3]):
+            output[:,:,i,j] = rot90(input[:,:,i,j], k)
+    directory = directory + "_rotate90_"+str(k) 
+    return output, directory
+
+##############################################################################################
+def transpose_scan_positions(input, directory = ""):
+    print("transpose scan positions...")
+    output = zeros((input.shape[0],input.shape[1],input.shape[3],input.shape[2]))
+
+    for i in range(output.shape[2]):
+        for j in range(output.shape[3]):
+            output[:,:,i,j] = input[:,:,j,i]
+
+    directory = directory + "_transposeScanPos"
+    return output, directory
+
+##############################################################################################
+def flip_scan_positions(input, flipType, directory = ""):
+    assert flipType in ['lr','ud'], "Flip type %s not known!" % flipType
+    if flipType=="ud":
+        print("flip scan positions up and down (y-axis, third dimension)")
+        output = input[:,:,::-1,:]
+    if flipType=="lr":
+        print("flip scan positions left and right (x-axis, forth dimension)")
+        output = input[:,:,:,::-1]
+
+    directory = directory + "_flipScanPos_" + flipType
+    return output, directory
+##############################################################################################
+def normalize_wave_function(input, dx):
+    output = input.copy()
+    c = sqrt( 1.0/ ( np.sum(np.abs(input)**2 * dx**2 ) ))
+    output = output * c
+    return output
+
+##############################################################################################
+def normalize_cbed(input, dk, directory = ""):
+    print("normalizing cbed patterns: sum(dp) = 1")
+    
+    for i in range(input.shape[2]):
+        for j in range(input.shape[3]):
+            c = 1.0/(sum(abs(input[:,:,i,j])))
+            input[:,:,i,j] = input[:,:,i,j] * c
+    directory = directory + "_normCBED"
+    return input, directory
+    
+##############################################################################################
+def background_removal(input, bg_level, directory = ""):
+    print("removing background: threshold=", bg_level)
+    output = input.copy()
+    output[output<=bg_level] = 0
+    directory = directory + "_bgRemoval"+str(bg_level)
+    return output, directory
+
+##############################################################################################
+def background_subtraction(input, bg_level, directory = ""):
+    print("subtracting background: threshold=", bg_level)
+
+    output = input - bg_level
+    output[output<0] = 0
+
+    directory = directory + "_bgSubtraction"+str(bg_level)
+    return output, directory
+
+##############################################################################################
+def calculate_scan_positions(N_scan_x, N_scan_y, scanStepSize_x, scanStepSize_y, rot_angle_d = 0, directory = "", randomOffset = 0, ppX = 0, ppY = 0):
+    print("calculate scan positions")
+    print("N_scan_x =", N_scan_x, "scanStepSize_x =", scanStepSize_x)
+    print("N_scan_y =", N_scan_y, "scanStepSize_y =", scanStepSize_y)
+    print("rot_angle =", rot_angle_d)
+    rot_angle = rot_angle_d*pi/180.0
+
+    ppx = linspace(-floor(N_scan_x/2.0),ceil(N_scan_x/2.0)-1,N_scan_x)*scanStepSize_x
+    ppy = linspace(-floor(N_scan_y/2.0),ceil(N_scan_y/2.0)-1,N_scan_y)*scanStepSize_y
+    [ppX0, ppY0] = meshgrid(ppx,ppy)
+    if not isscalar(ppX): ppX0 = ppX
+    if not isscalar(ppY): ppY0 = ppY
+
+    if randomOffset > 0:
+        ppX0 = ppX0 + (np.random.rand(ppX0.shape[0], ppX0.shape[1])*2-1)*scanStepSize_x*randomOffset
+        ppY0 = ppY0 + (np.random.rand(ppY0.shape[0], ppY0.shape[1])*2-1)*scanStepSize_y*randomOffset
+
+    ppY_rot = ppX0*-sin(rot_angle) + ppY0*cos(rot_angle)
+    ppX_rot = ppX0*cos(rot_angle) + ppY0*sin(rot_angle)
+
+    directory = directory + "/scanStepSize" + str(np.around(scanStepSize_x,4)) + "_rotAngle"+str(rot_angle_d)
+    if randomOffset>0: directory = directory + "_randomOffset" + str(randomOffset)
+    if not isscalar(ppX) or not isscalar(ppY): directory = directory + "_externalCoord"
+
+    return ppX_rot, ppY_rot, directory
+
+##############################################################################################
+def gaussian(N, sigma):
+    x = linspace(-floor(N/2.0),ceil(N/2.0)-1, N)
+
+    [X,Y] = meshgrid(x,x)
+    g = exp(-(X**2+Y**2)/(2*sigma**2)) + np.zeros((N,N), dtype=np.complex128)
+    return g
+
+##############################################################################################
+def guess_bad_scan(I, threshold):
+    print('Determining bad scans')
+
+    I_pad = np.lib.pad(I, (1, 1), 'edge')
+    # calculate standard deviation in a 3 x 3 window
+    averageI2 = scipy.ndimage.filters.uniform_filter(I_pad ** 2)
+    averageI = scipy.ndimage.filters.uniform_filter(I_pad)
+    std = np.sqrt(abs(averageI2 - averageI**2))[1:-1, 1:-1]
+
+    medianI = scipy.ndimage.filters.median_filter(I_pad, 2)[1:-1, 1:-1]
+
+    return abs(I - medianI) > std * threshold
+

ptycho/python/ptycho_recon/utility_function_recon.py

@@ -0,0 +1,237 @@
+import numpy as np
+from numpy import *
+from scipy import ndimage
+import scipy.ndimage
+import pyfftw
+from numpy.fft import *
+import multiprocessing
+
+class auxiliary_function:
+      def __init__(self, paraDict):
+          self.paraDict = paraDict          
+          self.createFourierCoord(self.paraDict['N_roi'])
+          self.initializeFFTW(self.paraDict['N_roi'])
+          self.initializeScanPositions()
+          self.initializeDiffractionPatterns()
+          self.alpha = self.paraDict['alpha']
+          self.beta = self.paraDict['beta']
+
+      def calculateUpdate(self, psi, delta_psi, type):
+          psi_mag = abs(psi)**2
+          if type=='o': #calculate update for object
+              w = self.alpha
+          elif type=='p': #calculate update for probe/psi
+              w = self.beta
+          else: raise RuntimeError('Invalid input!')
+
+          return w * conj(psi)/ amax(psi_mag) * delta_psi
+
+      def calculateMixedStatesUpdate(self, psi, delta_psi, type):
+          psi_tot = sum(abs(psi)**2,axis=0)
+          #psi_mag = abs(psi)**2
+          if type=='o': #calculate update for object
+              w = self.alpha
+          elif type=='p': #calculate update for probe/psi
+              w = self.beta
+          else: raise RuntimeError('Invalid input!')
+
+          return w * sum(conj(psi)*delta_psi, axis = 0)/amax(psi_tot)
+
+      #################### object ####################
+      def initializeObject(self):
+          if 'previous_obj' in self.paraDict:
+              O = self.paraDict['previous_obj']
+          else:
+              N_image = self.paraDict['N_image']
+              if self.paraDict['uniformInitialObject']:
+                  O = np.ones((N_image, N_image), dtype=np.complex128)
+              else:
+                  O = np.random.rand(N_image,N_image) + 1j*np.random.rand(N_image,N_image)
+              O = O/abs(O)
+          return O
+
+      def updateObj(self, O, objUpdate, ind_dp):
+          O[self.ind_y_lb_s[ind_dp]:self.ind_y_ub_s[ind_dp],self.ind_x_lb_s[ind_dp]:self.ind_x_ub_s[ind_dp]] += objUpdate
+
+      def getObjectROI(self, O, ind_dp):
+          return O[self.ind_y_lb_s[ind_dp]:self.ind_y_ub_s[ind_dp],self.ind_x_lb_s[ind_dp]:self.ind_x_ub_s[ind_dp]]
+
+      #################### probe ####################
+      def shiftProb(self, probe, ind_dp, direction, checkFilter=False):
+          if direction=='toScanPosition': #from origin to scan position
+              px = self.px_f[ind_dp]
+              py = self.py_f[ind_dp]
+          elif direction=='toOrigin': #from scan position back to origin
+              px = -self.px_f[ind_dp]
+              py = -self.py_f[ind_dp]
+          else: raise RuntimeError('Invalid input!')
+
+          self.r[:,:] = probe
+          self.fft_forward.update_arrays(self.r, self.f)
+          self.fft_forward.execute()
+          self.f = self.f*exp(-2*pi*1j*px*self.kX)*exp(-2*pi*1j*py*self.kY)
+          if checkFilter and 'filter_f_probe' in self.paraDict: 
+              self.f = self.f * self.paraDict['filter_f_probe']
+          self.fft_inverse.update_arrays(self.f, self.r)
+          self.fft_inverse.execute();
+          return self.r / self.N_tot #fix normalization
+
+      def orthoProbe(self, probes):
+          probes_temp = gramschmidt(probes.reshape(paraDict['N_probe'], paraDict['N_roi']**2))
+          probes[:,:,:] = probes_temp.reshape(paraDict['N_probe'], paraDict['N_roi'], paraDict['N_roi'])
+          #sort probes based on power
+          power = sum(abs(probes)**2, axis=(1,2))
+          power_ind = argsort(-power)
+          probes[:,:,:] = probes[power_ind,:,:]
+          return probes
+
+      def processProbe(self, probe):
+          if 'filter_r_probe' in self.paraDict: probe = probe * paraDict['filter_r_probe']
+          if 'probe_profile' in self.paraDict:
+              probe_mag_sum = sum(abs(probe))
+              probe = probe / abs(probe) * self.paraDict['probe_profile']
+              probe = probe / sum(abs(probe)) * probe_mag_sum
+          return probe
+
+      def FFTpsi(self, psi):
+          self.r[:,:] = psi
+          self.fft_forward.update_arrays(self.r, self.f)
+          self.fft_forward.execute()
+          self.f = fftshift(self.f)
+          psi_f_cbed_region_mag = abs(self.f[self.ind_dp_lb:self.ind_dp_ub, self.ind_dp_lb:self.ind_dp_ub])
+          psi_f = self.f
+          return psi_f, psi_f_cbed_region_mag
+
+      def updateFourierIntensity(self, psi_f, dp, denominator):
+          #psi_f: wave function in Fourier space
+          self.f[:,:] = psi_f;
+          f_cbed = self.f[self.ind_dp_lb:self.ind_dp_ub, self.ind_dp_lb:self.ind_dp_ub]
+          
+          f_cbed[self.paraDict['badPixels']==0] = f_cbed[self.paraDict['badPixels']==0]/(denominator[self.paraDict['badPixels']==0]+1e-16) * dp[self.paraDict['badPixels']==0]
+          self.f[self.ind_dp_lb:self.ind_dp_ub, self.ind_dp_lb:self.ind_dp_ub] = f_cbed
+          self.f = ifftshift(self.f)
+          if 'filter_f_psi' in self.paraDict: self.f = self.f * self.paraDict['filter_f_psi']
+          self.fft_inverse.update_arrays(self.f,self.r)
+          self.fft_inverse.execute();          
+          return self.r / self.N_tot
+
+      #################### position correction ####################
+      def gradPositionCorrection(self, probe, O, ind_dp, delta_psi):
+          dx_O, dy_O = self.getObjectGradient(O)
+          
+          dx_OP = dx_O*probe
+          shift_x = np.sum(real(conj(dx_OP)*delta_psi))/np.sum(abs(dx_OP)**2)
+          
+          dy_OP = dy_O*probe
+          shift_y = np.sum(real(conj(dy_OP)*delta_psi))/np.sum(abs(dy_OP)**2)
+
+          #update position
+          #print(shift_y)
+          self.paraDict['ppY'][ind_dp] = self.paraDict['ppY'][ind_dp] + shift_y*self.dx_y 
+          self.paraDict['ppX'][ind_dp] = self.paraDict['ppX'][ind_dp] + shift_x*self.dx_x
+
+          #position = pi + pf = integer + fraction
+          py_i = np.round(self.paraDict['ppY'][ind_dp] / self.dx_y)
+          self.py_f[ind_dp] = self.paraDict['ppY'][ind_dp] - py_i * self.dx_y 
+          px_i = np.round(self.paraDict['ppX'][ind_dp] / self.dx_x)
+          self.px_f[ind_dp] = self.paraDict['ppX'][ind_dp] - px_i * self.dx_x
+
+          #calculate ROI indices in the whole fov
+          self.ind_x_lb_s[ind_dp] = (px_i - floor(self.paraDict['N_roi']/2.0) + self.center_index_image).astype(np.int)
+          self.ind_x_ub_s[ind_dp] = (px_i + ceil(self.paraDict['N_roi']/2.0) + self.center_index_image).astype(np.int)
+          self.ind_y_lb_s[ind_dp] = (py_i - floor(self.paraDict['N_roi']/2.0) + self.center_index_image).astype(np.int)
+          self.ind_y_ub_s[ind_dp] = (py_i + ceil(self.paraDict['N_roi']/2.0) + self.center_index_image).astype(np.int)
+
+      #################### initialization ####################
+      def initializeDataError(self):
+          if 's' in self.paraDict:
+              if self.paraDict['Niter'] > self.paraDict['s'].size:
+                  s = self.pad(paraDict['s'].flatten(),(0, self.paraDict['Niter']-self.paraDict['s'].size),'constant')
+              else:
+                  s = self.paraDict['s'].flatten()
+          else:
+              s = zeros(self.paraDict['Niter']) #data error
+          if 'dp_error_old' in self.paraDict:
+              dp_error_old  = self.paraDict['dp_error_old'].flatten()
+          else:
+              dp_error_old = full(self.paraDict['N_scan'], np.inf)
+          return s, dp_error_old
+
+      def initializeScanPositions(self):
+          self.center_index_image = int(self.paraDict['N_image']/2)
+          self.dx_x = 1.0/(self.paraDict['dk_x'] * self.paraDict['N_roi'])
+          self.dx_y = 1.0/(self.paraDict['dk_y'] * self.paraDict['N_roi'])
+
+          #position = pi + pf = integer + fraction
+          py_i = np.round(self.paraDict['ppY'] / self.dx_y)
+          self.py_f = self.paraDict['ppY'] - py_i * self.dx_y 
+          px_i = np.round(self.paraDict['ppX'] / self.dx_x)
+          self.px_f = self.paraDict['ppX'] - px_i * self.dx_x
+
+          #calculate ROI indices in entire image
+          self.ind_x_lb_s = (px_i - floor(self.paraDict['N_roi']/2.0) + self.center_index_image).astype(np.int)
+          self.ind_x_ub_s = (px_i + ceil(self.paraDict['N_roi']/2.0) + self.center_index_image).astype(np.int)
+          self.ind_y_lb_s = (py_i - floor(self.paraDict['N_roi']/2.0) + self.center_index_image).astype(np.int)
+          self.ind_y_ub_s = (py_i + ceil(self.paraDict['N_roi']/2.0) + self.center_index_image).astype(np.int)
+
+      def initializeDiffractionPatterns(self):
+          center_index_roi = int(self.paraDict['N_roi']/2)
+          #calculate dp indices in ROI image
+          self.ind_dp_lb = int(-floor(self.paraDict['N_dp']/2.0) + center_index_roi)
+          self.ind_dp_ub = int(ceil(self.paraDict['N_dp']/2.0) + center_index_roi)
+
+      def createFourierCoord(self, N):
+          kx = linspace(-floor(N/2.0),ceil(N/2.0)-1,N)
+          kx = ifftshift(kx)
+          [self.kX, self.kY] = meshgrid(kx,kx)
+          self.kX = self.kX * self.paraDict['dk_x']
+          self.kY = self.kY * self.paraDict['dk_y']
+
+      def printStatus(self, timeLeft, iter, extraMessage=''):
+          timeLeftMin, timeLeftSec = divmod(timeLeft, 60)
+          timeLeftHour, timeLeftMin = divmod(timeLeftMin, 60)
+          print(self.paraDict['printID'] + extraMessage + '-Iter:%d Time remain: %02d:%02d:%02d' %(iter,timeLeftHour,timeLeftMin,timeLeftSec))
+
+      #################### FFT ####################
+      def initializeFFTW(self, N):
+          self.f = pyfftw.empty_aligned((N,N),dtype='complex128',n=16) 
+          self.r = pyfftw.empty_aligned((N,N),dtype='complex128',n=16)
+          self.N_tot = N*N
+          self.fft_forward = pyfftw.FFTW(self.r, self.f, axes=(0,1))
+          self.fft_inverse = pyfftw.FFTW(self.f, self.r, direction='FFTW_BACKWARD', axes=(0,1))
+
+      def shift(self, func, px, py):
+          """Shift function via FFT"""
+          self.r[:,:] = ifftshift(func)
+          self.fft_forward.update_arrays(self.r, self.f)
+          self.fft_forward.execute()
+          self.f = self.f*exp(-2*pi*1j*px*self.kX)*exp(-2*pi*1j*py*self.kY)
+          self.fft_inverse.update_arrays(self.f, self.r)
+          self.fft_inverse.execute();
+          return fftshift(self.r) / self.N_tot #fix normalization
+
+      ####################  ####################
+      def proj(u, v):
+          return u * np.vdot(u,v) / np.vdot(u,u)  
+
+      def gramschmidt(V):
+          U = np.copy(V)
+          for i in range(1, V.shape[0]):
+              for j in range(i):
+                  U[i,:] -= proj(U[j,:], V[i,:])
+          return U
+
+      def getObjectGradient(self, O):
+          Ny, Nx = O.shape
+          kx = fftshift(linspace(0,Nx-1,Nx)*1.0/Nx-0.5)
+          ky = fftshift(linspace(0,Ny-1,Ny)*1.0/Ny-0.5)
+          [kX, kY] = meshgrid(kx,ky)
+
+          O_fx = fft(O,axis=1)
+          O_fy = fft(O,axis=0)
+
+          O_dx = ifft(O_fx*kX*2j*pi,axis=1)
+          O_dy = ifft(O_fy*kY*2j*pi,axis=0)
+
+          return O_dx, O_dy
+

ptycho/python/recon_example.py

@@ -0,0 +1,126 @@
+import numpy as np
+import os #for change directory
+import scipy.io as sio #for read/write matlab file
+import scipy.ndimage.filters as sfilter
+import scipy.ndimage
+import scipy.misc
+
+import ptycho_recon.ptycho as pty
+import ptycho_recon.utility_function as utils
+import ptycho_recon.probe as probe
+
+#import h5py
+
+from numpy import *
+import time
+
+##################################### read data ########################################
+print("load data")
+#data_name = 'data_ws2_wse2_80keV_09_roi2.mat'
+data_name = 'data_mos2_sample.mat'
+currentdir = os.getcwd() #current directory
+#os.chdir(data_dir) #change to data directory
+data = sio.loadmat(data_name)
+dp = data['dp']*1.0
+
+sx = 0 if not 'sx' in data else int(squeeze(data['sx']))
+sy = 0 if not 'sy' in data else int(squeeze(data['sy']))
+#if not 'sx' in data else int(squeeze(data['sx']))
+#sy = 0 if not data.has_key('sy') else int(squeeze(data['sy']))
+######################################### Parameters ##################################### 
+ADU_background_cutoff = 20
+ADU_electronCount_ratio = 149.0
+
+N_roi = 128
+
+##############################################################################################
+N_dp = dp.shape[0]
+voltage = squeeze(data['voltage']) #kev
+alpha_max = squeeze(data['alpha_max']) #mrad
+df = squeeze(data['df'])
+cs = squeeze(data['cs']) #angstrom
+#scanStepSize_x = squeeze(data['scanStepSize_x']) #angstrom
+#scanStepSize_y = squeeze(data['scanStepSize_y'])
+scanStepSize_x = 0.21 #angstrom
+scanStepSize_y = 0.21
+
+dk = squeeze(data['dk'])
+
+print("dk =", dk)
+dx = 1.0/dk/N_roi
+print("dx =", dx)
+
+################################## data processing ##########################################
+rot_angle_d = 30
+
+print("processing data")
+result_dir_extra = "/preprocessCBED"
+dp, result_dir_extra = utils.transpose_cbed(dp, result_dir_extra)
+dp, result_dir_extra = utils.background_removal(dp, ADU_background_cutoff, result_dir_extra)
+
+print("recon data size:", dp.shape)
+################################## make initial probe function ##################################
+probe_init = probe.STEMprobe()
+probe_init.df = df
+probe_init.cs = cs
+probe_init.alpha_max = alpha_max
+probe_init.voltage = voltage
+
+################################## calculate scan positions ####################################
+#calculate scan positions
+N_scan_y = dp.shape[2]
+N_scan_x = dp.shape[3]
+
+ppX, ppY, result_dir_extra = utils.calculate_scan_positions(N_scan_x, N_scan_y, scanStepSize_x, scanStepSize_y, rot_angle_d, result_dir_extra)
+
+Ny_max = max([abs(round(np.min(ppY)/dx)-floor(N_roi/2.0)), abs(round(np.max(ppY)/dx)+ceil(N_roi/2.0))])*2+1
+Nx_max = max([abs(round(np.min(ppX)/dx)-floor(N_roi/2.0)), abs(round(np.max(ppX)/dx)+ceil(N_roi/2.0))])*2+1
+N_image = int(max([Ny_max,Nx_max]))+20
+print("Image size:", N_image)
+'''
+######### reshape dp and scan positions #########
+N_scan_y = dp.shape[2]
+N_scan_x = dp.shape[3]
+N_scan_tot = N_scan_y * N_scan_x
+
+dp_temp = zeros((N_scan_tot,dp.shape[0],dp.shape[1]))
+ppX = ppX.reshape(N_scan_tot)
+ppY = ppY.reshape(N_scan_tot)
+
+for i in range(N_scan_y):
+    for j in range(N_scan_x):
+        index = i*N_scan_x + j
+        dp_temp[index,:,:] = sqrt(dp[:,:,i,j])
+'''
+########################################### reconstruction #####################################
+print('Reconstruction')
+reconObject = pty.ptycho(dp, dk, probe_init, ppX, ppY)
+reconObject.paraDict['N_image'] = N_image
+reconObject.paraDict['N_roi'] = N_roi
+reconObject.paraDict['Niter'] = 50
+reconObject.paraDict['Niter_update_probe'] = 0
+reconObject.paraDict['Niter_save'] = 5
+
+reconObject.paraDict['rotationAngle']  = rot_angle_d
+
+reconObject.paraDict['printID'] = 'MoS2'
+
+############## position correction ##############
+reconObject.paraDict['Niter_update_position'] = 30
+
+############## mixed-states recon ##############
+reconObject.paraDict['N_probe'] = 2
+reconObject.paraDict['Niter_update_states'] = 10
+
+result_dir = currentdir + "/mos2"
+result_dir_extra = reconObject.initialize(result_dir)
+
+start_time = time.time()
+reconObject.recon() #start reconstruction
+total_time = time.time() - start_time
+
+timeLeftMin, timeLeftSec = divmod(total_time, 60)
+timeLeftHour, timeLeftMin = divmod(timeLeftMin, 60)
+print('Total recon time: %02d:%02d:%02d' %(timeLeftHour,timeLeftMin,timeLeftSec))
+
+