mirror of
https://github.com/c-sooyoung/fold_slice.git
synced 2026-09-18 00:49:09 +09:00
127 lines
4.2 KiB
Python
127 lines
4.2 KiB
Python
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))
|
|
|
|
|