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 test_simple():
# Need at least 32 blocks since the blocksize in
# io.get_hdf5_reader() is hard-coded to that.
shape = (32, 576, 576)
data = np.arange(np.prod(shape)).reshape(shape)
dataset_name = 'frames'
# Stempy uses uint16. Let's prevent overflow...
data %= np.iinfo(np.uint16).max
mean = data.mean(axis=0)
# Fake stuff that io.get_hdf5_reader() needs...
fake_scan_name = 'stem/images'
fake_scan = np.zeros((4, 8, 1))
temp = tempfile.NamedTemporaryFile(delete=False)
file_name = temp.name
try:
with h5py.File(file_name, 'w') as wf:
wf.create_dataset(dataset_name, data=data)
wf.create_dataset(fake_scan_name, data=fake_scan)
with h5py.File(file_name, 'r') as rf:
reader = io.get_hdf5_reader(rf)
average = image.calculate_average(reader)
finally:
del temp
if os.path.exists(file_name):
os.remove(file_name)
# These should be equal...
assert np.array_equal(average, mean)
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 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)
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)