def test_slice(tmpdir_factory):
directory = tmpdir_factory.mktemp("proc")
filename = os.path.join(directory, "temp.h5")
box = (400, 400, 400)
centre = (200, 200, 200)
shape = {"type": "cube", "cube": {"length": 400}}
sample = parakeet.sample.new(filename, box, centre, shape)
parakeet.sample.add_single_molecule(sample, "4v5d")
# Create the system configuration
system_conf = multem.SystemConfiguration()
system_conf.precision = "float"
if multem.is_gpu_available():
system_conf.device = "device"
else:
system_conf.device = "host"
# Create the input multislice configuration
n_phonons = 50
input_multislice = create_input_multislice(n_phonons, False)
input_multislice.spec_atoms = sample.get_atoms().to_multem()
input_multislice.spec_lx = sample.containing_box[1][0]
input_multislice.spec_ly = sample.containing_box[1][1]
input_multislice.spec_lz = sample.containing_box[1][2]
input_multislice.spec_dz = 3
print("Standard")
output = multem.simulate(system_conf, input_multislice)
print("Subslicing")
# Create the input multislice configuration
n_slices = 4
# Slice the sample
subslices = list(
multem.slice_spec_atoms(
input_multislice.spec_atoms, input_multislice.spec_lz, n_slices
)
)
# Do the slices simulation
sliced_output = multem.simulate(system_conf, input_multislice, iter(subslices))
# Print the difference
a = np.abs(np.array(output.data[-1].psi_coh)) ** 2
b = np.abs(np.array(sliced_output.data[-1].psi_coh)) ** 2
diff = np.max(np.abs(a - b))
input_multislice.obj_lens_c_23 = 0.0
input_multislice.obj_lens_phi_23 = 0.0
input_multislice.obj_lens_inner_aper_ang = 0.0
input_multislice.obj_lens_outer_aper_ang = 0.0
# defocus spread function
dsf_sigma = multem.iehwgd_to_sigma(32)
input_multislice.obj_lens_dsf_sigma = dsf_sigma
input_multislice.obj_lens_dsf_npoints = 5
# zero defocus reference
input_multislice.obj_lens_zero_defocus_type = "First"
input_multislice.obj_lens_zero_defocus_plane = 0
# Do the simulation
output_multislice = multem.simulate(system_conf, input_multislice)
data = {
"dx":
output_multislice.dx,
"dy":
output_multislice.dy,
"x":
numpy.array(output_multislice.x, dtype=numpy.float64),
"y":
numpy.array(output_multislice.y, dtype=numpy.float64),
"thick":
numpy.array(output_multislice.thick, dtype=numpy.float64),
"data": [{
"m2psi_tot": numpy.array(d.m2psi_tot, numpy.float64)
} for d in output_multislice.data],
def compute_exit_wave(atom_data, pixel_size):
"""
Compute the exit wave
"""
# Get the dimensions
x_min = atom_data.data["x"].min()
x_max = atom_data.data["x"].max()
y_min = atom_data.data["y"].min()
y_max = atom_data.data["y"].max()
z_min = atom_data.data["z"].min()
z_max = atom_data.data["z"].max()
x_size = x_max - x_min
y_size = y_max - y_min
select = ((atom_data.data["x"] > x_min + x_size / 6)
& (atom_data.data["x"] < x_max - x_size / 6)
& (atom_data.data["y"] > y_min + y_size / 6)
& (atom_data.data["y"] < y_max - y_size / 6))
atom_data = parakeet.sample.AtomData(data=atom_data.data[select])
x_min = atom_data.data["x"].min()
x_max = atom_data.data["x"].max()
y_min = atom_data.data["y"].min()
y_max = atom_data.data["y"].max()
z_min = atom_data.data["z"].min()
z_max = atom_data.data["z"].max()
x_size = x_max - x_min
y_size = y_max - y_min
z_size = z_max - z_min
# Translate to centre
x_box_size = x_size
y_box_size = y_size
z_box_size = z_size
# Create the system configuration
system_conf = multem.SystemConfiguration()
system_conf.precision = "float"
system_conf.device = "device"
# Create the input multislice configuration
input_multislice = create_input_multislice()
# Compute the number of pixels
nx = next_power_2(int(x_box_size / pixel_size))
ny = next_power_2(int(y_box_size / pixel_size))
assert nx <= 4096
assert ny <= 4096
x_box_size = nx * pixel_size
y_box_size = ny * pixel_size
x_trans = (x_box_size - x_size) / 2.0 - x_min
y_trans = (y_box_size - y_size) / 2.0 - y_min
z_trans = (z_box_size - z_size) / 2.0 - z_min
atom_data.translate((x_trans, y_trans, z_trans))
# Create the specimen size
input_multislice.nx = nx
input_multislice.ny = ny
input_multislice.spec_lx = x_box_size
input_multislice.spec_ly = y_box_size
input_multislice.spec_lz = z_box_size
input_multislice.spec_dz = 5
# Set the specimen atoms
input_multislice.spec_atoms = atom_data.to_multem()
# Run the simulation
output_multislice = multem.simulate(system_conf, input_multislice)
# Get the image
physical_image = numpy.array(output_multislice.data[0].psi_coh).T
# Create the masker
masker = multem.Masker(input_multislice.nx, input_multislice.ny,
pixel_size)
# Create the size of the cuboid
masker.set_cuboid(
(
x_box_size / 2 - x_size / 2,
y_box_size / 2 - y_size / 2,
z_box_size / 2 - z_size / 2,
),
(x_size, y_size, z_size),
)
# Run the simulation
input_multislice.spec_atoms = multem.AtomList()
output_multislice = multem.simulate(system_conf, input_multislice, masker)
# Get the image
random_image = numpy.array(output_multislice.data[0].psi_coh).T
# Return the images
x0 = numpy.array(
(x_box_size / 2 - x_size / 2, y_box_size / 2 - y_size / 2))
x1 = numpy.array(
(x_box_size / 2 + x_size / 2, y_box_size / 2 + y_size / 2))
return physical_image, random_image, x0, x1
import multem input_multislice = multem.Input() config = multem.SystemConfiguration() # print(multem.is_gpu_available()) # if not multem.is_gpu_available(): config.device = "host" result = multem.simulate(config, input_multislice) # print(result)
rms3d = 0.085
(
input_multislice.spec_atoms,
input_multislice.spec_lx,
input_multislice.spec_ly,
input_multislice.spec_lz,
a,
b,
c,
input_multislice.spec_dz,
) = SrTiO3001_crystal(na, nb, nc, ncu, rms3d)
print("Standard")
start_time = time.time()
ewrs = multem.simulate(system_conf, input_multislice)
print("Time taken: ", time.time() - start_time)
print("Subslicing")
start_time = time.time()
# Create the input multislice configuration
n_slices = 4
subslices = iter(
list(
multem.slice_spec_atoms(input_multislice.spec_atoms,
input_multislice.spec_lz, n_slices)))
output_multislice = multem.simulate(system_conf, input_multislice,
subslices)
def __call__(self, index):
"""
Simulate a single frame
Args:
simulation (object): The simulation object
index (int): The frame number
Returns:
tuple: (angle, image)
"""
# Get the specimen atoms
logger.info(f"Simulating image {index+1}")
# Get the rotation angle
angle = self.scan.angles[index]
position = self.scan.positions[index]
# Add the beam drift
if (
self.microscope.beam.drift
and not self.microscope.beam.drift["type"] is None
):
if self.microscope.beam.drift["type"] == "random":
shiftx, shifty = numpy.random.normal(
0, self.microscope.beam.drift["magnitude"], size=2
)
logger.info("Adding drift of %f, %f " % (shiftx, shifty))
elif self.microscope.beam.drift["type"] == "sinusoidal":
shiftx = sin(angle * pi / 180) * self.microscope.beam.drift["magnitude"]
shifty = shiftx
logger.info("Adding drift of %f, %f " % (shiftx, shifty))
else:
raise RuntimeError("Unknown drift type")
else:
shiftx = 0
shifty = 0
# The field of view
nx = self.microscope.detector.nx
ny = self.microscope.detector.ny
pixel_size = self.microscope.detector.pixel_size
origin = numpy.array(self.microscope.detector.origin)
margin = self.simulation["margin"]
padding = self.simulation["padding"]
x_fov = nx * pixel_size
y_fov = ny * pixel_size
margin_offset = margin * pixel_size
# padding_offset = padding * pixel_size
offset = (padding + margin) * pixel_size
# Set the rotation angle
# input_multislice.spec_rot_theta = angle
# input_multislice.spec_rot_u0 = simulation.scan.axis
# Create the sample extractor
x0 = numpy.array((-margin_offset, position - margin_offset))
x1 = numpy.array((x_fov + margin_offset, position + y_fov + margin_offset))
x0 += origin
x1 += origin
# thickness = self.simulation["division_thickness"]
# extractor = parakeet.sample.AtomSliceExtractor(
# sample=self.sample,
# translation=position,
# rotation=angle,
# x0=x0,
# x1=x1,
# thickness=thickness,
# )
# Create the multem system configuration
system_conf = create_system_configuration(self.device)
# The Z centre
z_centre = self.sample.centre[2]
# Create the multem input multislice object
input_multislice = create_input_multislice(
self.microscope,
self.simulation["slice_thickness"],
self.simulation["margin"] + self.simulation["padding"],
"EWRS",
z_centre,
)
# Set the specimen size
input_multislice.spec_lx = x_fov + offset * 2
input_multislice.spec_ly = y_fov + offset * 2
input_multislice.spec_lz = self.sample.containing_box[1][2]
# Compute the B factor
if self.simulation["radiation_damage_model"]:
sigma_B = sqrt(
self.simulation["sensitivity_coefficient"]
* self.microscope.beam.electrons_per_angstrom
* (index + 1)
)
else:
sigma_B = 0
# Either slice or don't
# if True: # len(extractor) == 1:
# Set the atoms in the input after translating them for the offset
# zslice = extractor[0]
atoms = self.sample.get_atoms_in_fov(x0, x1)
logger.info("Simulating with %d atoms" % atoms.data.shape[0])
# logger.info(
# " Simulating z slice %f -> %f with %d atoms"
# % (zslice.x_min[2], zslice.x_max[2], zslice.atoms.data.shape[0])
# )
# Set atom sigma
atoms.data["sigma"] = sigma_B
if len(atoms.data) > 0:
coords = atoms.data[["x", "y", "z"]].to_numpy()
coords = (
Rotation.from_rotvec((0, angle * pi / 180, 0)).apply(
coords - self.sample.centre
)
+ self.sample.centre
- (shiftx, shifty + position, 0)
).astype("float32")
atoms.data["x"] = coords[:, 0]
atoms.data["y"] = coords[:, 1]
atoms.data["z"] = coords[:, 2]
# atoms.data = atoms.data.append(parakeet.sample.AtomData(atomic_number=[1,1], x=[750,750], y=[750,750], z=[0,4000],sigma=[0,0],occupancy=[1,1],charge=[0,0]).data)
input_multislice.spec_atoms = atoms.translate(
(offset - origin[0], offset - origin[1], 0)
).to_multem()
logger.info(" Got spec atoms")
if self.simulation["ice"] == True:
# Create the masker
masker = multem.Masker(input_multislice.nx, input_multislice.ny, pixel_size)
# Get the sample centre
shape = self.sample.shape
centre = self.sample.centre
centre = (
centre[0] + offset - shiftx - origin[0],
centre[1] + offset - shifty - position - origin[1],
centre[2],
)
# Set the shape
if shape["type"] == "cube":
length = shape["cube"]["length"]
masker.set_cuboid(
(
centre[0] - length / 2,
centre[1] - length / 2,
centre[2] - length / 2,
),
(length, length, length),
)
elif shape["type"] == "cuboid":
length_x = shape["cuboid"]["length_x"]
length_y = shape["cuboid"]["length_y"]
length_z = shape["cuboid"]["length_z"]
masker.set_cuboid(
(
centre[0] - length_x / 2,
centre[1] - length_y / 2,
centre[2] - length_z / 2,
),
(length_x, length_y, length_z),
)
elif shape["type"] == "cylinder":
radius = shape["cylinder"]["radius"]
if not isinstance(radius, Iterable):
radius = [radius]
length = shape["cylinder"]["length"]
offset_x = shape["cylinder"].get("offset_x", [0] * len(radius))
offset_z = shape["cylinder"].get("offset_z", [0] * len(radius))
axis = shape["cylinder"].get("axis", (0, 1, 0))
masker.set_cylinder(
(centre[0], centre[1] - length / 2, centre[2]),
axis,
length,
list(radius),
list(offset_x),
list(offset_z),
)
# Rotate
origin = centre
masker.set_rotation(origin, (0, angle * pi / 180.0, 0))
# Run the simulation
output_multislice = multem.simulate(system_conf, input_multislice, masker)
else:
# Run the simulation
logger.info("Simulating")
output_multislice = multem.simulate(system_conf, input_multislice)
# else:
# pass
# # Slice the specimen atoms
# def slice_generator(extractor):
# # Get the data from the data buffer and return
# def prepare(data_buffer):
# # Extract the data
# atoms = parakeet.sample.AtomData(
# data=pandas.concat([d.atoms.data for d in data_buffer])
# )
# z_min = min([d.x_min[2] for d in data_buffer])
# z_max = max([d.x_max[2] for d in data_buffer])
# assert z_min < z_max
# # Print some info
# logger.info(
# " Simulating z slice %f -> %f with %d atoms"
# % (z_min, z_max, atoms.data.shape[0])
# )
# # Cast the atoms
# atoms = atoms.translate((offset, offset, 0)).to_multem()
# # Return the Z-min, Z-max and atoms
# return (z_min, z_max, atoms)
# # Loop through the slices and gather atoms until we have more
# # than the maximum buffer size. There seems to be an overhead
# # to the simulation code so it's better to have as many atoms
# # as possible before calling. Doing this is much fast than
# # simulating with only a small number of atoms.
# max_buffer = 10_000_000
# data_buffer = []
# for zslice in extractor:
# data_buffer.append(zslice)
# if sum(d.atoms.data.shape[0] for d in data_buffer) > max_buffer:
# yield prepare(data_buffer)
# data_buffer = []
# # Simulate from the final buffer
# if len(data_buffer) > 0:
# yield prepare(data_buffer)
# data_buffer = []
# Run the simulation
# st = time.time()
# output_multislice = multem.simulate(
# system_conf, input_multislice, slice_generator(extractor)
# )
# logger.info(
# " Image %d simulated in %d seconds" % (index, time.time() - st)
# )
# Get the ideal image data
# Multem outputs data in column major format. In C++ and Python we
# generally deal with data in row major format so we must do a
# transpose here.
image = numpy.array(output_multislice.data[0].psi_coh).T
image = image[padding:-padding, padding:-padding]
# Print some info
psi_tot = numpy.abs(image) ** 2
logger.info(
"Ideal image min/max: %f/%f" % (numpy.min(psi_tot), numpy.max(psi_tot))
)
# Compute the image scaled with Poisson noise
return (index, angle, position, image, (shiftx, shifty))
def __call__(self, index):
"""
Simulate a single frame
Args:
simulation (object): The simulation object
index (int): The frame number
Returns:
tuple: (angle, image)
"""
# Get the rotation angle
angle = self.scan.angles[index]
position = self.scan.positions[index]
# The field of view
nx = self.microscope.detector.nx
ny = self.microscope.detector.ny
pixel_size = self.microscope.detector.pixel_size
margin = self.simulation["margin"]
x_fov = nx * pixel_size
y_fov = ny * pixel_size
offset = margin * pixel_size
# Get the specimen atoms
logger.info(f"Simulating image {index+1}")
# Set the rotation angle
# input_multislice.spec_rot_theta = angle
# input_multislice.spec_rot_u0 = simulation.scan.axis
# x0 = (-offset, -offset)
# x1 = (x_fov + offset, y_fov + offset)
# Create the multem system configuration
system_conf = create_system_configuration(self.device)
# Create the multem input multislice object
input_multislice = create_input_multislice(
self.microscope,
self.simulation["slice_thickness"],
self.simulation["margin"],
"EWRS",
)
# Set the specimen size
input_multislice.spec_lx = x_fov + offset * 2
input_multislice.spec_ly = y_fov + offset * 2
input_multislice.spec_lz = numpy.max(self.atoms.data["z"])
# Set the atoms in the input after translating them for the offset
input_multislice.spec_atoms = self.atoms.translate(
(offset, offset, 0)
).to_multem()
# Run the simulation
output_multislice = multem.simulate(system_conf, input_multislice)
# Get the ideal image data
# Multem outputs data in column major format. In C++ and Python we
# generally deal with data in row major format so we must do a
# transpose here.
image = numpy.array(output_multislice.data[0].psi_coh).T
# Print some info
psi_tot = numpy.abs(image) ** 2
logger.info(
"Ideal image min/max: %f/%f" % (numpy.min(psi_tot), numpy.max(psi_tot))
)
# Compute the image scaled with Poisson noise
return (index, angle, position, image, None)
