Python serialization.loadfn函数代码示例

def get_solute_def_profile(mpid, solute, solute_conc, T, def_file, sol_file, 
        trial_chem_pot):

    raw_energy_dict = loadfn(def_file,cls=MontyDecoder)
    sol_raw_energy_dict = loadfn(sol_file,cls=MontyDecoder)

    #try:
    e0 = raw_energy_dict[mpid]['e0']
    struct = raw_energy_dict[mpid]['structure']
    vacs = raw_energy_dict[mpid]['vacancies']
    antisites = raw_energy_dict[mpid]['antisites']
    solutes = sol_raw_energy_dict[mpid]['solutes']
    for vac_def in vacs:
        if not vac_def:
            print 'All vacancy defect energies not present'
            continue
    for antisite_def in antisites:
        if not antisite_def:
            print 'All antisite defect energies not preset'
            continue
    for solute_def in solutes:
        if not solute_def:
            print 'All solute defect energies not preset'
            continue

    try:
        def_conc = solute_defect_density(struct, e0, vacs, 
                antisites, solutes, solute_concen=solute_conc, T=T, 
                trial_chem_pot=trial_chem_pot, plot_style="gnuplot")
        return  def_conc
    except:
        raise

    def __init__(self,vasp_task=None,name='vaspfw',handlers=None, handler_params=None, config_file=None):
        self.name = name
        self.handlers=handlers if handlers else []
        self.handler_params=handler_params if handler_params else {}

        if config_file:
            config_dict = loadfn(config_file)
        elif os.path.exists(os.path.join(os.environ['HOME'], 'vasp_interface_defaults.yaml')):
            config_dict = loadfn(os.path.join(os.environ['HOME'], 'vasp_interface_defaults.yaml'))
        else:
            config_dict = {}

        if config_dict:
            self.custodian_opts = config_dict.get('CUSTODIAN_PARAMS', {})
            if self.custodian_opts.get('handlers', []):
                self.handlers.extend(self.custodian_opts.get('handlers', []))
            self.handler_params.update(self.custodian_opts.get('handler_params', {}))

        self.tasks=[vasp_task.input,RunCustodianTask(handlers=self.handlers, 
                handler_params=self.handler_params)] if isinstance(vasp_task, 
                VaspInputInterface) else [vasp_task]
        self.Firework=Firework(self.tasks,name=self.name)
        # Try to establish connection with Launchpad
        try:
            self.LaunchPad=LaunchPad.from_file(os.path.join(os.environ["HOME"], ".fireworks", "my_launchpad.yaml"))
        except:
            self.LaunchPad = None

    def run_task(self, fw_spec):
        # the FW.json/yaml file is mandatory to get the fw_id
        # no need to deserialize the whole FW
        try:
            fw_dict = loadfn('FW.json')
        except IOError:
            try:
                fw_dict = loadfn('FW.yaml')
            except IOError:
                raise RuntimeError("No FW.json nor FW.yaml file present: impossible to determine fw_id")

        fw_id = fw_dict['fw_id']
        lp = LaunchPad.auto_load()
        wf = lp.get_wf_by_fw_id_lzyfw(fw_id)
        wf_module = importlib.import_module(wf.metadata['workflow_module'])
        wf_class = getattr(wf_module, wf.metadata['workflow_class'])

        get_results_method = getattr(wf_class, 'get_final_structure_and_history')
        #TODO: make this more general ... just to test right now ...
        results = get_results_method(wf)

        database = MongoDatabase.from_dict(fw_spec['mongo_database'])

        database.insert_entry({'structure': results['structure'], 'history': results['history']})

        logging.info("Inserted data:\n something")


        return FWAction()

    def test_construction(self):
        edges_frag = {(e[0], e[1]): {"weight":1.0} for e in self.pc_frag1_edges}
        mol_graph = MoleculeGraph.with_edges(self.pc_frag1, edges_frag)
        #dumpfn(mol_graph.as_dict(), os.path.join(module_dir,"pc_frag1_mg.json"))
        ref_mol_graph = loadfn(os.path.join(module_dir, "pc_frag1_mg.json"))
        self.assertEqual(mol_graph, ref_mol_graph)
        self.assertEqual(mol_graph.graph.adj, ref_mol_graph.graph.adj)
        for node in mol_graph.graph.nodes:
            self.assertEqual(mol_graph.graph.node[node]["specie"],
                             ref_mol_graph.graph.node[node]["specie"])
            for ii in range(3):
                self.assertEqual(
                    mol_graph.graph.node[node]["coords"][ii],
                    ref_mol_graph.graph.node[node]["coords"][ii])

        edges_pc = {(e[0], e[1]): {"weight":1.0} for e in self.pc_edges}
        mol_graph = MoleculeGraph.with_edges(self.pc, edges_pc)
        #dumpfn(mol_graph.as_dict(), os.path.join(module_dir,"pc_mg.json"))
        ref_mol_graph = loadfn(os.path.join(module_dir, "pc_mg.json"))
        self.assertEqual(mol_graph, ref_mol_graph)
        self.assertEqual(mol_graph.graph.adj, ref_mol_graph.graph.adj)
        for node in mol_graph.graph:
            self.assertEqual(mol_graph.graph.node[node]["specie"],
                             ref_mol_graph.graph.node[node]["specie"])
            for ii in range(3):
                self.assertEqual(
                    mol_graph.graph.node[node]["coords"][ii],
                    ref_mol_graph.graph.node[node]["coords"][ii])

        mol_graph_edges = MoleculeGraph.with_edges(self.pc, edges=edges_pc)
        mol_graph_strat = MoleculeGraph.with_local_env_strategy(self.pc, OpenBabelNN(), reorder=False, extend_structure=False)
        self.assertTrue(mol_graph_edges.isomorphic_to(mol_graph_strat))

    def run_task(self, fw_spec):
        # the FW.json/yaml file is mandatory to get the fw_id
        # no need to deserialize the whole FW

        if '_add_launchpad_and_fw_id' in fw_spec:
            lp = self.launchpad
            fw_id = self.fw_id
        else:
            try:
                fw_dict = loadfn('FW.json')
            except IOError:
                try:
                    fw_dict = loadfn('FW.yaml')
                except IOError:
                    raise RuntimeError("Launchpad/fw_id not present in spec and No FW.json nor FW.yaml file present: "
                                       "impossible to determine fw_id")
            lp = LaunchPad.auto_load()
            fw_id = fw_dict['fw_id']

        wf = lp.get_wf_by_fw_id_lzyfw(fw_id)

        deleted_files = []
        # iterate over all the fws and launches
        for fw_id, fw in wf.id_fw.items():
            for l in fw.launches+fw.archived_launches:
                l_dir = l.launch_dir

                deleted_files.extend(self.delete_files(os.path.join(l_dir, TMPDIR_NAME)))
                deleted_files.extend(self.delete_files(os.path.join(l_dir, INDIR_NAME)))
                deleted_files.extend(self.delete_files(os.path.join(l_dir, OUTDIR_NAME), self.out_exts))

        logging.info("Deleted files:\n {}".format("\n".join(deleted_files)))

        return FWAction(stored_data={'deleted_files': deleted_files})

def get_solute_def_profile1(mpid, solute, solute_conc, T, def_file, sol_file):

    raw_energy_dict = loadfn(def_file,cls=MontyDecoder)
    sol_raw_energy_dict = loadfn(sol_file,cls=MontyDecoder)

    #try:
    print raw_energy_dict[mpid].keys()
    e0 = raw_energy_dict[mpid]['e0']
    struct = raw_energy_dict[mpid]['structure']
    vacs = raw_energy_dict[mpid]['vacancies']
    antisites = raw_energy_dict[mpid]['antisites']
    solutes = sol_raw_energy_dict[mpid]['solutes']
    for vac_def in vacs:
        if not vac_def:
            print 'All vacancy defect energies not present'
            continue
    for antisite_def in antisites:
        if not antisite_def:
            print 'All antisite defect energies not preset'
            continue
    for solute_def in solutes:
        if not solute_def:
            print 'All solute defect energies not preset'
            continue

    try:
        sol_conc = solute_site_preference_finder(struct, e0, T, vacs, 
                antisites, solutes, solute_conc)#, 
                #trial_chem_pot={'Al':-4.120, 'Ni':-6.5136, 'Ti':-7.7861})
        return sol_conc
    except:
        raise

 def test_check_acc_bzt_bands(self):
     structure = loadfn(os.path.join(test_dir,'boltztrap/structure_mp-12103.json'))
     sbs = loadfn(os.path.join(test_dir,'boltztrap/dft_bs_sym_line.json'))
     sbs_bzt = self.bz_bands.get_symm_bands(structure,-5.25204548)
     corr,werr_vbm,werr_cbm,warn = self.bz_bands.check_acc_bzt_bands(sbs_bzt,sbs)
     self.assertAlmostEqual(corr[2],9.16851750e-05)
     self.assertAlmostEqual(werr_vbm['K-H'],0.18260273521047862)
     self.assertAlmostEqual(werr_cbm['M-K'],0.071552669981356981)
     self.assertFalse(warn)

 def test_get_symm_bands(self):
     structure = loadfn(os.path.join(test_dir,'boltztrap/structure_mp-12103.json'))
     sbs = loadfn(os.path.join(test_dir,'boltztrap/dft_bs_sym_line.json'))
     kpoints = [kp.frac_coords for kp in sbs.kpoints]
     labels_dict = {k: sbs.labels_dict[k].frac_coords for k in sbs.labels_dict}
     for kpt_line,labels_dict in zip([None,sbs.kpoints,kpoints],[None,sbs.labels_dict,labels_dict]):
         print(kpt_line)
         sbs_bzt = self.bz_bands.get_symm_bands(structure,-5.25204548,kpt_line=kpt_line,labels_dict=labels_dict)
         self.assertAlmostEqual(len(sbs_bzt.bands[Spin.up]),20)
         self.assertAlmostEqual(len(sbs_bzt.bands[Spin.up][1]),143)

 def test_dumpf_loadf(self):
     d = {"hello": "world"}
     dumpfn(d, "monte_test.json", indent=4)
     d2 = loadfn("monte_test.json")
     self.assertEqual(d, d2)
     os.remove("monte_test.json")
     dumpfn(d, "monte_test.yaml", default_flow_style=False)
     d2 = loadfn("monte_test.yaml")
     self.assertEqual(d, d2)
     dumpfn(d, "monte_test.yaml", Dumper=Dumper)
     d2 = loadfn("monte_test.yaml")
     os.remove("monte_test.yaml")

    def setUp(self):
        bs = loadfn(os.path.join(test_dir, "PbTe_bandstructure.json"))
        bs_sp = loadfn(os.path.join(test_dir, "N2_bandstructure.json"))
        self.loader = BandstructureLoader(bs, vrun.structures[-1])
        self.assertIsNotNone(self.loader)

        self.loader_sp_up = BandstructureLoader(bs_sp, vrun_sp.structures[-1],spin=1)
        self.loader_sp_dn = BandstructureLoader(bs_sp, vrun_sp.structures[-1],spin=-1)
        self.assertTupleEqual(self.loader_sp_up.ebands.shape, (12, 198))
        self.assertTupleEqual(self.loader_sp_dn.ebands.shape, (12, 198))
        self.assertIsNotNone(self.loader_sp_dn)
        self.assertIsNotNone(self.loader_sp_up)
        
        warnings.simplefilter("ignore")

    def test_get_wf_from_spec_dict(self):
        d = loadfn(os.path.join(os.path.abspath(os.path.dirname(__file__)), "spec.yaml"))
        wf = get_wf_from_spec_dict(self.structure, d)
        self.assertEqual(len(wf.fws), 4)
        for f in wf.fws:
            self.assertEqual(f.spec['_tasks'][-1]["db_file"], "db.json")

        self.assertEqual(sorted([len(v) for v in wf.links.values()]),
                         [0, 0, 1, 2])

        self.assertEqual(wf.name, "Si:band structure")
        d = loadfn(os.path.join(os.path.abspath(os.path.dirname(__file__)),
                                "badspec.yaml"))

        self.assertRaises(ImportError, get_wf_from_spec_dict, self.structure, d)

    def test_get_corrections_for_Mo_Ta_and_W(self):
        os.chdir(os.path.join(ROOT, 'Mo_Ta_W_controls'))
        get_corrections(write_yaml=True)
        test_corrections = loadfn('ion_corrections.yaml')
        control_corrections = {'Mo': 0.379, 'Ta': 0.161, 'W': 0.635}
        test_mus = loadfn('chemical_potentials.yaml')
        control_mus = {'Mo': -8.885, 'Ta': -10.17, 'W': -11.179, 'O': -4.658}

        for elt in control_corrections:
            self.assertEqual(control_corrections[elt], test_corrections[elt])
        for elt in control_mus:
            self.assertEqual(control_mus[elt], test_mus[elt])
        os.system('rm ion_corrections.yaml')
        os.system('rm chemical_potentials.yaml')
        os.chdir(ROOT)

 def from_config(cls, config):
     db_yaml = os.path.expandvars(config.db_yaml)
     db_cfg = loadfn(db_yaml)
     client = MongoClient(db_cfg['host'], db_cfg['port'], j=False)
     db = client[db_cfg['db']]
     try:
         db.authenticate(db_cfg['username'], db_cfg['password'])
     except:
         logger.error('authentication failed for {}'.format(db_yaml))
         sys.exit(1)
     logger.debug('using DB from {}'.format(db_yaml))
     duplicates_file = os.path.expandvars(config.duplicates_file)
     duplicates = loadfn(duplicates_file) \
             if os.path.exists(duplicates_file) else {}
     return OstiMongoAdapter(db, duplicates, config.osti.elink)

def get_solute_def_profile(args):
    if not args.mpid:
        print ('===========\nERROR: mpid is not given.\n===========')
        return
    if not args.solute:
        print ('===========\nERROR: Solute atom is not given.\n===========')
        return

    mpid = args.mpid 
    solute = args.solute 
    solute_conc = args.solute_conc/100.0
    T = args.T 

    def_file = mpid + '_raw_defect_energy.json'
    raw_energy_dict = loadfn(def_file,cls=MontyDecoder)
    sol_file = mpid+'_solute-'+solute+'_raw_defect_energy.json'
    sol_raw_energy_dict = loadfn(sol_file,cls=MontyDecoder)

    #try:
    e0 = raw_energy_dict[mpid]['e0']
    struct = raw_energy_dict[mpid]['structure']
    vacs = raw_energy_dict[mpid]['vacancies']
    antisites = raw_energy_dict[mpid]['antisites']
    solutes = sol_raw_energy_dict[mpid]['solutes']

    for vac_def in vacs:
        if not vac_def:
            print('All vacancy defect energies not present')
            continue
    for antisite_def in antisites:
        if not antisite_def:
            print('All antisite defect energies not preset')
            continue
    for solute_def in solutes:
        if not solute_def:
            print('All solute defect energies not preset')
            continue

    try:
        def_conc = solute_defect_density(struct, e0, vacs, 
                antisites, solutes, solute_concen=solute_conc, T=T, 
                plot_style="gnuplot")
        fl_nm = args.mpid+'_solute-'+args.solute+'_def_concentration.dat'
        with open(fl_nm,'w') as fp: 
            for row in def_conc:
                print >> fp, row 
    except:
        raise

def get_def_profile(mpid, T,  file):

    raw_energy_dict = loadfn(file,cls=MontyDecoder)

    e0 = raw_energy_dict[mpid]['e0']
    struct = raw_energy_dict[mpid]['structure']
    vacs = raw_energy_dict[mpid]['vacancies']
    antisites = raw_energy_dict[mpid]['antisites']
    vacs.sort(key=lambda entry: entry['site_index'])
    antisites.sort(key=lambda entry: entry['site_index'])
    for vac_def in vacs:
        if not vac_def:
            print 'All vacancy defect energies not present'
            continue
    for antisite_def in antisites:
        if not antisite_def:
            print 'All antisite defect energies not preset'
            continue

    try:
        def_conc, def_en, mu = compute_defect_density(struct, e0, vacs, antisites, T,
                plot_style='gnuplot')
        return def_conc, def_en, mu
    except:
        raise