def test_wrong_number_of_sublattices(self):
sublattice = am.Sublattice(np.zeros((10, 2)), np.zeros((10, 10)))
sublattice_list = [
sublattice,
]
with pytest.raises(ValueError):
al.Dumbbell_Lattice(image=np.zeros((10, 10)),
sublattice_list=sublattice_list)
sublattice_list = [sublattice, sublattice, sublattice]
with pytest.raises(ValueError):
al.Dumbbell_Lattice(image=np.zeros((10, 10)),
sublattice_list=sublattice_list)
def test_dumbbell_distance(self):
positions0 = [[10, 2], [12, 4], [14, 5], [20, 30]]
positions1 = [[9, 2], [10, 4], [11, 5], [20, 10]]
sublattice0 = am.Sublattice(positions0, np.zeros((10, 10)))
sublattice1 = am.Sublattice(positions1, np.zeros((10, 10)))
sublattice_list = [sublattice0, sublattice1]
dumbbell_lattice = al.Dumbbell_Lattice(image=np.zeros((10, 10)),
sublattice_list=sublattice_list)
dumbbell_distance = dumbbell_lattice.dumbbell_distance
assert (dumbbell_distance == [1., 2., 3., 20.]).all()
def test_dumbbell_angle(self):
positions0 = [[10, 2], [9, 2], [5, 2]]
positions1 = [[9, 2], [10, 2], [5, 3]]
sublattice0 = am.Sublattice(positions0, np.zeros((10, 10)))
sublattice1 = am.Sublattice(positions1, np.zeros((10, 10)))
sublattice_list = [sublattice0, sublattice1]
dumbbell_lattice = al.Dumbbell_Lattice(image=np.zeros((10, 10)),
sublattice_list=sublattice_list)
dumbbell_angle = dumbbell_lattice.dumbbell_angle
print(dumbbell_angle)
assert (dumbbell_angle == [np.pi, 0., np.pi / 2]).all()
def test_dumbbell_x_y(self):
positions0 = [[10, 3], [11, 5]]
positions1 = [[9, 2], [10, 4]]
sublattice0 = am.Sublattice(positions0, np.zeros((10, 10)))
sublattice1 = am.Sublattice(positions1, np.zeros((10, 10)))
sublattice_list = [sublattice0, sublattice1]
dumbbell_lattice = al.Dumbbell_Lattice(image=np.zeros((10, 10)),
sublattice_list=sublattice_list)
dumbbell_x = dumbbell_lattice.dumbbell_x
assert (dumbbell_x == [9.5, 10.5]).all()
dumbbell_y = dumbbell_lattice.dumbbell_y
assert (dumbbell_y == [2.5, 4.5]).all()
def test_load_simple(self):
image_data = np.arange(10000).reshape(100, 100)
peaks = np.arange(20).reshape(10, 2)
sublattice = Sublattice(atom_position_list=peaks, image=image_data)
dumbbell_lattice = al.Dumbbell_Lattice(
image=image_data, sublattice_list=[sublattice, sublattice])
save_path = pjoin(self.tmpdir.name, "test_dumbbell_lattice_save.hdf5")
dumbbell_lattice.save(filename=save_path, overwrite=True)
dumbbell_lattice_load = load_atom_lattice_from_hdf5(
save_path, construct_zone_axes=False)
dl0_qualname = dumbbell_lattice.__class__.__qualname__
dl1_qualname = dumbbell_lattice_load.__class__.__qualname__
assert dl0_qualname == 'Dumbbell_Lattice'
assert dl0_qualname == dl1_qualname
def test_two_sublattices(self):
sublattice = am.Sublattice(np.zeros((10, 2)), np.zeros((10, 10)))
sublattice_list = [sublattice, sublattice]
dumbbell_lattice = al.Dumbbell_Lattice(image=np.zeros((10, 10)),
sublattice_list=sublattice_list)
hasattr(dumbbell_lattice, 'plot')
def test_empty(self):
dumbbell_lattice = al.Dumbbell_Lattice()
hasattr(dumbbell_lattice, 'plot')
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')
if 'atom_lattice_type' in h5f.attrs.keys():
atom_lattice_type = h5f.attrs['atom_lattice_type']
if 'Atom_Lattice' in atom_lattice_type:
atom_lattice = al.Atom_Lattice()
elif 'Dumbbell_Lattice' in atom_lattice_type:
atom_lattice = al.Dumbbell_Lattice()
else:
atom_lattice = al.Atom_Lattice()
else:
atom_lattice = al.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
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_data':
atom_lattice.image = h5f[group_name][:]
if group_name == 'image_data0':
atom_lattice.image = h5f[group_name][:]
if group_name == 'image_extra_data':
atom_lattice.image_extra = h5f[group_name][:]
if group_name == 'image_data1':
atom_lattice.image_extra = h5f[group_name][:]
if group_name == 'original_image_data':
atom_lattice.original_image = 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 get_dumbbell_heterostructure_dumbbell_lattice():
"""Get a dumbbell heterostructure dumbbell lattice.
Returns
-------
dumbbell_lattice : Atomap Dumbbell_Lattice
Example
-------
>>> dl = am.dummy_data.get_dumbbell_heterostructure_dumbbell_lattice()
>>> dl.plot()
"""
test_data = MakeTestData(500, 500)
x0, y0 = np.mgrid[10:500:24., 10:250:30.]
x1, y1 = np.mgrid[10:500:24., 18:250:30.]
x2, y2 = np.mgrid[22:500:24., 25:250:30.]
x3, y3 = np.mgrid[22:500:24., 33:250:30.]
x0, y0 = x0.flatten(), y0.flatten()
x1, y1 = x1.flatten(), y1.flatten()
x2, y2 = x2.flatten(), y2.flatten()
x3, y3 = x3.flatten(), y3.flatten()
x0 += np.random.random(size=x0.shape) * 0.5
y0 += np.random.random(size=y0.shape) * 0.5
x1 += np.random.random(size=x1.shape) * 0.5
y1 += np.random.random(size=y1.shape) * 0.5
x2 += np.random.random(size=x2.shape) * 0.5
y2 += np.random.random(size=y2.shape) * 0.5
x3 += np.random.random(size=x3.shape) * 0.5
y3 += np.random.random(size=y3.shape) * 0.5
test_data.add_atom_list(x0, y0, sigma_x=3, sigma_y=3, amplitude=40)
test_data.add_atom_list(x1, y1, sigma_x=3, sigma_y=3, amplitude=50)
test_data.add_atom_list(x2, y2, sigma_x=3, sigma_y=3, amplitude=40)
test_data.add_atom_list(x3, y3, sigma_x=3, sigma_y=3, amplitude=50)
x4, y4 = np.mgrid[10:500:24., 250:500:32.]
x5, y5 = np.mgrid[10:500:24., 259:500:32.]
x6, y6 = np.mgrid[22:500:24., 266:500:32.]
x7, y7 = np.mgrid[22:500:24., 275:500:32.]
x4, y4 = x4.flatten(), y4.flatten()
x5, y5 = x5.flatten(), y5.flatten()
x6, y6 = x6.flatten(), y6.flatten()
x7, y7 = x7.flatten(), y7.flatten()
x4 += np.random.random(size=x4.shape) * 0.5
y4 += np.random.random(size=y4.shape) * 0.5
x5 += np.random.random(size=x5.shape) * 0.5
y5 += np.random.random(size=y5.shape) * 0.5
x6 += np.random.random(size=x6.shape) * 0.5
y6 += np.random.random(size=y6.shape) * 0.5
x7 += np.random.random(size=x7.shape) * 0.5
y7 += np.random.random(size=y7.shape) * 0.5
test_data.add_atom_list(x4, y4, sigma_x=3, sigma_y=3, amplitude=70)
test_data.add_atom_list(x5, y5, sigma_x=3, sigma_y=3, amplitude=60)
test_data.add_atom_list(x6, y6, sigma_x=3, sigma_y=3, amplitude=70)
test_data.add_atom_list(x7, y7, sigma_x=3, sigma_y=3, amplitude=60)
test_data0 = MakeTestData(500, 500)
test_data0.add_atom_list(x0, y0, sigma_x=3, sigma_y=3, amplitude=40)
test_data0.add_atom_list(x2, y2, sigma_x=3, sigma_y=3, amplitude=40)
test_data0.add_atom_list(x4, y4, sigma_x=3, sigma_y=3, amplitude=70)
test_data0.add_atom_list(x6, y6, sigma_x=3, sigma_y=3, amplitude=70)
sublattice0 = test_data0.sublattice
sublattice0._plot_color = 'blue'
sublattice0.image = test_data.signal.data
sublattice0.original_image = test_data.signal.data
test_data1 = MakeTestData(500, 500)
test_data1.add_atom_list(x1, y1, sigma_x=3, sigma_y=3, amplitude=50)
test_data1.add_atom_list(x3, y3, sigma_x=3, sigma_y=3, amplitude=50)
test_data1.add_atom_list(x5, y5, sigma_x=3, sigma_y=3, amplitude=60)
test_data1.add_atom_list(x7, y7, sigma_x=3, sigma_y=3, amplitude=60)
sublattice1 = test_data1.sublattice
sublattice1.image = test_data.signal.data
sublattice1.original_image = test_data.signal.data
sublattice0.find_nearest_neighbors()
sublattice1.find_nearest_neighbors()
dumbbell_lattice = al.Dumbbell_Lattice(
image=test_data.signal.data,
sublattice_list=[sublattice0, sublattice1])
return dumbbell_lattice