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()
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()