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