initial commit

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))
+
+