initial commit

This commit is contained in:
2026-08-07 15:56:42 +09:00
commit 91ad25aca9
1012 changed files with 159314 additions and 0 deletions
@@ -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