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)