def setUpClass(cls): root = logging.getLogger() root.setLevel(logging.DEBUG) ch = logging.StreamHandler(sys.stdout) ch.setLevel(logging.DEBUG) formatter = logging.Formatter( '%(asctime)s - %(name)s - %(levelname)s - %(message)s') ch.setFormatter(formatter) root.addHandler(ch) cls.tasks = MongoStore("emmet_test", "tasks", lu_field="last_updated") cls.tasks.connect() cleardb(cls.tasks.collection.database) vaspdb = VaspCalcDb(database="emmet_test") tasks_dir = os.path.join(test_dir, "tasks") raw_tasks = glob.glob(os.path.join(test_dir, "tasks", "*.json.gz")) for task_path in raw_tasks: with zopen(task_path) as f: data = f.read().decode() task = json.loads(data) vaspdb.insert_task(task, parse_dos=True, parse_bs=True)
def test_mongodb_more_files(self): try: VaspCalcDb.from_db_file(DB_FILE) except: raise unittest.SkipTest( "Cannot connect to MongoDB! Is the database server running? " "Are the credentials correct?") copy_r(self.vasp_dir, ".") with self.assertWarnsRegex(UserWarning, ".*wrong_file.*"): ft = LobsterRunToDb( calc_loc=True, db_file=DB_FILE, additional_outputs=[ "ICOHPLIST.lobster", "COOPCAR.lobster", "wrong_file", ], ) ft.run_task({"calc_locs": [{"path": "test"}]}) coll = self.get_task_collection("lobster") load_dict = coll.find_one({"formula_pretty": "K2Sn2O3"}) self.assertEqual(load_dict["formula_pretty"], "K2Sn2O3") self.assertListEqual(load_dict["output"]["chargespilling"], [0.008, 0.008]) db = self.get_task_database() gfs = gridfs.GridFS(db, "lobster_files") results = gfs.find({}).count() self.assertEqual(results, 2) for fn in ["ICOHPLIST", "COOPCAR"]: oid = load_dict[fn.lower() + "_id"] results = gfs.find({"_id": oid}).count() self.assertEqual(results, 1)
def check_relax_path(relax_path, db_file, tag, run_isif2, pass_isif4): from atomate.vasp.database import VaspCalcDb if db_file != ">>db_file<<": vasp_db = VaspCalcDb.from_db_file(db_file, admin=True) else: t_file = loadfn(config_to_dict()["FWORKER_LOC"])["env"]["db_file"] vasp_db = VaspCalcDb.from_db_file(t_file, admin=True) if relax_path != '': if os.path.exists(relax_path): return (relax_path, run_isif2, pass_isif4) if vasp_db.db["relax"].count_documents({'metadata.tag': tag}) > 0: items = vasp_db.db["relax"].find({ 'metadata.tag': tag }).sort([('_id', -1)]).limit(1) if os.path.exists(items[0]['path']): print( 'Relax result "%s" with "run_isif2 = %s" and "run_isif4 = %s" has been found, and will be used for new static calculations.' % (items[0]['path'], items[0]['run_isif2'], items[0]['pass_isif4'])) return (items[0]['path'], items[0]['run_isif2'], items[0]['pass_isif4']) else: print( 'Relax result "%s" has been found but NOT exists. Change tag and try again!' % relax_path) return ('', run_isif2, pass_isif4) else: print('No relax result found.') return ('', run_isif2, pass_isif4)
def consistent_check_db(db_file, tag): ''' In the subsequent running(run DFTTK again with the same tag exists in Mongo DB), if phonon method is committed, it'd better to check the lengths of "task" and "phonon" collections. ''' from atomate.vasp.database import VaspCalcDb if db_file != ">>db_file<<": vasp_db = VaspCalcDb.from_db_file(db_file, admin=True) else: t_file = loadfn(config_to_dict()["FWORKER_LOC"])["env"]["db_file"] vasp_db = VaspCalcDb.from_db_file(t_file, admin=True) num_task = vasp_db.collection.count_documents( {'$and': [{ 'metadata.tag': tag }, { 'adopted': True }]}) num_phonon = vasp_db.db['phonon'].count_documents( {'$and': [{ 'metadata.tag': tag }, { 'adopted': True }]}) if num_task == num_phonon: return (True) else: print( 'The records length of "task"(%s) differs to the length of "phonon"(%s) in mongodb.' % (num_task, num_phonon)) return (False)
def mark_adopted(tag, db_file, volumes, phonon=False): mark_adopted_TF(tag, db_file, False, phonon=phonon) # Mark all as adpoted from atomate.vasp.database import VaspCalcDb if db_file != ">>db_file<<": vasp_db = VaspCalcDb.from_db_file(db_file, admin=True) else: t_file = loadfn(config_to_dict()["FWORKER_LOC"])["env"]["db_file"] vasp_db = VaspCalcDb.from_db_file(t_file, admin=True) for volume in volumes: vasp_db.collection.update( { '$and': [{ 'metadata.tag': tag }, { 'output.structure.lattice.volume': volume }] }, {'$set': { 'adopted': True }}, upsert=True, multi=False) # Mark only one if phonon: vasp_db.db['phonon'].update( {'$and': [{ 'metadata.tag': tag }, { 'volume': volume }]}, {'$set': { 'adopted': True }}, upsert=True, multi=False)
def ext_thfind(args, vasp_db=None): """ find the metadata tag that has finished. Parameters STR_FOLDER = args.STRUCTURE_FOLDER folder/file containing structures MATCH_PATTERN = args.MATCH_PATTERN Match patterns for structure file, e.g. *POSCAR RECURSIVE = args.RECURSIVE recursive or not WORKFLOW = args.WORKFLOW workflow, current only get_wf_gibbs """ if args.db_file is not None: vasp_db = VaspCalcDb.from_db_file(args.db_file, admin=False) elif vasp_db is None: db_file = loadfn(config_to_dict()["FWORKER_LOC"])["env"]["db_file"] vasp_db = VaspCalcDb.from_db_file(db_file, admin=False) proc=thfindMDB(args,vasp_db) tags = proc.run_console() if args.get: with open("runs.log", "a") as fp: fp.write ('\nPostprocessing run at {}\n\n'.format(datetime.now())) for t in tags: if isinstance(t,dict): print("\nDownloading data by metadata tag:", t['tag'], "\n") args.metatag = t['tag'] args.phasename = t['phasename'] ext_thelec(args, vasp_db=vasp_db) #args.metatag = None #args.phasename = None else: ext_thelec(args,plotfiles=t, vasp_db=vasp_db) elif args.qha_phonon_repair: from dfttk.scripts.qha_phonon_repair import QHAAnalysis_failure for t in tags: if isinstance(t,dict): tag = t['tag'] qha_phonon = list(vasp_db.db['qha_phonon'].find({'metadata.tag': tag})) if len(qha_phonon) : continue # already ok skip phonon_calculations = list(vasp_db.db['phonon'].find({'$and':[ {'metadata.tag': tag}, {'adopted': True} ]})) T = phonon_calculations[0]['temperatures'] t_min = min(T) t_max = max(T) t_step = T[1]-T[0] print("Repairing data by metadata tag:", t) proc = QHAAnalysis_failure(phonon=True, t_min=t_min, t_max=t_max, t_step=t_step, db_file=db_file, test_failure=False, admin=True, metadata={'tag':tag}, tag=tag) proc.run_task()
def test_QHAAnalysis(): tags = ['5e8b5a18-b2b9-4dcd-81a7-0bd75ee361ec'] print(db_file) for tag in tags: vasp_db = VaspCalcDb.from_db_file(db_file, admin=False) phonon_calculations = list(vasp_db.db['phonon'].find( {'$and': [{ 'metadata.tag': tag }, { 'adopted': True }]})) T = phonon_calculations[0]['temperatures'] t_min = min(T) everyT = 10 t_step = (T[1] - T[0]) * everyT t_max = max(T) nT = int((t_max - t_min) / t_step) t_max = t_min + nT * t_step print("testing tag=", tag) proc = QHAAnalysis(phonon=True, t_min=t_min, t_max=t_max, t_step=t_step, everyT=everyT, db_file=db_file, test=True, tag=tag) proc.run_task("")
def run_task(self, fw_spec): # Get num of electrons and bands from static calc uuid = self["wf_uuid"] db_file = env_chk(self.get("db_file"), fw_spec) db = VaspCalcDb.from_db_file(db_file, admin=True) db.collection = db.db["tasks"] task_doc = db.collection.find_one( {"wf_meta.wf_uuid": uuid, "task_label": "static"}, ["input.parameters"] ) nelec = int(task_doc["input"]["parameters"]["NELECT"]) nbands = int(task_doc["input"]["parameters"]["NBANDS"]) w90_file = [ "num_wann = %d" % (nelec), "num_bands = %d" % (nelec), # 1 band / elec with SOC "spinors=.true.", "num_iter 0", "shell_list 1", "exclude_bands %d-%d" % (nelec + 1, 2 * nbands), ] w90_file = "\n".join(w90_file) with open("wannier90.win", "w") as f: f.write(w90_file) return FWAction()
def run_task(self, fw_spec): wf_uuid = self["wf_uuid"] surfaces = ["kx_0", "kx_1", "ky_0", "ky_1", "kz_0", "kz_1"] d = { "wf_uuid": wf_uuid, "formula": fw_spec["formula"], "reduced_formula": fw_spec["reduced_formula"], "structure": fw_spec["structure"], } for surface in surfaces: if surface in fw_spec.keys(): d[surface] = fw_spec[surface] d = jsanitize(d) # store the results db_file = env_chk(self.get("db_file"), fw_spec) if not db_file: with open("z2pack.json", "w") as f: f.write(json.dumps(d, default=DATETIME_HANDLER)) else: db = VaspCalcDb.from_db_file(db_file, admin=True) db.collection = db.db["z2pack"] db.collection.insert_one(d) logger.info("Z2Pack surface calculation complete.") return FWAction()
def run_task(self, fw_spec): irvsp = self["irvsp_out"] or fw_spec["irvsp_out"] irvsp = jsanitize(irvsp) additional_fields = self.get("additional_fields", {}) d = additional_fields.copy() d["wf_uuid"] = self["wf_uuid"] d["formula"] = fw_spec["formula"] d["efermi"] = fw_spec["efermi"] d["structure"] = fw_spec["structure"] d["irvsp"] = irvsp # store the results db_file = env_chk(self.get("db_file"), fw_spec) if not db_file: with open("irvsp.json", "w") as f: f.write(json.dumps(d, default=DATETIME_HANDLER)) else: db = VaspCalcDb.from_db_file(db_file, admin=True) db.collection = db.db["irvsp"] db.collection.insert_one(d) logger.info("IRVSP calculation complete.") return FWAction()
def mark_adopted_TF(tag, db_file, adpoted, phonon=False): from atomate.vasp.database import VaspCalcDb vasp_db = VaspCalcDb.from_db_file(db_file, admin = True) if vasp_db: vasp_db.collection.update({'metadata.tag': tag}, {'$set': {'adopted': adpoted}}, upsert = True, multi = True) if phonon: vasp_db.db['phonon'].update({'metadata.tag': tag}, {'$set': {'adopted': adpoted}}, upsert = True, multi = True)
def __init__(self, db_file, properties, metadata, **kwargs): #super(RetrieveDataFromDB, self).__init__(db_file, properties, metadata, phonon) self.phonon = kwargs.get('phonon', False) self.properties = properties self.metadata = metadata self.t_min = kwargs.get('t_min', 5.) self.t_max = kwargs.get('t_max', 2000.) self.t_step = kwargs.get('t_step', 5.) self.poisson = kwargs.get('poisson', 0.363615) self.bp2gru = kwargs.get('bp2gru', 1) if phonon: qha_collectioin = 'qha_phonon' else: qha_collectioin = 'qha' property_map = { 'qha': qha_collectioin, 'dos': 'tasks', 'eos': qha_collectioin, 'structure': qha_collectioin, 'phonon': 'phonon', 'debye': qha_collectioin } vasp_db = VaspCalcDb.from_db_file(db_file, admin=True) self.vaspdb = vasp_db prop_items = list(vasp_db.db[property_map[properties]].find({ 'metadata': metadata }).sort('_id', 1)) self.prop_items = prop_items
def get_vol_ene(self): tag = self["tag"] admin = self.get('admin', False) _db_file = self.get('db_file', db_file) vasp_db = VaspCalcDb.from_db_file(db_file=_db_file, admin=admin) volumes, energies, _, _ = get_static_calculations(vasp_db, tag) return volumes, energies
def mark_adopted(tag, db_file, volumes, phonon=False): mark_adopted_TF(tag, db_file, False, phonon=phonon) # Mark all as adpoted from atomate.vasp.database import VaspCalcDb vasp_db = VaspCalcDb.from_db_file(db_file, admin=True) for volume in volumes: vasp_db.collection.update( { '$and': [{ 'metadata.tag': tag }, { 'output.structure.lattice.volume': volume }] }, {'$set': { 'adopted': True }}, upsert=True, multi=False) # Mark only one if phonon: vasp_db.db['phonon'].update( {'$and': [{ 'metadata.tag': tag }, { 'volume': volume }]}, {'$set': { 'adopted': True }}, upsert=True, multi=False)
def run_task(self, fw_spec): db_file = env_chk(self["db_file"], fw_spec) wf_uuid = self["wf_uuid"] hmodel = loadfn("heisenberg_model.json") # Exchange collection mmdb = VaspCalcDb.from_db_file(db_file, admin=True) mmdb.collection = mmdb.db["exchange"] hmodel_dict = hmodel.as_dict() parent_structure = hmodel.structures[0] formula_pretty = parent_structure.composition.reduced_formula wf_meta = {"wf_uuid": wf_uuid} task_doc = { "wf_meta": wf_meta, "formula_pretty": formula_pretty, "nn_cutoff": hmodel.cutoff, "nn_tol": hmodel.tol, "heisenberg_model": hmodel_dict, "task_name": "heisenberg model", } if fw_spec.get("tags", None): task_doc["tags"] = fw_spec["tags"] mmdb.collection.insert_one(task_doc)
def get_detour_workflow(self, next_steps, final_energy): # TODO: add all the necessary arguments and keyword arguments for the new Fireworks # TODO: add update metadata with the input metadata + the symmetry type for static # delayed imports to avoid circular import from fireworks import Workflow from .fworks import RobustOptimizeFW, StaticFW tol_energy = self.get("tol_energy", 0.025) tol_strain = self.get("tol_strain", 0.05) tol_bond = self.get("tol_bond", 0.10) symmetry_options = {"tol_energy": tol_energy, "tol_strain": tol_strain, "tol_bond": tol_bond} # Assume the data for the current step is already in the database db = VaspCalcDb.from_db_file(self.db_file, admin=True).db['relaxations'] def _get_input_structure_for_step(step): # Get the structure of "type" from the "isif" step. relax_data = db.find_one({'$and': [{'tag': self["tag"]}, {'isif': step["structure"]["isif"]}]}) return Structure.from_dict(relax_data[step["structure"]["type"]]) detour_fws = [] for step in next_steps: job_type = step["job_type"] inp_structure = _get_input_structure_for_step(step) if job_type == "static": common_copy = copy.deepcopy(self.get("common_kwargs", {})) md = common_copy.get("metadata", {}) md['symmetry_type'] = step["symmetry_type"] common_copy["metadata"] = md detour_fws.append(StaticFW(inp_structure, **common_copy)) elif job_type == "relax": detour_fws.append(RobustOptimizeFW(inp_structure, isif=step["isif"], override_symmetry_tolerances=symmetry_options, **self["common_kwargs"])) else: raise ValueError(f"Unknown job_type {job_type} for step {step}.") return detour_fws
def run_task(self, fw_spec): symm_check_data = self.check_symmetry() passed = symm_check_data["symmetry_checks_passed"] cur_isif = symm_check_data["isif"] next_steps = self.get_next_steps(passed, cur_isif) symm_check_data.update({ "tag": self["tag"], "metadata": self.get("metadata"), "version_info": { "atomate": atomate_ver, "dfttk": dfttk_ver, "pymatgen": pymatgen_ver, }, "next_steps": next_steps, }) # write to JSON for debugging purposes with open('relaxation_check_summary.json', 'w') as fp: json.dump(symm_check_data, fp) self.db_file = env_chk(self.get("db_file"), fw_spec) vasp_db = VaspCalcDb.from_db_file(self.db_file, admin=True) vasp_db.db['relaxations'].insert_one(symm_check_data) return FWAction(detours=self.get_detour_workflow(next_steps, symm_check_data['final_energy_per_atom']))
def run_task(self, fw_spec): tag = self["tag"] metadata = self.get('metadata', {}) metadata['tag'] = tag unitcell = Structure.from_file('POSCAR-unitcell') supercell_matrix = self['supercell_matrix'] temperatures, f_vib, s_vib, cv_vib, force_constants = get_f_vib_phonopy( unitcell, supercell_matrix, vasprun_path='vasprun.xml', t_min=self['t_min'], t_max=self['t_max'], t_step=self['t_step']) if isinstance(supercell_matrix, np.ndarray): supercell_matrix = supercell_matrix.tolist() # make serializable thermal_props_dict = { 'volume': unitcell.volume, 'F_vib': f_vib.tolist(), 'CV_vib': cv_vib.tolist(), 'S_vib': s_vib.tolist(), 'temperatures': temperatures.tolist(), 'force_constants': force_constants.tolist(), 'metadata': metadata, 'unitcell': unitcell.as_dict(), 'supercell_matrix': supercell_matrix, } # insert into database db_file = env_chk(self["db_file"], fw_spec) vasp_db = VaspCalcDb.from_db_file(db_file, admin=True) vasp_db.db['phonon'].insert_one(thermal_props_dict)
def get_eq_structure_by_metadata(metadata, db_file=None): ''' Get the equilibrium structure by metadata Parameters ---------- metadata: dict The metadata use for searching the database db_file: filepath-like The file path of db.json(the settings for mongodb file) if it is None or >>db_file<<, then using the db_file of fireworks configurations Returns ------- eq_structure: pymatgen.Strucutre object the equilibrium structure ''' if (db_file is None) or (db_file == '>>db_file<<'): db_file = loadfn(config_to_dict()["FWORKER_LOC"])["env"]["db_file"] vasp_db = VaspCalcDb.from_db_file(db_file, admin=True) static_items = list(vasp_db.db['tasks'].find({ 'metadata.tag': metadata['tag'] }).sort('_id', 1)) structure_list = [itemi['output']['structure'] for itemi in static_items] if structure_list: #not empty eq_structure = Structure.from_dict(structure_list[0]) return eq_structure
def test_chgcar_db_read_write(self): # generate a doc from the test folder drone = VaspDrone(parse_chgcar=True, parse_aeccar=True) print(ref_dirs_si['static']) doc = drone.assimilate(ref_dirs_si['static']+'/outputs') # insert the doc make sure that the cc = doc['calcs_reversed'][0]['chgcar'] self.assertAlmostEqual(cc.data['total'].sum()/cc.ngridpts, 8.0, 4) cc = doc['calcs_reversed'][0]['aeccar0'] self.assertAlmostEqual(cc.data['total'].sum()/cc.ngridpts, 23.253588293583313, 4) cc = doc['calcs_reversed'][0]['aeccar2'] self.assertAlmostEqual(cc.data['total'].sum()/cc.ngridpts, 8.01314480789829, 4) mmdb = VaspCalcDb.from_db_file(os.path.join(db_dir, "db.json")) t_id = mmdb.insert_task(doc, use_gridfs=True) # space is freed up after uploading the document self.assertRaises(KeyError, lambda: doc['calcs_reversed'][0]['chgcar']) self.assertRaises(KeyError, lambda: doc['calcs_reversed'][0]['aeccar0']) self.assertRaises(KeyError, lambda: doc['calcs_reversed'][0]['aeccar2']) cc = mmdb.get_chgcar(task_id=t_id) self.assertAlmostEqual(cc.data['total'].sum()/cc.ngridpts, 8.0, 4) dcc = mmdb.get_aeccar(task_id=t_id) self.assertAlmostEqual(dcc['aeccar0'].data['total'].sum()/cc.ngridpts, 23.253588293583313, 4) self.assertAlmostEqual(dcc['aeccar2'].data['total'].sum()/cc.ngridpts, 8.01314480789829, 4) # check the retrieve_task function for the same fake calculation ret_task = mmdb.retrieve_task(t_id) ret_chgcar = ret_task['calcs_reversed'][0]['chgcar'] ret_aeccar0 = ret_task['calcs_reversed'][0]['aeccar0'] ret_aeccar2 = ret_task['calcs_reversed'][0]['aeccar2'] ret_aeccar = ret_aeccar0 + ret_aeccar2 self.assertAlmostEqual(ret_chgcar.data['total'].sum()/ret_chgcar.ngridpts, 8.0, 4) self.assertAlmostEqual(ret_aeccar.data['total'].sum()/ret_aeccar.ngridpts, 31.2667331015, 4)
def get_orig_EV(self, db_file, tag): vasp_db = VaspCalcDb.from_db_file(db_file=db_file, admin=True) volumes, energies, dos_objs, _ \ = get_static_calculations(vasp_db, tag) print('%s Volumes = %s' % (len(volumes), volumes)) print('%s Energies = %s' % (len(energies), energies)) return (list(volumes), list(energies), dos_objs)
def run_task(self, fw_spec): from atomate.vasp.analysis.phonopy import get_phonopy_thermal_expansion tag = self["tag"] db_file = env_chk(self.get("db_file"), fw_spec) t_step = self.get("t_step", 10) t_min = self.get("t_min", 0) t_max = self.get("t_max", 1000) mesh = self.get("mesh", [20, 20, 20]) eos = self.get("eos", "vinet") pressure = self.get("pressure", 0.0) summary_dict = {} mmdb = VaspCalcDb.from_db_file(db_file, admin=True) # get the optimized structure d = mmdb.collection.find_one( {"task_label": "{} structure optimization".format(tag)}) structure = Structure.from_dict( d["calcs_reversed"][-1]["output"]['structure']) summary_dict["structure"] = structure.as_dict() summary_dict["formula_pretty"] = structure.composition.reduced_formula # get the data(energy, volume, force constant) from the deformation runs docs = mmdb.collection.find({ "task_label": { "$regex": "{} thermal_expansion*".format(tag) }, "formula_pretty": structure.composition.reduced_formula }) energies = [] volumes = [] force_constants = [] for d in docs: s = Structure.from_dict( d["calcs_reversed"][-1]["output"]['structure']) energies.append(d["calcs_reversed"][-1]["output"]['energy']) volumes.append(s.volume) force_constants.append( d["calcs_reversed"][-1]["output"]['force_constants']) summary_dict["energies"] = energies summary_dict["volumes"] = volumes summary_dict["force_constants"] = force_constants alpha, T = get_phonopy_thermal_expansion(energies, volumes, force_constants, structure, t_min, t_step, t_max, mesh, eos, pressure) summary_dict["alpha"] = alpha summary_dict["T"] = T with open("thermal_expansion.json", "w") as f: f.write(json.dumps(summary_dict, default=DATETIME_HANDLER)) # TODO: @matk86 - there needs to be a way to insert this into a database! And also # a builder to put it into materials collection... -computron logger.info("Thermal expansion coefficient calculation complete.")
def test_mongodb(self): try: VaspCalcDb.from_db_file(DB_FILE) except: raise unittest.SkipTest( "Cannot connect to MongoDB! Is the database server running? " "Are the credentials correct?") copy_r(self.vasp_dir, ".") ft = LobsterRunToDb(calc_loc=True, db_file=DB_FILE) ft.run_task({"calc_locs": [{"path": "test"}]}) coll = self.get_task_collection("lobster") load_dict = coll.find_one({"formula_pretty": "K2Sn2O3"}) self.assertEqual(load_dict["formula_pretty"], "K2Sn2O3") self.assertListEqual(load_dict["output"]["chargespilling"], [0.008, 0.008]) self.assertNotIn("lobster_icohplist", load_dict)
def test_insert_maggma_store(self): # generate a doc from the test folder doc = {"a": 1, "b": 2} with mock_s3(): conn = boto3.client("s3") conn.create_bucket(Bucket="test_bucket") mmdb = VaspCalcDb.from_db_file(os.path.join(db_dir, "db_aws.json")) mmdb_changed_prefix = VaspCalcDb.from_db_file( os.path.join(db_dir, "db_aws_prefix.json")) fs_id, compress_type = mmdb.insert_maggma_store(doc, 'store1', oid='1') fs_id, compress_type = mmdb_changed_prefix.insert_maggma_store( doc, 'store1', oid='1') assert fs_id == '1' assert compress_type == 'zlib' doc['task_id'] = 'mp-1' _, _ = mmdb.insert_maggma_store(doc, 'store2', oid='2') assert set(mmdb._maggma_stores.keys()) == {'store1', 'store2'} with mmdb._maggma_stores['store1'] as store: self.assertTrue(store.compress == True) self.assertTrue( store.query_one({'fs_id': '1'}) == { 'fs_id': '1', 'maggma_store_type': 'S3Store', 'compression': 'zlib', 'data': { 'a': 1, 'b': 2 } }) with mmdb._maggma_stores['store2'] as store: self.assertTrue(store.compress == True) self.assertTrue( store.query_one({'task_id': 'mp-1'}) == { 'fs_id': '2', 'maggma_store_type': 'S3Store', 'compression': 'zlib', 'data': { 'a': 1, 'b': 2, 'task_id': 'mp-1' }, 'task_id': 'mp-1' })
def run_task(self, fw_spec): additional_fields = self.get("additional_fields", {}) # pass the additional_fields first to avoid overriding BoltztrapAnalyzer items d = additional_fields.copy() btrap_dir = os.path.join(os.getcwd(), "boltztrap") d["boltztrap_dir"] = btrap_dir bta = BoltztrapAnalyzer.from_files(btrap_dir) d.update(bta.as_dict()) d["scissor"] = bta.intrans["scissor"] # trim the output for x in ['cond', 'seebeck', 'kappa', 'hall', 'mu_steps', 'mu_doping', 'carrier_conc']: del d[x] if not self.get("hall_doping"): del d["hall_doping"] bandstructure_dir = os.getcwd() d["bandstructure_dir"] = bandstructure_dir # add the structure v, o = get_vasprun_outcar(bandstructure_dir, parse_eigen=False, parse_dos=False) structure = v.final_structure d["structure"] = structure.as_dict() d["formula_pretty"] = structure.composition.reduced_formula d.update(get_meta_from_structure(structure)) # add the spacegroup sg = SpacegroupAnalyzer(Structure.from_dict(d["structure"]), 0.1) d["spacegroup"] = {"symbol": sg.get_space_group_symbol(), "number": sg.get_space_group_number(), "point_group": sg.get_point_group_symbol(), "source": "spglib", "crystal_system": sg.get_crystal_system(), "hall": sg.get_hall()} d["created_at"] = datetime.utcnow() db_file = env_chk(self.get('db_file'), fw_spec) if not db_file: del d["dos"] with open(os.path.join(btrap_dir, "boltztrap.json"), "w") as f: f.write(json.dumps(d, default=DATETIME_HANDLER)) else: mmdb = VaspCalcDb.from_db_file(db_file, admin=True) # dos gets inserted into GridFS dos = json.dumps(d["dos"], cls=MontyEncoder) fsid, compression = mmdb.insert_gridfs(dos, collection="dos_boltztrap_fs", compress=True) d["dos_boltztrap_fs_id"] = fsid del d["dos"] mmdb.db.boltztrap.insert(d)
def get_static_structure_by_metadata(metadata, db_file=None): ''' Get the static structure by metadata Parameters ---------- metadata: dict The metadata use for searching the database db_file: filepath-like The file path of db.json(the settings for mongodb file) if it is None or >>db_file<<, then using the db_file of fireworks configurations Returns ------- structure_list: list The list of different structures. The structures are sorted by volume ''' if (db_file is None) or (db_file == '>>db_file<<'): db_file = loadfn(config_to_dict()["FWORKER_LOC"])["env"]["db_file"] vasp_db = VaspCalcDb.from_db_file(db_file, admin=True) static_items = list( vasp_db.collection.find( {'$and': [{ 'metadata.tag': metadata['tag'] }, { 'adopted': True }]})) #static_items = list(vasp_db.db['tasks'].find({'metadata.tag': metadata})) structure_list = [ Structure.from_dict(itemi['output']['structure']) for itemi in static_items ] volumes = [ itemi['output']['structure']['lattice']['volume'] for itemi in static_items ] energies = [itemi['output']['energy_per_atom'] for itemi in static_items] band_gap = [] static_settings = {} for itemi in static_items: if itemi['output']['is_gap_direct']: band_gap.append(itemi['output']['bandgap']) else: band_gap.append(itemi['output']['direct_gap']) if not static_settings: pot = itemi['input']['pseudo_potential']['functional'].upper() if pot == "": pot = itemi['orig_inputs']['potcar']['functional'].upper() if pot == 'Perdew-Zunger81'.upper(): pot = "LDA" static_settings['user_potcar_functional'] = pot static_settings['user_incar_settings'] = itemi['input']['incar'] structure_list = sort_x_by_y(structure_list, volumes) band_gap = sort_x_by_y(band_gap, volumes) energies = sort_x_by_y(energies, volumes) volumes = sorted(volumes) return (structure_list, energies, band_gap, static_settings, volumes)
def run_task(self, fw_spec): db_file = env_chk(self["db_file"], fw_spec) vasp_db = VaspCalcDb.from_db_file(db_file, admin=True) content = {} content["metadata"] = self.get('metadata', {}) content["path"] = fw_spec["calc_locs"][-1]["path"] content["run_isif2"] = self.get('run_isif2') content["pass_isif4"] = self.get('pass_isif4') vasp_db.db["relax"].insert_one(content)
def run_task(self, fw_spec): from atomate.vasp.database import VaspCalcDb db_file = env_chk(self["db_file"], fw_spec) vasp_db = VaspCalcDb.from_db_file(db_file, admin=True) content = {} content["metadata"] = self.get('metadata', {}) content["path"] = fw_spec["calc_locs"][-1]["path"] content["Pos_Shape_relax"] = self.get('Pos_Shape_relax') vasp_db.db["relax"].insert_one(content)
def run_task(self, fw_spec): from pymatgen.analysis.eos import EOS eos = self.get("eos", "vinet") tag = self["tag"] db_file = env_chk(self.get("db_file"), fw_spec) summary_dict = {"eos": eos} to_db = self.get("to_db", True) # collect and store task_id of all related tasks to make unique links with "tasks" collection all_task_ids = [] mmdb = VaspCalcDb.from_db_file(db_file, admin=True) # get the optimized structure d = mmdb.collection.find_one({"task_label": "{} structure optimization".format(tag)}) all_task_ids.append(d["task_id"]) structure = Structure.from_dict(d["calcs_reversed"][-1]["output"]['structure']) summary_dict["structure"] = structure.as_dict() summary_dict["formula_pretty"] = structure.composition.reduced_formula # get the data(energy, volume, force constant) from the deformation runs docs = mmdb.collection.find({"task_label": {"$regex": "{} bulk_modulus*".format(tag)}, "formula_pretty": structure.composition.reduced_formula}) energies = [] volumes = [] for d in docs: s = Structure.from_dict(d["calcs_reversed"][-1]["output"]['structure']) energies.append(d["calcs_reversed"][-1]["output"]['energy']) volumes.append(s.volume) all_task_ids.append(d["task_id"]) summary_dict["energies"] = energies summary_dict["volumes"] = volumes summary_dict["all_task_ids"] = all_task_ids # fit the equation of state eos = EOS(eos) eos_fit = eos.fit(volumes, energies) summary_dict["bulk_modulus"] = eos_fit.b0_GPa # TODO: find a better way for passing tags of the entire workflow to db - albalu if fw_spec.get("tags", None): summary_dict["tags"] = fw_spec["tags"] summary_dict["results"] = dict(eos_fit.results) summary_dict["created_at"] = datetime.utcnow() # db_file itself is required but the user can choose to pass the results to db or not if to_db: mmdb.collection = mmdb.db["eos"] mmdb.collection.insert_one(summary_dict) else: with open("bulk_modulus.json", "w") as f: f.write(json.dumps(summary_dict, default=DATETIME_HANDLER)) # TODO: @matk86 - there needs to be a builder to put it into materials collection... -computron logger.info("Bulk modulus calculation complete.")
def calcdb_from_mgrant(spec_or_dbfile): if os.path.exists(spec_or_dbfile): return VaspCalcDb.from_db_file(spec_or_dbfile) client = Client() role = "rw" # NOTE need write access to source to ensure indexes host, dbname_or_alias = spec_or_dbfile.split("/", 1) auth = client.get_auth(host, dbname_or_alias, role) if auth is None: raise Exception("No valid auth credentials available!") return VaspCalcDb( auth["host"], 27017, auth["db"], "tasks", auth["username"], auth["password"], authSource=auth["db"], )
def run_task(self, fw_spec): nm_eigenvecs = np.array(fw_spec["normalmodes"]["eigenvecs"]) nm_eigenvals = np.array(fw_spec["normalmodes"]["eigenvals"]) nm_norms = np.linalg.norm(nm_eigenvecs, axis=2) structure = fw_spec["normalmodes"]["structure"] masses = np.array([site.specie.data['Atomic mass'] for site in structure]) nm_norms = nm_norms / np.sqrt(masses) # eigenvectors in vasprun.xml are not divided by sqrt(M_i) # To get the actual eigenvals, the values read from vasprun.xml must be multiplied by -1. # frequency_i = sqrt(-e_i) # To convert the frequency to THZ: multiply sqrt(-e_i) by 15.633 # To convert the frequency to cm^-1: multiply sqrt(-e_i) by 82.995 nm_frequencies = np.sqrt(np.abs(nm_eigenvals)) * 82.995 # cm^-1 d = {"structure": structure.as_dict(), "formula_pretty": structure.composition.reduced_formula, "normalmodes": {"eigenvals": fw_spec["normalmodes"]["eigenvals"], "eigenvecs": fw_spec["normalmodes"]["eigenvecs"] }, "frequencies": nm_frequencies.tolist()} # store the displacement & epsilon for each mode in a dictionary mode_disps = fw_spec["raman_epsilon"].keys() modes_eps_dict = defaultdict(list) for md in mode_disps: modes_eps_dict[fw_spec["raman_epsilon"][md]["mode"]].append( [fw_spec["raman_epsilon"][md]["displacement"], fw_spec["raman_epsilon"][md]["epsilon"]]) # raman tensor = finite difference derivative of epsilon wrt displacement. raman_tensor_dict = {} scale = np.sqrt(structure.volume/2.0) / 4.0 / np.pi for k, v in modes_eps_dict.items(): raman_tensor = (np.array(v[0][1]) - np.array(v[1][1])) / (v[0][0] - v[1][0]) # frequency in cm^-1 omega = nm_frequencies[k] if nm_eigenvals[k] > 0: logger.warn("Mode: {} is UNSTABLE. Freq(cm^-1) = {}".format(k, -omega)) raman_tensor = scale * raman_tensor * np.sum(nm_norms[k]) / np.sqrt(omega) raman_tensor_dict[str(k)] = raman_tensor.tolist() d["raman_tensor"] = raman_tensor_dict d["state"] = "successful" # store the results db_file = env_chk(self.get("db_file"), fw_spec) if not db_file: with open("raman.json", "w") as f: f.write(json.dumps(d, default=DATETIME_HANDLER)) else: db = VaspCalcDb.from_db_file(db_file, admin=True) db.collection = db.db["raman"] db.collection.insert_one(d) logger.info("Raman tensor calculation complete.") return FWAction()
def run_task(self, fw_spec): ref_file = self.get("json_filename", "task.json") calc_dir = self.get("calc_dir", os.getcwd()) with open(os.path.join(calc_dir, ref_file), "r") as fp: task_doc = json.load(fp) db_file = env_chk(self.get('db_file'), fw_spec) if not db_file: with open("task.json", "w") as f: f.write(json.dumps(task_doc, default=DATETIME_HANDLER)) else: mmdb = VaspCalcDb.from_db_file(db_file, admin=True) mmdb.insert(task_doc)
def run_task(self, fw_spec): from atomate.vasp.analysis.phonopy import get_phonopy_thermal_expansion tag = self["tag"] db_file = env_chk(self.get("db_file"), fw_spec) t_step = self.get("t_step", 10) t_min = self.get("t_min", 0) t_max = self.get("t_max", 1000) mesh = self.get("mesh", [20, 20, 20]) eos = self.get("eos", "vinet") pressure = self.get("pressure", 0.0) summary_dict = {} mmdb = VaspCalcDb.from_db_file(db_file, admin=True) # get the optimized structure d = mmdb.collection.find_one({"task_label": "{} structure optimization".format(tag)}) structure = Structure.from_dict(d["calcs_reversed"][-1]["output"]['structure']) summary_dict["structure"] = structure.as_dict() summary_dict["formula_pretty"] = structure.composition.reduced_formula # get the data(energy, volume, force constant) from the deformation runs docs = mmdb.collection.find({"task_label": {"$regex": "{} thermal_expansion*".format(tag)}, "formula_pretty": structure.composition.reduced_formula}) energies = [] volumes = [] force_constants = [] for d in docs: s = Structure.from_dict(d["calcs_reversed"][-1]["output"]['structure']) energies.append(d["calcs_reversed"][-1]["output"]['energy']) volumes.append(s.volume) force_constants.append(d["calcs_reversed"][-1]["output"]['force_constants']) summary_dict["energies"] = energies summary_dict["volumes"] = volumes summary_dict["force_constants"] = force_constants alpha, T = get_phonopy_thermal_expansion(energies, volumes, force_constants, structure, t_min, t_step, t_max, mesh, eos, pressure) summary_dict["alpha"] = alpha summary_dict["T"] = T with open("thermal_expansion.json", "w") as f: f.write(json.dumps(summary_dict, default=DATETIME_HANDLER)) # TODO: @matk86 - there needs to be a way to insert this into a database! And also # a builder to put it into materials collection... -computron logger.info("Thermal expansion coefficient calculation complete.")
def run_task(self, fw_spec): gibbs_dict = {} tag = self["tag"] t_step = self.get("t_step", 10) t_min = self.get("t_min", 0) t_max = self.get("t_max", 1000) mesh = self.get("mesh", [20, 20, 20]) eos = self.get("eos", "vinet") qha_type = self.get("qha_type", "debye_model") pressure = self.get("pressure", 0.0) poisson = self.get("poisson", 0.25) anharmonic_contribution = self.get("anharmonic_contribution", False) gibbs_dict["metadata"] = self.get("metadata", {}) db_file = env_chk(self.get("db_file"), fw_spec) mmdb = VaspCalcDb.from_db_file(db_file, admin=True) # get the optimized structure d = mmdb.collection.find_one({"task_label": "{} structure optimization".format(tag)}, {"calcs_reversed": 1}) structure = Structure.from_dict(d["calcs_reversed"][-1]["output"]['structure']) gibbs_dict["structure"] = structure.as_dict() gibbs_dict["formula_pretty"] = structure.composition.reduced_formula # get the data(energy, volume, force constant) from the deformation runs docs = mmdb.collection.find({"task_label": {"$regex": "{} gibbs*".format(tag)}, "formula_pretty": structure.composition.reduced_formula}, {"calcs_reversed": 1}) energies = [] volumes = [] force_constants = [] for d in docs: s = Structure.from_dict(d["calcs_reversed"][-1]["output"]['structure']) energies.append(d["calcs_reversed"][-1]["output"]['energy']) if qha_type not in ["debye_model"]: force_constants.append(d["calcs_reversed"][-1]["output"]['force_constants']) volumes.append(s.volume) gibbs_dict["energies"] = energies gibbs_dict["volumes"] = volumes if qha_type not in ["debye_model"]: gibbs_dict["force_constants"] = force_constants try: # use quasi-harmonic debye approximation if qha_type in ["debye_model"]: from pymatgen.analysis.quasiharmonic import QuasiharmonicDebyeApprox qhda = QuasiharmonicDebyeApprox(energies, volumes, structure, t_min, t_step, t_max, eos, pressure=pressure, poisson=poisson, anharmonic_contribution=anharmonic_contribution) gibbs_dict.update(qhda.get_summary_dict()) gibbs_dict["anharmonic_contribution"] = anharmonic_contribution gibbs_dict["success"] = True # use the phonopy interface else: from atomate.vasp.analysis.phonopy import get_phonopy_gibbs G, T = get_phonopy_gibbs(energies, volumes, force_constants, structure, t_min, t_step, t_max, mesh, eos, pressure) gibbs_dict["gibbs_free_energy"] = G gibbs_dict["temperatures"] = T gibbs_dict["success"] = True # quasi-harmonic analysis failed, set the flag to false except: import traceback logger.warn("Quasi-harmonic analysis failed!") gibbs_dict["success"] = False gibbs_dict["traceback"] = traceback.format_exc() gibbs_dict['metadata'].update({"task_label_tag": tag}) gibbs_dict["created_at"] = datetime.utcnow() gibbs_dict = jsanitize(gibbs_dict) # TODO: @matk86: add a list of task_ids that were used to construct the analysis to DB? # -computron if not db_file: dump_file = "gibbs.json" logger.info("Dumping the analysis summary to {}".format(dump_file)) with open(dump_file, "w") as f: f.write(json.dumps(gibbs_dict, default=DATETIME_HANDLER)) else: coll = mmdb.db["gibbs_tasks"] coll.insert_one(gibbs_dict) logger.info("Gibbs free energy calculation complete.") if not gibbs_dict["success"]: return FWAction(defuse_children=True)
def run_task(self, fw_spec): ref_struct = self['structure'] d = { "analysis": {}, "initial_structure": self['structure'].as_dict() } # Get optimized structure calc_locs_opt = [cl for cl in fw_spec.get('calc_locs', []) if 'optimiz' in cl['name']] if calc_locs_opt: optimize_loc = calc_locs_opt[-1]['path'] logger.info("Parsing initial optimization directory: {}".format(optimize_loc)) drone = VaspDrone() optimize_doc = drone.assimilate(optimize_loc) opt_struct = Structure.from_dict(optimize_doc["calcs_reversed"][0]["output"]["structure"]) d.update({"optimized_structure": opt_struct.as_dict()}) ref_struct = opt_struct eq_stress = -0.1*Stress(optimize_doc["calcs_reversed"][0]["output"]["ionic_steps"][-1]["stress"]) else: eq_stress = None if self.get("fw_spec_field"): d.update({self.get("fw_spec_field"): fw_spec.get(self.get("fw_spec_field"))}) # Get the stresses, strains, deformations from deformation tasks defo_dicts = fw_spec["deformation_tasks"].values() stresses, strains, deformations = [], [], [] for defo_dict in defo_dicts: stresses.append(Stress(defo_dict["stress"])) strains.append(Strain(defo_dict["strain"])) deformations.append(Deformation(defo_dict["deformation_matrix"])) # Add derived stresses and strains if symmops is present for symmop in defo_dict.get("symmops", []): stresses.append(Stress(defo_dict["stress"]).transform(symmop)) strains.append(Strain(defo_dict["strain"]).transform(symmop)) deformations.append(Deformation(defo_dict["deformation_matrix"]).transform(symmop)) stresses = [-0.1*s for s in stresses] pk_stresses = [stress.piola_kirchoff_2(deformation) for stress, deformation in zip(stresses, deformations)] d['fitting_data'] = {'cauchy_stresses': stresses, 'eq_stress': eq_stress, 'strains': strains, 'pk_stresses': pk_stresses, 'deformations': deformations } logger.info("Analyzing stress/strain data") # TODO: @montoyjh: what if it's a cubic system? don't need 6. -computron # TODO: Can add population method but want to think about how it should # be done. -montoyjh order = self.get('order', 2) if order > 2: method = 'finite_difference' else: method = self.get('fitting_method', 'finite_difference') if method == 'finite_difference': result = ElasticTensorExpansion.from_diff_fit( strains, pk_stresses, eq_stress=eq_stress, order=order) if order == 2: result = ElasticTensor(result[0]) elif method == 'pseudoinverse': result = ElasticTensor.from_pseudoinverse(strains, pk_stresses) elif method == 'independent': result = ElasticTensor.from_independent_strains(strains, pk_stresses, eq_stress=eq_stress) else: raise ValueError("Unsupported method, method must be finite_difference, " "pseudoinverse, or independent") ieee = result.convert_to_ieee(ref_struct) d.update({ "elastic_tensor": { "raw": result.voigt, "ieee_format": ieee.voigt } }) if order == 2: d.update({"derived_properties": ieee.get_structure_property_dict(ref_struct)}) else: soec = ElasticTensor(ieee[0]) d.update({"derived_properties": soec.get_structure_property_dict(ref_struct)}) d["formula_pretty"] = ref_struct.composition.reduced_formula d["fitting_method"] = method d["order"] = order d = jsanitize(d) # Save analysis results in json or db db_file = env_chk(self.get('db_file'), fw_spec) if not db_file: with open("elasticity.json", "w") as f: f.write(json.dumps(d, default=DATETIME_HANDLER)) else: db = VaspCalcDb.from_db_file(db_file, admin=True) db.collection = db.db["elasticity"] db.collection.insert_one(d) logger.info("Elastic analysis complete.") return FWAction()
def run_task(self, fw_spec): wfid = list(filter(lambda x: 'wfid' in x, fw_spec['tags'])).pop() db_file = env_chk(self.get("db_file"), fw_spec) vaspdb = VaspCalcDb.from_db_file(db_file, admin=True) # ferroelectric workflow groups calculations by generated wfid tag polarization_tasks = vaspdb.collection.find({"tags": wfid, "task_label": {"$regex": ".*polarization"}}) tasks = [] outcars = [] structure_dicts = [] sort_weight = [] energies_per_atom = [] energies = [] zval_dicts = [] for p in polarization_tasks: # Grab data from each polarization task energies_per_atom.append(p['calcs_reversed'][0]['output']['energy_per_atom']) energies.append(p['calcs_reversed'][0]['output']['energy']) tasks.append(p['task_label']) outcars.append(p['calcs_reversed'][0]['output']['outcar']) structure_dicts.append(p['calcs_reversed'][0]['input']['structure']) zval_dicts.append(p['calcs_reversed'][0]['output']['outcar']['zval_dict']) # Add weight for sorting # Want polarization calculations in order of nonpolar to polar for Polarization object # This number needs to be bigger than the number of calculations max_sort_weight = 1000000 if 'nonpolar_polarization' in p['task_label']: sort_weight.append(0) elif "polar_polarization" in p['task_label']: sort_weight.append(max_sort_weight) elif "interpolation_" in p['task_label']: num = 0 part = re.findall(r'interpolation_[0-9]+_polarization', p['task_label']) if part != []: part2 = re.findall(r'[0-9]+', part.pop()) if part2 != []: num = part2.pop() sort_weight.append(max_sort_weight - int(num)) # Sort polarization tasks # nonpolar -> interpolation_n -> interpolation_n-1 -> ... -> interpolation_1 -> polar data = zip(tasks, structure_dicts, outcars, energies_per_atom, energies, sort_weight) data = sorted(data,key=lambda x: x[-1]) # Get the tasks, structures, etc in sorted order from the zipped data. tasks, structure_dicts, outcars, energies_per_atom, energies, sort_weight = zip(*data) structures = [Structure.from_dict(structure) for structure in structure_dicts] # If LCALCPOL = True then Outcar will parse and store the pseudopotential zvals. zval_dict = zval_dicts.pop() # Assumes that we want to calculate the ionic contribution to the dipole moment. # VASP's ionic contribution is sometimes strange. # See pymatgen.analysis.ferroelectricity.polarization.Polarization for details. p_elecs = [outcar['p_elec'] for outcar in outcars] p_ions = [get_total_ionic_dipole(structure, zval_dict) for structure in structures] polarization = Polarization(p_elecs, p_ions, structures) p_change = polarization.get_polarization_change().A1.tolist() p_norm = polarization.get_polarization_change_norm() polarization_max_spline_jumps = polarization.max_spline_jumps() same_branch = polarization.get_same_branch_polarization_data(convert_to_muC_per_cm2=True) raw_elecs, raw_ions = polarization.get_pelecs_and_pions() quanta = polarization.get_lattice_quanta(convert_to_muC_per_cm2=True) energy_trend = EnergyTrend(energies_per_atom) energy_max_spline_jumps = energy_trend.max_spline_jump() polarization_dict = {} def split_abc(var, var_name): d = {} for i, j in enumerate('abc'): d.update({var_name + "_{}".format(j): var[:, i].A1.tolist()}) return d # Add some sort of id for the structures? Like cid but more general? # polarization_dict.update({'cid': cid}) # General information polarization_dict.update({'pretty_formula': structures[0].composition.reduced_formula}) polarization_dict.update({'wfid': wfid}) polarization_dict.update({'task_label_order': tasks}) # Polarization information polarization_dict.update({'polarization_change': p_change}) polarization_dict.update({'polarization_change_norm': p_norm}) polarization_dict.update({'polarization_max_spline_jumps': polarization_max_spline_jumps}) polarization_dict.update(split_abc(same_branch, "same_branch_polarization")) polarization_dict.update(split_abc(raw_elecs, "raw_electron_polarization")) polarization_dict.update(split_abc(raw_ions, "raw_electron_polarization")) polarization_dict.update(split_abc(quanta, "polarization_quanta")) polarization_dict.update({"zval_dict": zval_dict}) # Energy information polarization_dict.update({'energy_per_atom_max_spline_jumps': energy_max_spline_jumps}) polarization_dict.update({"energies": energies}) polarization_dict.update({"energies_per_atom": energies_per_atom}) polarization_dict.update({'outcars': outcars}) polarization_dict.update({"structures": structure_dicts}) # Write all the info to db. coll = vaspdb.db["polarization_tasks"] coll.insert_one(polarization_dict)
def run_task(self, fw_spec): # get the directory that contains the VASP dir to parse calc_dir = os.getcwd() if "calc_dir" in self: calc_dir = self["calc_dir"] elif self.get("calc_loc"): calc_dir = get_calc_loc(self["calc_loc"], fw_spec["calc_locs"])["path"] # parse the VASP directory logger.info("PARSING DIRECTORY: {}".format(calc_dir)) drone = VaspDrone(additional_fields=self.get("additional_fields"), parse_dos=self.get("parse_dos", False), bandstructure_mode=self.get("bandstructure_mode", False), parse_chgcar=self.get("parse_chgcar", False), parse_aeccar=self.get("parse_aeccar", False)) # assimilate (i.e., parse) task_doc = drone.assimilate(calc_dir) # Check for additional keys to set based on the fw_spec if self.get("fw_spec_field"): task_doc.update(fw_spec[self.get("fw_spec_field")]) # get the database connection db_file = env_chk(self.get('db_file'), fw_spec) # db insertion or taskdoc dump if not db_file: with open("task.json", "w") as f: f.write(json.dumps(task_doc, default=DATETIME_HANDLER)) else: mmdb = VaspCalcDb.from_db_file(db_file, admin=True) t_id = mmdb.insert_task( task_doc, use_gridfs=self.get("parse_dos", False) or bool(self.get("bandstructure_mode", False)) or self.get("parse_chgcar", False) or self.get("parse_aeccar", False)) logger.info("Finished parsing with task_id: {}".format(t_id)) defuse_children = False if task_doc["state"] != "successful": defuse_unsuccessful = self.get("defuse_unsuccessful", DEFUSE_UNSUCCESSFUL) if defuse_unsuccessful is True: defuse_children = True elif defuse_unsuccessful is False: pass elif defuse_unsuccessful == "fizzle": raise RuntimeError( "VaspToDb indicates that job is not successful " "(perhaps your job did not converge within the " "limit of electronic/ionic iterations)!") else: raise RuntimeError("Unknown option for defuse_unsuccessful: " "{}".format(defuse_unsuccessful)) task_fields_to_push = self.get("task_fields_to_push", None) update_spec = {} if task_fields_to_push: if isinstance(task_fields_to_push, dict): for key, path_in_task_doc in task_fields_to_push.items(): if has(task_doc, path_in_task_doc): update_spec[key] = get(task_doc, path_in_task_doc) else: logger.warn("Could not find {} in task document. Unable to push to next firetask/firework".format(path_in_task_doc)) else: raise RuntimeError("Inappropriate type {} for task_fields_to_push. It must be a " "dictionary of format: {key: path} where key refers to a field " "in the spec and path is a full mongo-style path to a " "field in the task document".format(type(task_fields_to_push))) return FWAction(stored_data={"task_id": task_doc.get("task_id", None)}, defuse_children=defuse_children, update_spec=update_spec)