def setup_method(self):
image_data = np.arange(10000).reshape(100, 100)
peaks0 = np.arange(20).reshape(10, 2)
peaks1 = np.arange(26).reshape(13, 2)
sublattice0 = Sublattice(atom_position_list=peaks0, image=image_data)
sublattice1 = Sublattice(atom_position_list=peaks1, image=image_data)
self.atom_lattice = Atom_Lattice()
self.atom_lattice.sublattice_list.extend([sublattice0, sublattice1])
self.atom_lattice.image0 = image_data
def test_save_load_atom_lattice_atom_values(self):
image_data = np.arange(10000).reshape(100, 100)
atom0_pos = np.random.random(size=(30, 2)) * 10
atom0_sigma_x = np.random.random(size=30)
atom0_sigma_y = np.random.random(size=30)
atom0_rot = np.random.random(size=30)
atom1_pos = np.random.random(size=(30, 2)) * 10
atom1_sigma_x = np.random.random(size=30)
atom1_sigma_y = np.random.random(size=30)
atom1_rot = np.random.random(size=30)
sublattice0 = Sublattice(atom_position_list=atom0_pos,
image=image_data)
sublattice1 = Sublattice(atom_position_list=atom1_pos,
image=image_data)
for i, atom in enumerate(sublattice0.atom_list):
atom.sigma_x = atom0_sigma_x[i]
atom.sigma_y = atom0_sigma_y[i]
atom.rotation = atom0_rot[i]
for i, atom in enumerate(sublattice1.atom_list):
atom.sigma_x = atom1_sigma_x[i]
atom.sigma_y = atom1_sigma_y[i]
atom.rotation = atom1_rot[i]
atom_lattice = Atom_Lattice()
atom_lattice.sublattice_list.extend([sublattice0, sublattice1])
atom_lattice.image0 = image_data
tmpdir = tempfile.TemporaryDirectory()
save_path = os.path.join(tmpdir.name, "atomic_lattice.hdf5")
atom_lattice.save(filename=save_path, overwrite=True)
atom_lattice_load = load_atom_lattice_from_hdf5(
save_path, construct_zone_axes=False)
sl0 = atom_lattice_load.sublattice_list[0]
sl1 = atom_lattice_load.sublattice_list[1]
assert (sl0.x_position == atom0_pos[:, 0]).all()
assert (sl0.y_position == atom0_pos[:, 1]).all()
assert (sl1.x_position == atom1_pos[:, 0]).all()
assert (sl1.y_position == atom1_pos[:, 1]).all()
assert (sl0.sigma_x == atom0_sigma_x).all()
assert (sl0.sigma_y == atom0_sigma_y).all()
assert (sl1.sigma_x == atom1_sigma_x).all()
assert (sl1.sigma_y == atom1_sigma_y).all()
assert (sl0.rotation == atom0_rot).all()
assert (sl1.rotation == atom1_rot).all()
def get_simple_atom_lattice_two_sublattices(image_noise=False):
"""Returns a simple atom_lattice with two sublattices.
Parameters
----------
image_noise : bool, default False
Returns
-------
simple_atom_lattice : Atom_Lattice object
Examples
--------
>>> import temul.external.atomap_devel_012.api as am
>>> al = am.dummy_data.get_simple_atom_lattice_two_sublattices()
>>> al.plot()
"""
test_data_1 = MakeTestData(300, 300)
test_data_2 = MakeTestData(300, 300)
test_data_3 = MakeTestData(300, 300)
if image_noise:
test_data_3.add_image_noise(mu=0, sigma=0.02)
_make_atom_lists_two_sublattices(test_data_1, test_data_2)
_make_atom_lists_two_sublattices(test_data_3)
sublattice_1 = test_data_1.sublattice
sublattice_2 = test_data_2.sublattice
sublattice_2._plot_color = 'b'
image = test_data_3.signal.data
atom_lattice = Atom_Lattice(image=image,
name='Simple Atom Lattice',
sublattice_list=[sublattice_1, sublattice_2])
return (atom_lattice)
def setup_method(self):
pos0 = np.array([[5, 10], [10, 15]])
pos1 = np.array([[20, 25], [30, 35]])
sublattice0 = Sublattice(pos0, np.zeros((40, 40)))
sublattice1 = Sublattice(pos1, np.zeros((40, 40)))
self.atom_lattice = Atom_Lattice(
np.zeros((40, 40)), sublattice_list=[sublattice0, sublattice1])
self.x_pos = np.concatenate((pos0[:, 0], pos1[:, 0]))
self.y_pos = np.concatenate((pos0[:, 1], pos1[:, 1]))
def get_polarisation_dummy_dataset(image_noise=False):
"""Get an Atom Lattice with two sublattices resembling a perovskite film.
Similar to a perovskite oxide thin film, where the B cations
are shifted in the film.
Parameters
----------
image_noise : bool, default False
Returns
-------
simple_atom_lattice : Atom_Lattice object
Examples
--------
>>> import numpy as np
>>> from temul.dummy_data import get_polarisation_dummy_dataset
>>> atom_lattice = get_polarisation_dummy_dataset()
>>> atom_lattice.plot()
>>> sublatticeA = atom_lattice.sublattice_list[0]
>>> sublatticeB = atom_lattice.sublattice_list[1]
>>> sublatticeA.construct_zone_axes()
>>> za0, za1 = sublatticeA.zones_axis_average_distances[0:2]
>>> s_p = sublatticeA.get_polarization_from_second_sublattice(
... za0, za1, sublatticeB, color='blue')
>>> s_p.plot()
>>> vector_list = s_p.metadata.vector_list
>>> x, y = [i[0] for i in vector_list], [i[1] for i in vector_list]
>>> u, v = [i[2] for i in vector_list], [i[3] for i in vector_list]
>>> u, v = -np.asarray(u), -np.asarray(v)
You can they use the plot_polarisation_vectors function to visualise:
>>> from temul.polarisation import plot_polarisation_vectors
>>> ax = plot_polarisation_vectors(x, y, u, v, image=sublatticeA.image,
... unit_vector=False, plot_style="vector",
... overlay=True, color='yellow',
... degrees=False, save=None, monitor_dpi=50)
"""
test_data0 = _make_rigid_sublattice(image_noise=image_noise)
test_data1 = _make_polarised_sublattice(image_noise=image_noise)
image = test_data0.signal.data + test_data1.signal.data
sublattice0 = test_data0.sublattice
sublattice1 = test_data1.sublattice
sublattice0.image = image
sublattice1.image = image
sublattice0.original_image = image
sublattice1.original_image = image
sublattice1._plot_color = 'b'
atom_lattice = Atom_Lattice(image=image,
name='Test Polarisation Dataset',
sublattice_list=[sublattice0, sublattice1])
return atom_lattice
def test_create_atom_plane_object(self):
atom_list = [
Atom_Position(1, 2),
Atom_Position(2, 4),
]
zone_vector = (1, 2)
atom_list[0]._start_atom = [zone_vector]
atom_list[1]._end_atom = [zone_vector]
atom_lattice = Atom_Lattice()
Atom_Plane(atom_list, zone_vector, atom_lattice)
def test_scaling(self):
sublattice = Sublattice([
[10, 10],
],
np.ones((20, 20)),
pixel_size=0.2)
atom_lattice = Atom_Lattice(np.ones((100, 100)),
sublattice_list=[sublattice])
signal = atom_lattice.signal
assert signal.axes_manager.signal_axes[0].scale == 0.2
assert signal.axes_manager.signal_axes[1].scale == 0.2
def setup_method(self):
atom_lattice = Atom_Lattice()
atom_plane_list = []
for i in range(2):
atom_list = [
Atom_Position(1, 2),
Atom_Position(2, 4),
]
zone_vector = (1, 2)
atom_list[0]._start_atom = [zone_vector]
atom_list[1]._end_atom = [zone_vector]
atom_plane = Atom_Plane(atom_list, (1, 2), atom_lattice)
atom_plane_list.append(atom_plane)
self.atom_plane_list = atom_plane_list
def setup_method(self):
x, y = [1, 2], [2, 4]
sX, sY, r = [3.1, 1.2], [2.2, 1.1], [0.5, 0.4]
A_g = [10.2, 5.2]
atom_list = [
Atom_Position(x[0], y[0], sX[0], sY[0], r[0]),
Atom_Position(x[1], y[1], sX[1], sY[1], r[1]),
]
atom_list[0].amplitude_gaussian = A_g[0]
atom_list[1].amplitude_gaussian = A_g[1]
zone_vector = (1, 2)
atom_list[0]._start_atom = [zone_vector]
atom_list[1]._end_atom = [zone_vector]
self.atom_plane = Atom_Plane(atom_list, zone_vector, Atom_Lattice())
self.x, self.y, self.sX, self.sY, self.r = x, y, sX, sY, r
self.A_g = A_g
def get_polarisation_dummy_dataset_bora(image_noise=False):
test_data0 = _make_rigid_sublattice_bora(image_noise=image_noise)
test_data1 = _make_polarised_sublattice_bora(image_noise=image_noise)
image = test_data0.signal.data + test_data1.signal.data
sublattice0 = test_data0.sublattice
sublattice1 = test_data1.sublattice
sublattice0.image = image
sublattice1.image = image
sublattice0.original_image = image
sublattice1.original_image = image
sublattice1._plot_color = 'b'
atom_lattice = Atom_Lattice(image=image,
name='Test Polarisation Dataset',
sublattice_list=[sublattice0, sublattice1])
# atom_lattice.plot()
return atom_lattice
def _get_fantasite_atom_lattice():
test_data1 = MakeTestData(500, 500)
test_data1 = _add_fantasite_sublattice_A(test_data1)
test_data2 = MakeTestData(500, 500)
test_data2 = _add_fantasite_sublattice_B(test_data2)
test_data3 = MakeTestData(500, 500)
test_data3 = _add_fantasite_sublattice_A(test_data3)
test_data3 = _add_fantasite_sublattice_B(test_data3)
test_data3.add_image_noise(mu=0, sigma=0.01, random_seed=0)
sublattice_1 = test_data1.sublattice
sublattice_2 = test_data2.sublattice
sublattice_1._plot_color = 'b'
image = test_data3.signal.data
atom_lattice = Atom_Lattice(image=image,
name='Fantasite Atom Lattice',
sublattice_list=[sublattice_1, sublattice_2])
return (atom_lattice)
def get_polarization_film_atom_lattice(image_noise=False):
"""Get an Atom Lattice with two sublattices resembling a perovskite film.
Similar to a perovskite oxide thin film, where the B cations
are shifted in the film.
Parameters
----------
image_noise : bool, default False
Returns
-------
simple_atom_lattice : Atom_Lattice object
Examples
--------
>>> import temul.external.atomap_devel_012.api as am
>>> atom_lattice = am.dummy_data.get_polarization_film_atom_lattice()
>>> atom_lattice.plot()
With image noise
>>> atom_lattice = am.dummy_data.get_polarization_film_atom_lattice(
... image_noise=True)
>>> atom_lattice.plot()
"""
test_data0 = _make_polarization_film_A(image_noise=image_noise)
test_data1 = _make_polarization_film_B(image_noise=image_noise)
image = test_data0.signal.data + test_data1.signal.data
sublattice0 = test_data0.sublattice
sublattice1 = test_data1.sublattice
sublattice0.image = image
sublattice1.image = image
sublattice0.original_image = image
sublattice1.original_image = image
sublattice1._plot_color = 'b'
atom_lattice = Atom_Lattice(image=image,
name='Perovskite film',
sublattice_list=[sublattice0, sublattice1])
return atom_lattice
def test_manual_processing(self):
s_adf_filename = os.path.join(my_path, "datasets",
"test_ADF_cropped.hdf5")
s = load(s_adf_filename)
s.change_dtype('float32')
atom_positions = afr.get_atom_positions(
signal=s,
separation=17,
threshold_rel=0.02,
)
sublattice = Sublattice(atom_position_list=atom_positions,
image=s.data)
sublattice.find_nearest_neighbors()
sublattice.refine_atom_positions_using_center_of_mass(
sublattice.image, percent_to_nn=0.4)
sublattice.refine_atom_positions_using_2d_gaussian(sublattice.image,
percent_to_nn=0.4)
Atom_Lattice(image=s.data, sublattice_list=[sublattice])
sublattice.construct_zone_axes()
def load_atom_lattice_from_hdf5(filename, construct_zone_axes=True):
"""
Load an Atomap HDF5-file, restoring a saved Atom_Lattice.
Parameters
----------
filename : string
Filename of the HDF5-file.
construct_zone_axes : bool
If True, find relations between atomic positions by
constructing atomic planes. Default True.
Returns
-------
Atomap Atom_Lattice object
"""
h5f = h5py.File(filename, 'r')
atom_lattice = Atom_Lattice()
sublattice_list = []
sublattice_index_list = []
for group_name in h5f:
if ('atom_lattice' in group_name) or ('sublattice' in group_name):
sublattice_set = h5f[group_name]
modified_image_data = sublattice_set['modified_image_data'][:]
original_image_data = sublattice_set['original_image_data'][:]
atom_position_array = sublattice_set['atom_positions'][:]
if 'sublattice_index' in sublattice_set.attrs.keys():
sublattice_index_list.append(
sublattice_set.attrs['sublattice_index'])
sublattice = Sublattice(atom_position_array, modified_image_data)
sublattice.original_image = original_image_data
if 'sigma_x' in sublattice_set.keys():
sigma_x_array = sublattice_set['sigma_x'][:]
for atom, sigma_x in zip(sublattice.atom_list, sigma_x_array):
atom.sigma_x = sigma_x
if 'sigma_y' in sublattice_set.keys():
sigma_y_array = sublattice_set['sigma_y'][:]
for atom, sigma_y in zip(sublattice.atom_list, sigma_y_array):
atom.sigma_y = sigma_y
if 'rotation' in sublattice_set.keys():
rotation_array = sublattice_set['rotation'][:]
for atom, rotation in zip(sublattice.atom_list,
rotation_array):
atom.rotation = rotation
####
if 'amplitude_max_intensity' in sublattice_set.keys():
amplitude_max_intensity_array = sublattice_set[
'amplitude_max_intensity'][:]
for atom, amplitude_max_intensity in zip(
sublattice.atom_list, amplitude_max_intensity_array):
atom.amplitude_max_intensity = amplitude_max_intensity
if 'amplitude_mean_intensity' in sublattice_set.keys():
amplitude_mean_intensity_array = sublattice_set[
'amplitude_mean_intensity'][:]
for atom, amplitude_mean_intensity in zip(
sublattice.atom_list, amplitude_mean_intensity_array):
atom.amplitude_mean_intensity = amplitude_mean_intensity
if 'amplitude_min_intensity' in sublattice_set.keys():
amplitude_min_intensity_array = sublattice_set[
'amplitude_min_intensity'][:]
for atom, amplitude_min_intensity in zip(
sublattice.atom_list, amplitude_min_intensity_array):
atom.amplitude_min_intensity = amplitude_min_intensity
if 'amplitude_total_intensity' in sublattice_set.keys():
amplitude_total_intensity_array = sublattice_set[
'amplitude_total_intensity'][:]
for atom, amplitude_total_intensity in zip(
sublattice.atom_list, amplitude_total_intensity_array):
atom.amplitude_total_intensity = amplitude_total_intensity
if 'elements' in sublattice_set.keys():
elements_array = sublattice_set['elements'][:]
for atom, elements in zip(sublattice.atom_list,
elements_array):
atom.elements = elements
if 'z_height' in sublattice_set.keys():
z_height_array = sublattice_set['z_height'][:]
# z_height_array_2 = [] # first loop needed because i can't eval() the z_height itself, don't really know why
# for i in range(0, len(z_height_array)):
# z_h = ast.literal_eval(z_height_array[i])
# z_height_array_2.append(z_h)
for atom, z_height in zip(sublattice.atom_list,
z_height_array):
atom.z_height = z_height
####
sublattice.pixel_size = sublattice_set.attrs['pixel_size']
if 'tag' in sublattice_set.attrs.keys():
sublattice.name = sublattice_set.attrs['tag']
elif 'name' in sublattice_set.attrs.keys():
sublattice.name = sublattice_set.attrs['name']
else:
sublattice.name = ''
if type(sublattice.name) == bytes:
sublattice.name = sublattice.name.decode()
sublattice._plot_color = sublattice_set.attrs['plot_color']
if type(sublattice._plot_color) == bytes:
sublattice._plot_color = sublattice._plot_color.decode()
if 'pixel_separation' in sublattice_set.attrs.keys():
sublattice._pixel_separation = sublattice_set.attrs[
'pixel_separation']
else:
sublattice._pixel_separation = 0.0
if construct_zone_axes:
construct_zone_axes_from_sublattice(sublattice)
if 'zone_axis_names_byte' in sublattice_set.keys():
zone_axis_list_byte = sublattice_set.attrs[
'zone_axis_names_byte']
zone_axis_list = []
for zone_axis_name_byte in zone_axis_list_byte:
zone_axis_list.append(zone_axis_name_byte.decode())
sublattice.zones_axis_average_distances_names = zone_axis_list
sublattice_list.append(sublattice)
if group_name == 'image_data0':
atom_lattice.image0 = h5f[group_name][:]
atom_lattice.image = atom_lattice.image0
if group_name == 'image_data1':
atom_lattice.image1 = h5f[group_name][:]
sorted_sublattice_list = []
if not sublattice_index_list: # Support for older hdf5 files
sublattice_index_list = range(len(sublattice_list))
for sublattice_index in sublattice_index_list:
sorted_sublattice_list.append(sublattice_list[sublattice_index])
atom_lattice.sublattice_list.extend(sorted_sublattice_list)
if 'name' in h5f.attrs.keys():
atom_lattice.name = h5f.attrs['name']
elif 'path_name' in h5f.attrs.keys():
atom_lattice.name = h5f.attrs['path_name']
if 'pixel_separation' in h5f.attrs.keys():
atom_lattice._pixel_separation = h5f.attrs['pixel_separation']
else:
atom_lattice._pixel_separation = 0.8 / sublattice.pixel_size
if type(atom_lattice.name) == bytes:
atom_lattice.name = atom_lattice.name.decode()
h5f.close()
return (atom_lattice)
def atom_lattice(self):
sublattice = self.sublattice
atom_lattice = Atom_Lattice(image=sublattice.image,
sublattice_list=[sublattice])
return atom_lattice
def make_atom_lattice_from_image(s_image0,
process_parameter=None,
pixel_separation=None,
s_image1=None,
debug_plot=False):
if s_image0.data.dtype == 'float16':
raise ValueError(
"s_image0 has the dtype float16, which is not supported. "
"Convert it to something else, for example using "
"s_image0.change_dtype('float64')")
image0_filename = _get_signal_name(s_image0)
name = image0_filename
s_image0 = s_image0.deepcopy()
s_image0_modified = run_image_filtering(s_image0)
if process_parameter is None:
process_parameter = pp.GenericStructure()
image0_scale = s_image0.axes_manager[0].scale
if pixel_separation is None:
if process_parameter.peak_separation is None:
raise ValueError("pixel_separation is not set.\
Either set it in the process_parameter.peak_separation\
or pixel_separation parameter")
else:
pixel_separation = process_parameter.peak_separation / image0_scale
initial_atom_position_list = get_atom_positions(
s_image0_modified, separation=pixel_separation)
if s_image1 is not None:
if s_image1.data.dtype == 'float16':
raise ValueError(
"s_image1 has the dtype float16, which is not supported. "
"Convert it to something else, for example using "
"s_image1.change_dtype('float64')")
s_image1 = s_image1.deepcopy()
s_image1.data = 1. / s_image1.data
image1_data = s_image1.data
#################################
image0_data = s_image0.data
image0_data_modified = s_image0_modified.data
atom_lattice = Atom_Lattice(name=name)
atom_lattice._original_filename = image0_filename
atom_lattice.image0 = image0_data
if s_image1 is not None:
atom_lattice.image1 = image1_data
atom_lattice._pixel_separation = pixel_separation
for sublattice_index in range(process_parameter.number_of_sublattices):
sublattice_para = process_parameter.get_sublattice_from_order(
sublattice_index)
if sublattice_para.image_type == 0:
s_image = s_image0
image_data = image0_data
image_data_modified = image0_data_modified
if sublattice_para.image_type == 1:
if s_image1 is not None:
s_image = s_image1
image_data = image1_data
image_data_modified = image1_data
else:
break
if sublattice_para.sublattice_order == 0:
sublattice = Sublattice(initial_atom_position_list,
image_data_modified)
else:
temp_sublattice = atom_lattice.get_sublattice(
sublattice_para.sublattice_position_sublattice)
temp_zone_vector_index = temp_sublattice.get_zone_vector_index(
sublattice_para.sublattice_position_zoneaxis)
zone_vector = temp_sublattice.zones_axis_average_distances[
temp_zone_vector_index]
atom_list = temp_sublattice.find_missing_atoms_from_zone_vector(
zone_vector)
sublattice = Sublattice(atom_list, image_data)
zone_axis_para_list = False
if hasattr(sublattice_para, 'zone_axis_list'):
zone_axis_para_list = sublattice_para.zone_axis_list
sublattice._plot_color = sublattice_para.color
sublattice.name = sublattice_para.name
sublattice.pixel_size = s_image.axes_manager[0].scale
sublattice._pixel_separation = pixel_separation
sublattice.original_image = image_data
atom_lattice.sublattice_list.append(sublattice)
if debug_plot:
sublattice.plot_atom_list_on_image_data(figname=sublattice.name +
"_initial_position.jpg")
for atom in sublattice.atom_list:
atom.sigma_x = sublattice._pixel_separation / 10.
atom.sigma_y = sublattice._pixel_separation / 10.
if not (sublattice_para.sublattice_order == 0):
construct_zone_axes_from_sublattice(
sublattice, zone_axis_para_list=zone_axis_para_list)
atom_subtract_config = sublattice_para.atom_subtract_config
image_data = sublattice.image
for atom_subtract_para in atom_subtract_config:
temp_sublattice = atom_lattice.get_sublattice(
atom_subtract_para['sublattice'])
neighbor_distance = atom_subtract_para['neighbor_distance']
image_data = remove_atoms_from_image_using_2d_gaussian(
image_data,
temp_sublattice,
percent_to_nn=neighbor_distance)
sublattice.image = image_data
sublattice.original_image = image_data
refinement_config = sublattice_para.refinement_config
refinement_neighbor_distance = refinement_config['neighbor_distance']
refinement_steps = refinement_config['config']
for refinement_step in refinement_steps:
if refinement_step[0] == 'image_data':
refinement_step[0] = sublattice.original_image
elif refinement_step[0] == 'image_data_modified':
refinement_step[0] = sublattice.image
else:
refinement_step[0] = sublattice.original_image
refine_sublattice(sublattice, refinement_steps,
refinement_neighbor_distance)
if sublattice_para.sublattice_order == 0:
sublattice.construct_zone_axes(
zone_axis_para_list=zone_axis_para_list)
return (atom_lattice)
def test_create_empty_atom_lattice_object(self):
Atom_Lattice()
class TestAtomLatticeInputOutput:
def setup_method(self):
image_data = np.arange(10000).reshape(100, 100)
peaks0 = np.arange(20).reshape(10, 2)
peaks1 = np.arange(26).reshape(13, 2)
sublattice0 = Sublattice(atom_position_list=peaks0, image=image_data)
sublattice1 = Sublattice(atom_position_list=peaks1, image=image_data)
self.atom_lattice = Atom_Lattice()
self.atom_lattice.sublattice_list.extend([sublattice0, sublattice1])
self.atom_lattice.image0 = image_data
def test_save_load_atom_lattice_simple(self):
tmpdir = tempfile.TemporaryDirectory()
save_path = os.path.join(tmpdir.name, "test_atomic_lattice_save.hdf5")
self.atom_lattice.save(filename=save_path, overwrite=True)
load_atom_lattice_from_hdf5(save_path, construct_zone_axes=False)
tmpdir.cleanup()
def test_save_load_atom_lattice_check_metadata_values(self):
sublattice0 = self.atom_lattice.sublattice_list[0]
sublattice1 = self.atom_lattice.sublattice_list[1]
sublattice0.name = "test 0"
sublattice1.name = "test 1"
sublattice0._plot_color = "blue"
sublattice1._plot_color = "green"
assert len(sublattice0.atom_list) == 10
assert len(sublattice1.atom_list) == 13
tmpdir = tempfile.TemporaryDirectory()
save_path = os.path.join(tmpdir.name, "test_atomic_lattice_save.hdf5")
self.atom_lattice.save(filename=save_path, overwrite=True)
atom_lattice_load = load_atom_lattice_from_hdf5(
save_path, construct_zone_axes=False)
sl0 = atom_lattice_load.sublattice_list[0]
sl1 = atom_lattice_load.sublattice_list[1]
assert len(sl0.atom_list) == 10
assert len(sl1.atom_list) == 13
assert sl0.name == "test 0"
assert sl1.name == "test 1"
assert sl0._plot_color == "blue"
assert sl1._plot_color == "green"
tmpdir.cleanup()
def test_save_load_atom_lattice_atom_values(self):
image_data = np.arange(10000).reshape(100, 100)
atom0_pos = np.random.random(size=(30, 2)) * 10
atom0_sigma_x = np.random.random(size=30)
atom0_sigma_y = np.random.random(size=30)
atom0_rot = np.random.random(size=30)
atom1_pos = np.random.random(size=(30, 2)) * 10
atom1_sigma_x = np.random.random(size=30)
atom1_sigma_y = np.random.random(size=30)
atom1_rot = np.random.random(size=30)
sublattice0 = Sublattice(atom_position_list=atom0_pos,
image=image_data)
sublattice1 = Sublattice(atom_position_list=atom1_pos,
image=image_data)
for i, atom in enumerate(sublattice0.atom_list):
atom.sigma_x = atom0_sigma_x[i]
atom.sigma_y = atom0_sigma_y[i]
atom.rotation = atom0_rot[i]
for i, atom in enumerate(sublattice1.atom_list):
atom.sigma_x = atom1_sigma_x[i]
atom.sigma_y = atom1_sigma_y[i]
atom.rotation = atom1_rot[i]
atom_lattice = Atom_Lattice()
atom_lattice.sublattice_list.extend([sublattice0, sublattice1])
atom_lattice.image0 = image_data
tmpdir = tempfile.TemporaryDirectory()
save_path = os.path.join(tmpdir.name, "atomic_lattice.hdf5")
atom_lattice.save(filename=save_path, overwrite=True)
atom_lattice_load = load_atom_lattice_from_hdf5(
save_path, construct_zone_axes=False)
sl0 = atom_lattice_load.sublattice_list[0]
sl1 = atom_lattice_load.sublattice_list[1]
assert (sl0.x_position == atom0_pos[:, 0]).all()
assert (sl0.y_position == atom0_pos[:, 1]).all()
assert (sl1.x_position == atom1_pos[:, 0]).all()
assert (sl1.y_position == atom1_pos[:, 1]).all()
assert (sl0.sigma_x == atom0_sigma_x).all()
assert (sl0.sigma_y == atom0_sigma_y).all()
assert (sl1.sigma_x == atom1_sigma_x).all()
assert (sl1.sigma_y == atom1_sigma_y).all()
assert (sl0.rotation == atom0_rot).all()
assert (sl1.rotation == atom1_rot).all()
def test_save_atom_lattice_already_exist(self):
tmpdir = tempfile.TemporaryDirectory()
save_path = os.path.join(tmpdir.name, "test_atomic_lattice_io.hdf5")
self.atom_lattice.save(filename=save_path, overwrite=True)
with pytest.raises(FileExistsError):
self.atom_lattice.save(filename=save_path)
tmpdir.cleanup()
def test_create_atom_lattice_object(self):
atom_lattice = Atom_Lattice()
atom_lattice.sublattice_list.append(self.sublattice)
def test_get_sublattice_atom_list_on_image(self):
atom_lattice = Atom_Lattice()
atom_lattice.image0 = self.sublattice.image
atom_lattice.sublattice_list.append(self.sublattice)
atom_lattice.get_sublattice_atom_list_on_image()
def test_simple(self):
atom_lattice = Atom_Lattice(np.ones((100, 100)))
signal = atom_lattice.signal
assert_array_equal(atom_lattice.image, signal.data)
def test_no_image(self):
atom_lattice = Atom_Lattice()
with pytest.raises(ValueError):
atom_lattice.signal