def write_vasp_volumetric(self, filename="CHGCAR", normalize=False):
"""
Writes volumetric data into a VASP CHGCAR format
Args:
filename (str): Filename of the new file
normalize (bool): True if the data is to be normalized by the volume
"""
write_poscar(structure=self.atoms, filename=filename)
with open(filename, "a") as f:
f.write("\n")
f.write(" ".join(list(np.array(self.total_data.shape, dtype=str))))
f.write("\n")
_, n_y, n_z = self.total_data.shape
flattened_data = np.hstack([
self.total_data[:, i, j] for j in range(n_z)
for i in range(n_y)
])
if normalize:
flattened_data /= self.atoms.get_volume()
num_lines = int(len(flattened_data) / 5) * 5
reshaped_data = np.reshape(flattened_data[0:num_lines], (-1, 5))
np.savetxt(f, reshaped_data, fmt="%.12f")
if len(flattened_data) % 5 > 0:
np.savetxt(f, [flattened_data[num_lines:]], fmt="%.12f")
def test_vasp_sorter(self):
write_poscar(structure=self.structure,
filename=posixpath.join(self.file_location,
"POSCAR_test"))
test_atoms = read_atoms(
posixpath.join(self.file_location, "POSCAR_test"))
vasp_order = vasp_sorter(self.structure)
self.assertEqual(len(self.structure), len(test_atoms))
self.assertEqual(self.structure[vasp_order], test_atoms)
os.remove(posixpath.join(self.file_location, "POSCAR_test"))
def test_write_poscar(self):
write_poscar(structure=self.structure,
filename=posixpath.join(self.file_location,
"POSCAR_test"))
test_atoms = read_atoms(
posixpath.join(self.file_location, "POSCAR_test"))
self.assertEqual(self.structure.get_chemical_formula(),
test_atoms.get_chemical_formula())
struct = self.structure.copy()
struct.add_tag(selective_dynamics=[True, True, True])
write_poscar(structure=struct,
filename=posixpath.join(self.file_location,
"POSCAR_test"))
test_atoms = read_atoms(
posixpath.join(self.file_location, "POSCAR_test"))
truth_array = np.empty_like(struct.positions, dtype=bool)
truth_array[:] = [True, True, True]
self.assertTrue(
np.array_equal(np.array(test_atoms.selective_dynamics),
truth_array))
os.remove(posixpath.join(self.file_location, "POSCAR_test"))
def test_manip_contcar(self):
for f in self.file_list:
if "CONTCAR_Mg" in f:
struct = read_atoms(f)
Mg_indices = struct.select_index("Mg")
add_pos = np.zeros_like(struct.positions)
max_Mg = np.argmax(struct.positions[Mg_indices, 2])
init_z = struct.positions[max_Mg, 2]
add_pos[np.argsort(vasp_sorter(struct))[max_Mg], 2] += 5.0
manip_contcar(filename=f,
new_filename="manip_file",
add_pos=add_pos)
new_struct = read_atoms("manip_file")
Mg_indices = new_struct.select_index("Mg")
max_Mg = np.argmax(new_struct.positions[Mg_indices, 2])
final_z = new_struct.positions[max_Mg, 2]
self.assertEqual(round(final_z - init_z, 3), 5.0)
os.remove("manip_file")
break
positions = np.ones((3, 3))
positions[0] = [5., 5., 5.]
positions[1] = [5., 5.7, 5.7]
positions[2] = [5., -5.7, -5.7]
struct = Atoms(["O", "H", "H"],
positions=positions,
cell=10. * np.eye(3))
write_poscar(structure=struct, filename="simple_water")
add_pos = np.zeros_like(positions)
poscar_order = np.argsort(vasp_sorter(struct))
add_pos[poscar_order[struct.select_index("O")], 2] += 3
manip_contcar("simple_water", "simple_water_new", add_pos)
new_struct = read_atoms("simple_water_new")
self.assertEqual(new_struct.positions[new_struct.select_index("O"), 2],
8)
os.remove("simple_water")
os.remove("simple_water_new")