def load_trajectories_from_gfs(runs, db_file):
fs_id = []
fs = []
for run in runs:
if "INCAR" in run.keys():
# for backwards compatibility with older version of mpmorph
fs_id.append(run["ionic_steps_fs_id"])
fs.append('previous_runs_gfs')
elif "input" in run.keys():
fs_id.append(
run["calcs_reversed"][0]["output"]["ionic_steps_fs_id"])
fs.append('structures_fs')
for i, v in enumerate(fs_id):
mmdb = VaspMDCalcDb.from_db_file(db_file, admin=True)
ionic_steps_dict = load_ionic_steps(v, mmdb.db, fs[i])
if i == 0:
trajectory = Trajectory.from_structures([
Structure.from_dict(i['structure']) for i in ionic_steps_dict
])
else:
trajectory.extend(
Trajectory.from_structures([
Structure.from_dict(i['structure'])
for i in ionic_steps_dict
]))
return trajectory
def test_slice(self):
sliced_traj = self.traj[2:99:3]
sliced_traj_from_structs = Trajectory.from_structures(
self.structures[2:99:3])
if len(sliced_traj) == len(sliced_traj_from_structs):
self.assertTrue(
all([
sliced_traj[i] == sliced_traj_from_structs[i]
for i in range(len(sliced_traj))
]))
else:
self.assertTrue(False)
sliced_traj = self.traj[:-4:2]
sliced_traj_from_structs = Trajectory.from_structures(
self.structures[:-4:2])
if len(sliced_traj) == len(sliced_traj_from_structs):
self.assertTrue(
all([
sliced_traj[i] == sliced_traj_from_structs[i]
for i in range(len(sliced_traj))
]))
else:
self.assertTrue(False)
def calcs_reversed_to_trajectory(calcs_reversed: List[dict]):
"""
Converts data from calc_reversed to pymatgen Trajectory objects
that contain structure, energy, force and stress data for each
ionic step.
Args:
calcs_reversed: List of dictionaries in calcs_reversed entry
of a task document.
"""
trajectories = []
for calculation in calcs_reversed:
for step in calculation["output"]["ionic_steps"]:
structures = []
frame_props = defaultdict(list) # type: dict
structures.append(Structure.from_dict(step["structure"]))
frame_props["e_fr_energy"].append(step["e_fr_energy"])
frame_props["e_wo_entrp"].append(step["e_wo_entrp"])
frame_props["e_0_energy"].append(step["e_wo_entrp"])
frame_props["forces"].append(step["forces"])
frame_props["stresses"].append(step["stress"])
traj = Trajectory.from_structures(structures,
frame_properties=frame_props,
time_step=None).as_dict()
trajectories.append(traj)
return trajectories
def get_pymatgen_trajectory(self, include_shell=False):
table = self.get_structure_table(input=True,
include_shell=include_shell)
lattices, constant_lattice = self.get_md_cell()
frames = [x[table.index] for x in self.get_md_coords()]
time = self.get_step_props()["time"]
time_step = time[1] - time[0]
species_labels = self.get_species_labels()
symbols = table["label"].map(species_labels).values.tolist()
structures = [
Structure(
lattice=lattice,
species=symbols,
coords=coords,
coords_are_cartesian=True,
site_properties={"gulp_labels": table["label"]},
) for lattice, coords in zip(lattices, frames)
]
return Trajectory.from_structures(structures,
time_step=time_step,
constant_lattice=constant_lattice)
def test_variable_lattice(self):
structure = self.structures[0]
# Generate structures with different lattices
structures = []
for i in range(10):
new_lattice = np.dot(structure.lattice.matrix,
np.diag(1 + np.random.random_sample(3) / 20))
temp_struct = structure.copy()
temp_struct.lattice = Lattice(new_lattice)
structures.append(temp_struct)
traj = Trajectory.from_structures(structures, constant_lattice=False)
# Check if lattices were properly stored
self.assertTrue(
all(
np.allclose(struct.lattice.matrix,
structures[i].lattice.matrix)
for i, struct in enumerate(traj)))
# Check if the file is written correctly when lattice is not constant.
traj.write_Xdatcar(filename="traj_test_XDATCAR")
# Load trajectory from written xdatcar and compare to original
written_traj = Trajectory.from_file("traj_test_XDATCAR",
constant_lattice=False)
self._check_traj_equality(traj, written_traj)
os.remove("traj_test_XDATCAR")
def test_list_slice(self):
sliced_traj = self.traj[[10, 30, 70]]
sliced_traj_from_structs = Trajectory.from_structures([self.structures[i] for i in [10, 30, 70]])
if len(sliced_traj) == len(sliced_traj_from_structs):
self.assertTrue(all([sliced_traj[i] == sliced_traj_from_structs[i] for i in range(len(sliced_traj))]))
else:
self.assertTrue(False)
def test_slice(self):
sliced_traj = self.traj[2:99:3]
sliced_traj_from_structs = Trajectory.from_structures(self.structures[2:99:3])
if len(sliced_traj) == len(sliced_traj_from_structs):
self.assertTrue(all([sliced_traj[i] == sliced_traj_from_structs[i] for i in range(len(sliced_traj))]))
else:
self.assertTrue(False)
def test_displacements(self):
poscar = Poscar.from_file(os.path.join(PymatgenTest.TEST_FILES_DIR, "POSCAR"))
structures = [poscar.structure]
displacements = np.zeros((11, *np.shape(structures[-1].frac_coords)))
for i in range(10):
displacement = np.random.random_sample(np.shape(structures[-1].frac_coords)) / 20
new_coords = displacement + structures[-1].frac_coords
structures.append(Structure(structures[-1].lattice, structures[-1].species, new_coords))
displacements[i + 1, :, :] = displacement
traj = Trajectory.from_structures(structures, constant_lattice=True)
traj.to_displacements()
self.assertTrue(np.allclose(traj.frac_coords, displacements))
def test_displacements(self):
poscar = Poscar.from_file(os.path.join(test_dir, "POSCAR"))
structures = [poscar.structure]
displacements = np.zeros((11, *np.shape(structures[-1].frac_coords)))
for i in range(10):
displacement = np.random.random_sample(np.shape(structures[-1].frac_coords)) / 20
new_coords = displacement + structures[-1].frac_coords
structures.append(Structure(structures[-1].lattice, structures[-1].species, new_coords))
displacements[i+1, :, :] = displacement
traj = Trajectory.from_structures(structures, constant_lattice=True)
traj.to_displacements()
self.assertTrue(np.allclose(traj.frac_coords, displacements))
def test_changing_lattice(self):
structure = self.structures[0]
# Generate structures with different lattices
structures = []
for i in range(10):
new_lattice = np.dot(structure.lattice.matrix, np.diag(1 + np.random.random_sample(3)/20))
temp_struct = structure.copy()
temp_struct.modify_lattice(Lattice(new_lattice))
structures.append(temp_struct)
traj = Trajectory.from_structures(structures, constant_lattice=False)
# Check if lattices were properly stored
self.assertTrue(
all([np.allclose(struct.lattice.matrix, structures[i].lattice.matrix) for i, struct in enumerate(traj)]))
def insert_task(self,
task_doc,
parse_dos=False,
parse_bs=False,
md_structures=False):
"""
Inserts a task document (e.g., as returned by Drone.assimilate()) into the database.
Handles putting DOS and band structure into GridFS as needed.
Args:
task_doc: (dict) the task document
parse_dos: (bool) attempt to parse dos in task_doc and insert into Gridfs
parse_bs: (bool) attempt to parse bandstructure in task_doc and insert into Gridfs
Returns:
(int) - task_id of inserted document
"""
# insert dos into GridFS
if parse_dos and "calcs_reversed" in task_doc:
if "dos" in task_doc["calcs_reversed"][0]: # only store idx=0 DOS
dos = json.dumps(task_doc["calcs_reversed"][0]["dos"],
cls=MontyEncoder)
gfs_id, compression_type = self.insert_gridfs(dos, "dos_fs")
task_doc["calcs_reversed"][0][
"dos_compression"] = compression_type
task_doc["calcs_reversed"][0]["dos_fs_id"] = gfs_id
del task_doc["calcs_reversed"][0]["dos"]
# insert band structure into GridFS
if parse_bs and "calcs_reversed" in task_doc:
if "bandstructure" in task_doc["calcs_reversed"][
0]: # only store idx=0 BS
bs = json.dumps(task_doc["calcs_reversed"][0]["bandstructure"],
cls=MontyEncoder)
gfs_id, compression_type = self.insert_gridfs(
bs, "bandstructure_fs")
task_doc["calcs_reversed"][0][
"bandstructure_compression"] = compression_type
task_doc["calcs_reversed"][0]["bandstructure_fs_id"] = gfs_id
del task_doc["calcs_reversed"][0]["bandstructure"]
# insert structures at each ionic step into GridFS
if md_structures and "calcs_reversed" in task_doc:
# insert ionic steps into gridfs
#TODO: Depricate this and move to only storing trajectory
ionic_steps_json = json.dumps(
task_doc["calcs_reversed"][0]['output']['ionic_steps'],
cls=MontyEncoder)
gfs_id, compression_type = self.insert_gridfs(
ionic_steps_json, "structures_fs")
task_doc["calcs_reversed"][0]['output'][
'ionic_steps_compression'] = compression_type
task_doc["calcs_reversed"][0]['output'][
'ionic_steps_fs_id'] = gfs_id
# Aggregate a trajectory
## Convert from a list of dictionaries to a dictionary of lists
ionic_steps_dict = task_doc["calcs_reversed"][0]['output'][
'ionic_steps']
del task_doc["calcs_reversed"][0]['output']['ionic_steps']
ionic_steps_defaultdict = defaultdict(list)
for d in ionic_steps_dict:
for key, val in d.items():
ionic_steps_defaultdict[key].append(val)
ionic_steps = dict(ionic_steps_defaultdict.items())
## extract structures from dictionary
structures = [
Structure.from_dict(struct)
for struct in ionic_steps['structure']
]
del ionic_steps['structure']
frame_properties = {}
for key in [
'e_fr_energy', 'e_wo_entrp', 'e_0_energy', 'kinetic',
'lattice kinetic', 'nosepot', 'nosekinetic', 'total'
]:
frame_properties[key] = ionic_steps[key]
# Create trajectory
trajectory = Trajectory.from_structures(
structures,
constant_lattice=True,
frame_properties=frame_properties,
time_step=task_doc['input']['incar']['POTIM'])
traj_dict = json.dumps(trajectory.as_dict(), cls=MontyEncoder)
gfs_id, compression_type = self.insert_gridfs(
traj_dict, "trajectories_fs")
task_doc['trajectory'] = {
'formula': trajectory[0].composition.formula.replace(' ', ''),
'temperature': int(task_doc["input"]["incar"]["TEBEG"]),
'compression': compression_type,
'fs_id': gfs_id,
'dimension': list(np.shape(trajectory.frac_coords)),
'time_step': task_doc["input"]["incar"]["POTIM"],
}
# insert the task document and return task_id
return self.insert(task_doc)
