3.2. Example of simulation parameter generation
3.2.1. Prerequisites for running this example
Prior to running any scripts or Jupyter notebooks in the directory
<root>/examples, where <root> is the root of the prismatique
repository, a set of Python libraries need to be installed in the Python
environment within which any such scripts or Jupyter notebooks are to be
executed. See this page
for instructions on how to do so.
3.2.2. Example description
In this example, we generate simulation parameters that can be used: to generate the potential slices for the bilayer \(\text{MoS}_2\) sample that we defined in the Example of atomic coordinate generation page; to generate the \(S\)-matrices of said sample; to simulate a STEM experiment involving said sample; and to simulate a HRTEM experiment involving said sample.
See the documentation for the class prismatique.discretization.Params
for a discussion on potential slices, and the subpackage prismatique.stem
for a discussion on \(S\)-matrices.
3.2.3. Code
Below is the code that generates the aforementioned simulation parameters. You
can also find the same code in the file
<root>/examples/sim_param_generator/generate.py of the repository, where
<root> is the root of the prismatique repository. To run the script from
the terminal, change into the directory containing said script, and then issue
the following command:
python generate.py
The output files generated by this script are saved in the directory
<root>/examples/data/sim_param_generator_output, where <root> is the
root of the prismatique repository. To analyze the output, use the Jupyter
notebook
<root>/examples/output_data_analysis_notebooks/analyzing_sim_params.ipynb.
If you would like to modify this script for your own work, it is recommended that you copy the original script and save it elsewhere outside of the git repository so that the changes made are not tracked by git.
# -*- coding: utf-8 -*-
# Copyright 2024 Matthew Fitzpatrick.
#
# This program is free software: you can redistribute it and/or modify it under
# the terms of the GNU General Public License as published by the Free Software
# Foundation, version 3.
#
# This program is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along with
# this program. If not, see <https://www.gnu.org/licenses/gpl-3.0.html>.
"""An example of generating simulation parameters to be used in future
simulations.
See the link
https://mrfitzpa.github.io/prismatique/examples/sim-param-generator/generate.html
for a description of the example.
A NOTE BEFORE STARTING
----------------------
To run this script from the terminal as is, i.e. without modifications, change
into the directory containing said script, and then issue the following
command::
python generate.py
The output files generated by this script are saved in the directory
``<root>/examples/data/sim-param-generator-output``, where ``<root>`` is the
root of the ``prismatique`` repository. To analyze the output, use the Jupyter
notebook
``<root>/examples/output-data-analysis-notebooks/analyzing-sim-params.ipynb``.
If you would like to modify this script for your own work, it is recommended
that you copy the original script and save it elsewhere outside of the git
repository so that the changes made are not tracked by git.
"""
#####################################
## Load libraries/packages/modules ##
#####################################
# For making directories and creating path objects.
import pathlib
# For spawning shell processes and getting the path to current script.
import os
# For using special mathematical constants.
import numpy as np
# For modelling electron guns, lenses, and probes.
import embeam
# For postprocessing diffraction patterns (DPs) and images.
import empix
# For setting the simulation parameters.
import prismatique
###############################################
## Define classes, functions, and contstants ##
###############################################
###########################
## Define error messages ##
###########################
#########################
## Main body of script ##
#########################
msg = ("Generating the simulation parameter sets for the various simulated "
"experiments involving the MoS2 sample...")
print(msg)
print()
# First, we specify the "worker" parameters, which are the parameters related to
# how computations are delegated to CPUs and (if available) GPUs. See the
# documentation for the classes :class:`prismatique.worker.cpu.Params`,
# :class:`prismatique.worker.gpu.Params`, and :class:`prismatique.worker.Params`
# for details on the various worker parameters.
# The CPU worker parameters.
kwargs = {"enable_workers": True,
"num_worker_threads": 32,
"batch_size": 1,
"early_stop_count": 100}
cpu_params = prismatique.worker.cpu.Params(**kwargs)
# The GPU worker parameters. Note that if no GPUs are available then the
# following parameters are essentially ignored.
kwargs = {"num_gpus": 4,
"batch_size": 1,
"data_transfer_mode": "auto",
"num_streams_per_gpu": 3}
gpu_params = prismatique.worker.gpu.Params(**kwargs)
# Put the CPU and GPU worker parameters together.
kwargs = {"cpu_params": cpu_params, "gpu_params": gpu_params}
worker_params = prismatique.worker.Params(**kwargs)
# Next, we specify the simulation parameters related to the discretization of
# real-space and Fourier/k-space. See the documentation for the class
# :class:`prismatique.discretization.Params` for details on these parameters and
# how real-space and k-space are discretized.
sample_supercell_reduced_xy_dims_in_pixels = (32, 32)
kwargs = {"z_supersampling": \
16,
"sample_supercell_reduced_xy_dims_in_pixels": \
sample_supercell_reduced_xy_dims_in_pixels,
"interpolation_factors": \
(1, 1),
"num_slices": \
16}
discretization_params = prismatique.discretization.Params(**kwargs)
# Next, we specify the simulation parameters related to the thermal properties
# of the sample and its environment. See the documentation for the class
# :class:`prismatique.thermal.Params` for details on these parameters and how
# [among other things] thermal effects are modelled.
kwargs = {"enable_thermal_effects": True,
"num_frozen_phonon_configs_per_subset": 5,
"num_subsets": 2,
"rng_seed": 555}
thermal_params = prismatique.thermal.Params(**kwargs)
# Next, we specify the file which contains the zero-temperature expectation
# value of the atomic coordinates, as well as the root-mean-square
# x-displacement, for each atom in the unit cell of our sample of interest. If
# the file does not already exist, we must generate it. This can be done using
# the example script ``<root>/examples/atomic-coord-generator/generate.py``,
# where ``<root>`` is the root of the ``prismatique`` repository.
path_to_current_script = pathlib.Path(os.path.realpath(__file__))
path_to_examples_dir = str(path_to_current_script.parent.parent)
path_to_data_dir = path_to_examples_dir + "/data"
atomic_coords_filename = path_to_data_dir + "/atomic_coords.xyz"
try:
prismatique.sample.check_atomic_coords_file_format(atomic_coords_filename)
except:
msg = ("Need to generate the file {} from which to read-in the "
"atomic coordinates of the MoS2 "
"sample.".format(atomic_coords_filename))
print(msg)
print()
script_to_execute = (path_to_examples_dir
+ "/atomic_coord_generator/generate.py")
os.system("python " + script_to_execute)
msg = ("Resuming the generation of the simulation parameter sets for the "
"various simulated experiments involving the MoS2 sample...")
print(msg)
print()
# In generating the atomic coordinates discussed above, we made use of a
# quantity called the "atomic potential extent", which is one of the simulation
# parameters related to the modelling of the sample. The script referenced above
# that generates the atomic coordinates also saves the value of the atomic
# potential extent, which we read-in here for consistency between scripts.
with open(atomic_coords_filename, "r") as file_obj:
atomic_potential_extent = float(file_obj.readline().split("=")[-1])
# Next, we specify the simulation parameters related to the modelling of the
# sample. See the documentation for the class
# :class:`prismatique.sim.ModelParams` for details on these parameters.
kwargs = {"atomic_coords_filename": atomic_coords_filename,
"unit_cell_tiling": (1, 1, 1),
"discretization_params": discretization_params,
"atomic_potential_extent": atomic_potential_extent,
"thermal_params": thermal_params}
sample_model_params = prismatique.sample.ModelParams(**kwargs)
# Next, we specify the simulation parameters related to the modelling of the
# electron gun. See the documentation for the class
# :class:`embeam.gun.ModelParams` for details on these parameters.
kwargs = {"mean_beam_energy": 20, # In keV.
"intrinsic_energy_spread": 0.6e-3, # In keV.
"accel_voltage_spread": 0.0} # In keV.
gun_model_params = embeam.gun.ModelParams(**kwargs)
# Next, we specify a set of coherent lens aberrations, subject to which are the
# probe forming and objective lenses. See the documentation for the class
# :class:`embeam.coherent.Aberration` for details on aberration model
# parameters.
Delta_f = -4000 # Defocus strength in Å. Negative sign indicates underfocus.
C_s = 8 # Spherical aberration strength in mm.
# Beam wavelength in Å.
wavelength = embeam.wavelength(gun_model_params.core_attrs["mean_beam_energy"])
# Aberration strengths in dimensionless units.
C_2_0_mag = (2*np.pi*Delta_f)/(2*wavelength)
C_4_0_mag = (2*np.pi*(C_s*1e7))/(4*wavelength)
defocus_aberration = embeam.coherent.Aberration(m=2,
n=0,
C_mag=C_2_0_mag,
C_ang=0)
spherical_aberration = embeam.coherent.Aberration(m=4,
n=0,
C_mag=C_4_0_mag,
C_ang=0)
coherent_aberrations = (defocus_aberration, spherical_aberration)
# Next, we specify the simulation parameters related to the modelling of both
# the probe forming and the objective lenses. For simplicity, we use the same
# model for both lenses. See the documentation for the class
# :class:`embeam.lens.ModelParams` for details on these parameters.
kwargs = {"coherent_aberrations": coherent_aberrations,
"chromatic_aberration_coef": 8, # In mm.
"mean_current": 50, # In pA.
"std_dev_current": 0.0} # In pA.
lens_model_params = embeam.lens.ModelParams(**kwargs)
# The simulation parameters that we have specified above are common to both the
# STEM and HRTEM simulations. The following parameters are specific to the STEM
# simulations.
# Next, we specify the simulation parameters related to the modelling of the
# STEM probe. See the documentation for the class
# :class:`embeam.stem.probe.ModelParams` for details on these parameters.
convergence_semiangle = 4.5 # In mrads.
defocal_offset_supersampling = 3
kwargs = {"lens_model_params": lens_model_params,
"gun_model_params": gun_model_params,
"convergence_semiangle": convergence_semiangle,
"defocal_offset_supersampling": defocal_offset_supersampling}
probe_model_params = embeam.stem.probe.ModelParams(**kwargs)
# Next, we specify the simulation parameters related to the probe scan pattern
# to be used in our STEM simulations, namely a rectangular grid-like
# pattern. See the documentation for the class
# :class:`prismatique.scan.rectangular.Params` for details on these parameters.
kwargs = {"step_size": (5, 5), # In Å.
"window": (0.25, 0.75, 0.25, 0.75),
"jitter": 0.25, # Dimensionless.
"rng_seed": 5555}
scan_config = prismatique.scan.rectangular.Params(**kwargs)
# Next, we specify the simulation parameters related to the modelling of the
# STEM system. See the documentation for the class
# :class:`prismatique.stem.system.ModelParams` for details on these parameters.
kwargs = {"sample_specification": sample_model_params,
"probe_model_params": probe_model_params,
"specimen_tilt": (0, 0), # In mrads.
"scan_config": scan_config}
stem_system_model_params = prismatique.stem.system.ModelParams(**kwargs)
# Next, we specify the simulation parameters related to the convergent beam
# electron diffraction patterns to be generated. See the documentation for the
# class :class:`prismatique.cbed.Params` for details on these parameters. Note
# that we also make use of classes from the ``empix`` library below. See the
# documentation of said library for details on these classes.
window_dims = (sample_supercell_reduced_xy_dims_in_pixels[0],
sample_supercell_reduced_xy_dims_in_pixels[1])
kwargs = {"center": (0, 0),
"window_dims": window_dims,
"pad_mode": "zeros",
"apply_symmetric_mask": True}
optional_cropping_params = empix.OptionalCroppingParams(**kwargs)
kwargs = {"block_dims": (2, 2),
"padding_const": 0,
"downsample_mode": "mean"} # Must be set to ``"mean"``.
optional_downsampling_params = empix.OptionalDownsamplingParams(**kwargs)
postprocessing_seq = (optional_cropping_params, optional_downsampling_params)
kwargs = {"postprocessing_seq": postprocessing_seq,
"avg_num_electrons_per_postprocessed_dp": 500000,
"apply_shot_noise": True,
"save_wavefunctions": True,
"save_final_intensity": True}
cbed_params = prismatique.cbed.Params(**kwargs)
# Next, we specify the base output parameters for the STEM simulations. Note
# that we are setting up two STEM simulations which simulate the same experiment
# but are implemented using different algorithms. See the documentation for the
# class :class:`prismatique.stem.output.base.Params` for details on these
# parameters.
output_dirname = path_to_data_dir + "/multislice_stem_sim_output"
kwargs = {"output_dirname": output_dirname,
"max_data_size": 5e9, # In bytes.
"cbed_params": cbed_params,
"radial_step_size_for_3d_stem": 1, # In mrads.
"radial_range_for_2d_stem": (0, convergence_semiangle), # In mrads.
"save_com": True,
"save_potential_slices": False}
base_stem_output_params_1 = prismatique.stem.output.base.Params(**kwargs)
output_dirname = path_to_data_dir + "/prism_stem_sim_output"
kwargs["output_dirname"] = output_dirname
base_stem_output_params_2 = prismatique.stem.output.base.Params(**kwargs)
# Next, we specify the output parameters that are applicable to the STEM
# simulation implemented using the multislice algorithm. See the documentation
# for the class :class:`prismatique.stem.output.multislice.Params` for details
# on these parameters.
kwargs = \
{"num_slices_per_output": 4,
"z_start_output": 0} # In Å.
output_params_specific_to_stem_multislice_alg = \
prismatique.stem.output.multislice.Params(**kwargs)
# Next, we specify the output parameters that are applicable to the STEM
# simulation implemented using the PRISM algorithm. See the documentation for
# the class :class:`prismatique.stem.output.prism.Params` for details on these
# parameters.
kwargs = \
{"enable_S_matrix_refocus": True, "save_S_matrices": False}
output_params_specific_to_stem_prism_alg = \
prismatique.stem.output.prism.Params(**kwargs)
# Next, we group the above output parameters into two sets: one set is the
# output parameter set for the STEM simulation implemented using the multislice
# algorithm, and the other set is that for the STEM simulation implemented using
# the PRISM algorithm. See the documentation for the class
# :class:`prismatique.stem.output.Params` for details on these parameters.
kwargs = {"base_params": base_stem_output_params_1,
"alg_specific_params": output_params_specific_to_stem_multislice_alg}
stem_output_params_1 = prismatique.stem.output.Params(**kwargs)
kwargs = {"base_params": base_stem_output_params_2,
"alg_specific_params": output_params_specific_to_stem_prism_alg}
stem_output_params_2 = prismatique.stem.output.Params(**kwargs)
# Next, we group all the above simulation parameters related to STEM simulations
# into two sets: one set is the complete STEM simulation parameter set for the
# STEM simulation implemented using the multislice algorithm, and the other set
# is that for the STEM simulation implemented using the PRISM algorithm. See the
# documentation for the class :class:`prismatique.stem.sim.Params` for details
# on these parameters.
kwargs = {"stem_system_model_params": stem_system_model_params,
"output_params": stem_output_params_1,
"worker_params": worker_params}
stem_sim_params_1 = prismatique.stem.sim.Params(**kwargs)
kwargs["output_params"] = stem_output_params_2
stem_sim_params_2 = prismatique.stem.sim.Params(**kwargs)
# The following parameters are specific to the HRTEM simulation.
# Next, we specify the simulation parameters related to the objective
# aperture. See the documentation for the class
# :class:`prismatique.aperture.Params` for details on these parameters.
kwargs = {"offset": (0, 0), # In mrads.
"window": (0, 10*convergence_semiangle)} # In mrads.
objective_aperture_params = prismatique.aperture.Params(**kwargs)
# Next, we specify the simulation parameters related to the beam tilt series in
# the HRTEM simulation. Here, we use the beam tilt series to model a single
# spatially incoherent HRTEM beam. See the documentation for the class
# :class:`prismatique.tilt.Params` for details on these parameters.
kwargs = {"offset": (0, 0), # In mrads.
"window": (0, 2*convergence_semiangle),
"spread": 3} # In mrads.
tilt_params = prismatique.tilt.Params(**kwargs)
# Next, we specify the simulation parameters related to the modelling of the
# HRTEM system. See the documentation for the class
# :class:`prismatique.hrtem.system.ModelParams` for details on these parameters.
kwargs = {"sample_specification": sample_model_params,
"gun_model_params": gun_model_params,
"lens_model_params": lens_model_params,
"tilt_params": tilt_params,
"objective_aperture_params": objective_aperture_params,
"defocal_offset_supersampling": defocal_offset_supersampling}
hrtem_system_model_params = prismatique.hrtem.system.ModelParams(**kwargs)
# Next, we specify the simulation parameters related to the HRTEM image
# wavefunctions and intensities to be generated. See the documentation for the
# class :class:`prismatique.hrtem.image.Params` for details on these
# parameters.
kwargs = {"center": None,
"window_dims": window_dims,
"pad_mode": "zeros",
"apply_symmetric_mask": True}
optional_cropping_params = empix.OptionalCroppingParams(**kwargs)
postprocessing_seq = (optional_cropping_params,)
kwargs = {"postprocessing_seq": postprocessing_seq,
"avg_num_electrons_per_postprocessed_image": 500000,
"apply_shot_noise": True,
"save_wavefunctions": True,
"save_final_intensity": True}
image_params = prismatique.hrtem.image.Params(**kwargs)
# Next, we specify the output parameters for the HRTEM simulation. See the
# documentation for the class :class:`prismatique.hrtem.output.Params` for
# details on these parameters.
output_dirname = path_to_data_dir + "/hrtem_sim_output"
kwargs = {"output_dirname": output_dirname,
"max_data_size": 5e9, # In bytes.
"image_params": image_params,
"save_potential_slices": False}
hrtem_output_params = prismatique.hrtem.output.Params(**kwargs)
# Next, we group all of the HRTEM simulation parameters together. See the
# documentation for the class :class:`prismatique.hrtem.sim.Params` for details
# on these parameters.
kwargs = {"hrtem_system_model_params": hrtem_system_model_params,
"output_params": hrtem_output_params,
"worker_params": worker_params}
hrtem_sim_params = prismatique.hrtem.sim.Params(**kwargs)
# Lastly, with all of the simulation parameters set, we can save them to JSON
# files for future use.
output_dirname = path_to_data_dir + "/sim_param_generator_output"
pathlib.Path(output_dirname).mkdir(parents=True, exist_ok=True)
filename = output_dirname + "/multislice_stem_sim_params.json"
stem_sim_params_1.dump(filename, overwrite=True)
filename = output_dirname + "/prism_stem_sim_params.json"
stem_sim_params_2.dump(filename, overwrite=True)
filename = output_dirname + "/hrtem_sim_params.json"
hrtem_sim_params.dump(filename, overwrite=True)
msg = ("Finished generating the simulation parameter sets for the various "
"simulated experiments involving the MoS2 sample: "
"all output files were saved in the directory {}"
".".format(output_dirname))
print(msg)
print()