def run_benchmarks(files, dark_sample, num_runs):
"""Run benchmarks using the files given in the arguments
Example: "python benchmark_stem_image.py -d darksample.dat /path/to/data/data*.dat"
"""
if len(files) == 0:
sys.exit('No files found')
times = []
for i in range(num_runs):
# You can use a command like this to clear the caches, where
# `clearcaches` should be a bash script.
# subprocess.Popen('clearcaches').wait()
start = time.time()
reader = io.reader(dark_sample, version=io.FileVersion.VERSION2)
dark = image.calculate_average(reader)
reader = io.reader(files, version=io.FileVersion.VERSION2)
frame_events = image.electron_count(reader, dark, scan_width=40,
scan_height=40)
end = time.time()
times.append(end - start)
print('Run ' + str(len(times)) + ': {:0.2f} seconds'.format(times[-1]))
print('Number of runs was:', num_runs)
print('Average time: {:0.2f} seconds'.format(sum(times) / num_runs))
def main(input_path, output_path, scan_num, threshold):
if input_path is None:
input_path = Path(f'/mnt/nvmedata1/temp/data_scan{scan_num}.h5')
if output_path is None:
output_path = Path(f'/mnt/hdd1/data_scan{scan_num}_th{threshold}_electrons.h5')
print(output_path)
dark0 = np.zeros((576,576))
print('Opening: {}'.format(input_path))
with h5py.File(input_path, 'r') as f0:
sReader = stio.reader(f0)
print('start counting')
t0 = time.time()
ee = stim.electron_count(sReader, dark0, number_of_samples=1200,
verbose=True,
xray_threshold_n_sigma=175,
background_threshold_n_sigma=threshold)
t1 = time.time()
print('total time = {}'.format(t1-t0))
ii = 0
while output_path.exists():
ii += 1
output_path = Path(output_path.stem + '_{:03d}'.format(ii))
print('Saving to {}'.format(output_path))
stio.save_electron_counts(str(output_path), ee)
def run_benchmarks(files, num_runs):
"""Run benchmarks using the files given in the arguments
Example: "python benchmark_stem_image.py /path/to/data/data*.dat"
"""
if len(files) == 0:
sys.exit('No files found')
times = []
for i in range(num_runs):
start = time.time()
reader = io.reader(files)
img = image.create_stem_image(reader,
40,
288,
scan_dimensions=(160, 160))
end = time.time()
times.append(end - start)
print('Run ' + str(len(times)) + ': {:0.2f} seconds'.format(times[-1]))
print('Number of runs was:', num_runs)
print('Average time: {:0.2f} seconds'.format(sum(times) / num_runs))
def get_worker_reader(path, version):
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
world_size = comm.Get_size()
files = glob.glob(path)[rank::world_size]
if (len(files) == 0):
return None
return io.reader(files, version=int(version))
def main(argv):
print_help()
dataDir = argv[0]
outDir = argv[1]
# create output directory if it does not exit
if not os.path.exists(outDir):
os.mkdir(outDir)
print('Output directory', outDir, 'created')
else:
print('Output directory', outDir, 'already exists')
# get the all the data files
files = []
for root, dirs, fs in os.walk(dataDir):
for f in fs:
files.append(os.path.join(root, f))
# inner and outer radius of mask
inner_radii = [0, 40]
outer_radii = [288, 288]
# file reader
reader = io.reader(files)
# generate histograms
numBins = 100
print('Generating histograms for input data')
all_bins, all_freqs = image.create_stem_histogram(numBins,
reader,
inner_radii,
outer_radii,
width=160,
height=160)
# plot histogram
for i in range(len(all_bins)):
# obtain current bins and freq
bins = [str(element) for element in all_bins[i]]
freq = all_freqs[i]
# init figure
fig = plt.figure(1, figsize=(16, 8))
myHist = fig.add_subplot(111)
# plt.bar considers the left boundary
x = np.arange(numBins + 1)
myHist.bar(x[:-1], freq, align='edge')
plt.xticks(x[::20], bins[::20])
plt.title('Histogram of STEM image with inner radius = ' +
str(inner_radii[i]) + ', outer radius = ' +
str(outer_radii[i]))
plt.xlabel('Value')
plt.ylabel('Frequency')
# save to local
suffix = str(inner_radii[i]) + '_' + str(outer_radii[i]) + '.png'
plt.savefig(outDir + '/histogram_' + suffix)
plt.show()
def main(files, dark_file, output_file):
"""
Example of calculating maximum diffracton pattern using MPI.
"""
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
world_size = comm.Get_size()
if (world_size > len(files)):
if rank == 0:
print('Error: number of MPI processes,', world_size, ', exceeds',
'the number of files:', len(files))
return
comm.Barrier()
start = MPI.Wtime()
# Split up the files among processes
files = get_files(files)
# Create local maximum diffraction pattern
reader = io.reader(files, version=io.FileVersion.VERSION3)
mdp = image.maximum_diffraction_pattern(reader)
mdp = mdp.data
# Now reduce to root
global_mdp = np.zeros_like(mdp)
comm.Reduce(mdp, global_mdp, op=MPI.MAX)
if dark_file is not None:
reader = io.reader(dark_file, version=io.FileVersion.VERSION3)
dark = image.calculate_average(reader)
mdp -= dark
comm.Barrier()
end = MPI.Wtime()
if rank == 0:
print('time: %s' % (end - start))
if rank == 0:
# Write out the MDP
np.save(output_file, global_mdp)
def main(output_path, scan_num, threshold):
if output_path is None:
output_path = Path(f'/mnt/hdd1/data_scan{scan_num}_th{threshold}_electrons.h5')
scanName = f'data_scan{scan_num}_'
# Setup the data drives
drives = []
for ii in range(1,5):
drives.append( (Path('/mnt/nvmedata{}/'.format(ii))))
print('Looking for files in:')
for d in drives:
print(d)
dark0 = np.zeros((576,576))
iFiles = []
for drive in drives:
files = drive.glob(scanName + '*.data')
for f in files:
iFiles.append(str(f))
iFiles = sorted(iFiles)
# Electron count the data
sReader = stio.reader(iFiles,stio.FileVersion.VERSION5, backend='thread')
print('start counting')
t0 = time.time()
ee = stim.electron_count(sReader,dark0,number_of_samples=1200,
verbose=False,threshold_num_blocks=20,
xray_threshold_n_sigma=175,
background_threshold_n_sigma=threshold)
t1 = time.time()
print('total time = {}'.format(t1-t0))
ii = 0
while output_path.exists():
ii += 1
output_path = Path(output_path.stem + '_{:03d}'.format(ii))
print('Saving to {}'.format(output_path))
stio.save_electron_counts(str(output_path), ee)
for drive in drives:
files = drive.glob(scanName + '*.data')
for f in files:
iFiles.append(str(f))
# Sort the files
iFiles = sorted(iFiles)
if args.verbose:
print('Number of files = {}'.format(len(iFiles)))
# Electron count the data
if args.verbose:
#print('NOTE: Using file version 4')
print(f'backend = {backend}')
sReader = stio.reader(iFiles, stio.FileVersion.VERSION5, backend=backend)
if args.verbose:
print('start counting')
t0 = time.time()
ee = stim.electron_count(sReader,
dark0,
gain=gain0,
number_of_samples=1200,
verbose=args.verbose,
threshold_num_blocks=20,
xray_threshold_n_sigma=175,
background_threshold_n_sigma=threshold)
t1 = time.time()
full_time = t1 - t0
max = np.max(stem_image_data)
stem_image_data = stem_image_data.reshape((160, 160))
stem_image_data = np.interp(stem_image_data, [min, max], [0, 256])
stem_image_data = stem_image_data.astype(np.uint8)
img = Image.fromarray(stem_image_data)
img.save(name)
stem_image_data_day = np.zeros((160 * 160, ), dtype=float)
stem_image_data_night = np.zeros((160 * 160, ), dtype=float)
mask_size = 20
files = []
for f in glob.glob('/data/4dstem/smallScanningDiffraction/data*.dat'):
files.append(f)
inner_radii = [0, 40]
outer_radii = [288, 288]
reader = io.reader(files)
imgs = image.create_stem_images(reader,
inner_radii,
outer_radii,
width=160,
height=160)
for i, img in enumerate(imgs):
suffix = str(inner_radii[i]) + '_' + str(outer_radii[i]) + '.png'
save_img(img, 'img_' + suffix)
def main(files, center, inner_radii, outer_radii):
center = center.split(',')
if len(center) != 2:
raise click.ClickException(
'Center must be of the form: center_x,center_y.')
center = tuple(int(x) for x in center)
inner_radii = inner_radii.split(',')
outer_radii = outer_radii.split(',')
if len(inner_radii) != len(outer_radii):
raise click.ClickException(
'Number of inner and outer radii must match')
inner_radii = [int(x) for x in inner_radii]
outer_radii = [int(x) for x in outer_radii]
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
world_size = comm.Get_size()
comm.Barrier()
start = MPI.Wtime()
files = get_files(files)
# Create local sum
reader = io.reader(files, version=io.FileVersion.VERSION3)
# Get the scan image size
block = reader.read()
scan_dimensions = block.header.scan_dimensions
reader.reset()
local_stems = image.create_stem_images(reader,
inner_radii,
outer_radii,
center=center)
# Now reduce to root
global_stems = [
np.zeros(scan_dimensions[0] * scan_dimensions[1], dtype='uint64')
for _ in range(len(inner_radii))
]
for i in range(len(inner_radii)):
comm.Reduce(local_stems[i], global_stems[i], op=MPI.SUM)
# Save out the image
if rank == 0:
for global_stem, inner, outer in zip(global_stems, inner_radii,
outer_radii):
thr = global_stem[:] > 0
vmin = global_stem[thr].min()
cmap = plt.cm.viridis
norm = plt.Normalize(vmin=vmin)
stem_image = cmap(
norm(
global_stem.reshape(scan_dimensions[1],
scan_dimensions[0])))
filename = 'stem_%d_%d.png' % (inner, outer)
plt.imsave(filename, stem_image)
comm.Barrier()
end = MPI.Wtime()
if rank == 0:
print('time: %s' % (end - start))
def make_stem_hdf5(files, dark_sample, width, height, inner_radius,
outer_radius, reader_version, dark_reader_version, save_raw,
zip_raw, output):
"""Make an HDF5 file containing a STEM image
Example: "python create_hdf5.py -d darksample.dat /path/to/data/data*.dat"
"""
if len(files) == 0:
sys.exit('No files found')
if reader_version == 1:
reader_version = io.FileVersion.VERSION1
elif reader_version == 2:
reader_version = io.FileVersion.VERSION2
else:
sys.exit('Unknown reader version:', reader_version)
if dark_reader_version == 1:
dark_reader_version = io.FileVersion.VERSION1
elif dark_reader_version == 2:
dark_reader_version = io.FileVersion.VERSION2
else:
sys.exit('Unknown dark reader version:', dark_reader_version)
scan_dimensions = (width, height)
if dark_sample:
reader = io.reader(dark_sample, version=dark_reader_version)
dark = image.calculate_average(reader)
else:
# Get the frame dimensions from a block header, and use zeros
# for the dark sample.
reader = io.reader(files, version=reader_version)
frame_dimensions = reader.read().header.frame_dimensions
dark = np.zeros(frame_dimensions)
reader = io.reader(files, version=reader_version)
data = image.electron_count(reader, dark, scan_dimensions=scan_dimensions)
frame_dimensions = data.frame_dimensions
inner_radii = [0, inner_radius]
outer_radii = [outer_radius, outer_radius]
names = ['bright', 'dark']
reader.reset()
imgs = image.create_stem_images(reader,
inner_radii,
outer_radii,
scan_dimensions=scan_dimensions)
io.save_electron_counts(output, data)
io.save_stem_images(output, imgs, names)
if save_raw:
reader.reset()
# In order to avoid two copies of the data, we must allocate
# space for the large numpy array first.
raw_data = np.zeros((np.prod(scan_dimensions), frame_dimensions[1],
frame_dimensions[0]),
dtype=np.uint16)
# Make sure the frames are sorted in order
for block in reader:
for i in range(len(block.header.image_numbers)):
num = block.header.image_numbers[i]
raw_data[num] = block.data[i]
io.save_raw_data(output, raw_data, zip_data=zip_raw)
from stempy import io
import glob
for i, f in enumerate(glob.glob('/data/4dstem/smallScanningDiffraction/data*.dat')):
print(f)
reader = io.reader(f)
reader.process(url='http://localhost:5000', stream_id=i)
from stempy import io
import glob
import os
node_id = os.environ['SLURM_NODEID']
path = '/global/project/projectdirs/ncemhub/simData/smallScanningDiffraction/data00%s.dat' % node_id.zfill(
2)
reader = io.reader(path)
reader.process(url='http://128.55.206.19:60048/', stream_id=int(node_id))
def make_stem_hdf5(files, dark_sample, width, height, inner_radius,
outer_radius, reader_version, dark_reader_version, save_raw,
zip_raw, output):
"""Make an HDF5 file containing a STEM image
Example: "python create_hdf5.py -d darksample.dat /path/to/data/data*.dat"
"""
if len(files) == 0:
sys.exit('No files found')
if reader_version == 1:
reader_version = io.FileVersion.VERSION1
elif reader_version == 2:
reader_version = io.FileVersion.VERSION2
else:
sys.exit('Unknown reader version:', reader_version)
if dark_reader_version == 1:
dark_reader_version = io.FileVersion.VERSION1
elif dark_reader_version == 2:
dark_reader_version = io.FileVersion.VERSION2
else:
sys.exit('Unknown dark reader version:', dark_reader_version)
reader = io.reader(dark_sample, version=dark_reader_version)
dark = image.calculate_average(reader)
reader = io.reader(files, version=reader_version)
frame_events = image.electron_count(reader,
dark,
scan_width=width,
scan_height=height)
# Read one block in to get the detector frames
reader.reset()
block = reader.read()
detector_nx = block.header.frame_width
detector_ny = block.header.frame_height
inner_radii = [0, inner_radius]
outer_radii = [outer_radius, outer_radius]
names = ['bright', 'dark']
reader.reset()
imgs = image.create_stem_images(reader,
inner_radii,
outer_radii,
width=width,
height=height)
io.save_electron_counts(output, frame_events, width, height, detector_nx,
detector_ny)
io.save_stem_images(output, imgs, names)
if save_raw:
reader.reset()
blocks = [block for block in reader]
raw_data = np.concatenate([block.data for block in blocks])
io.save_raw_data(output, raw_data, zip_raw)
def main(files, dark_file, center, inner_radii, outer_radii, output_file,
generate_sparse):
center = center.split(',')
if len(center) != 2:
msg = 'Center must be of the form: center_x,center_y.'
raise click.ClickException(msg)
center = tuple(int(x) for x in center)
inner_radii = inner_radii.split(',')
outer_radii = outer_radii.split(',')
if len(inner_radii) != len(outer_radii):
msg = 'Number of inner and outer radii must match'
raise click.ClickException(msg)
inner_radii = [int(x) for x in inner_radii]
outer_radii = [int(x) for x in outer_radii]
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
world_size = comm.Get_size()
if (world_size > len(files)):
if rank == 0:
print('Error: number of MPI processes,', world_size, ', exceeds',
'the number of files:', len(files))
return
comm.Barrier()
start = MPI.Wtime()
if dark_file is not None:
# Every process will do the dark field reference average for now
reader = io.reader(dark_file, version=io.FileVersion.VERSION3)
dark = image.calculate_average(reader)
else:
dark = np.zeros((576, 576))
# Split up the files among processes
files = get_files(files)
# Create local electron count
reader = io.reader(files, version=io.FileVersion.VERSION3)
electron_counted_data = image.electron_count(reader, dark, verbose=True)
local_frame_events = electron_counted_data.data
# Now reduce to root
global_frame_events = reduce_to_root_method1(local_frame_events)
# global_frame_events = reduce_to_root_method2(local_frame_events)
comm.Barrier()
end = MPI.Wtime()
if rank == 0:
print('time: %s' % (end - start))
if rank == 0:
# Create new electron counted data with the global frame events
data = namedtuple('ElectronCountedData',
['data', 'scan_dimensions', 'frame_dimensions'])
data.data = global_frame_events
data.scan_dimensions = electron_counted_data.scan_dimensions
data.frame_dimensions = electron_counted_data.frame_dimensions
# Write out the HDF5 file
io.save_electron_counts(output_file, data)
if generate_sparse:
# Save out the sparse image
stem_imgs = image.create_stem_images(data,
inner_radii,
outer_radii,
center=center)
for i, img in enumerate(stem_imgs):
fig, ax = plt.subplots(figsize=(12, 12))
ax.matshow(img)
name = 'sparse_stem_image_' + str(i) + '.png'
plt.savefig(name, dpi=300)