Python pymatgen.Structure类代码示例

    def run_task(self, fw_spec):

        from pymatgen.analysis.eos import EOS

        tag = self["tag"]
        db_file = env_chk(self.get("db_file"), fw_spec)
        summary_dict = {"eos": self["eos"]}

        mmdb = MMVaspDb.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()

        # 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)
        summary_dict["energies"] = energies
        summary_dict["volumes"] = volumes

        # fit the equation of state
        eos = EOS(self["eos"])
        eos_fit = eos.fit(volumes, energies)
        summary_dict["results"] = dict(eos_fit.results)

        with open("bulk_modulus.json", "w") as f:
            f.write(json.dumps(summary_dict, default=DATETIME_HANDLER))

        logger.info("BULK MODULUS CALCULATION COMPLETE")

    def _match_material(self, doc):
        """
        Returns the material_id that has the same structure as this doc as
         determined by the structure matcher. Returns None if no match.

        Args:
            doc: a JSON-like document

        Returns:
            (int) matching material_id or None
        """
        formula = doc["formula_reduced_abc"]
        sgnum = doc["spacegroup"]["number"]

        for m in self._materials.find({"formula_reduced_abc": formula, "sg_number": sgnum},
                                      {"structure": 1, "material_id": 1}):

            m_struct = Structure.from_dict(m["structure"])
            t_struct = Structure.from_dict(doc["structure"])

            sm = StructureMatcher(ltol=0.2, stol=0.3, angle_tol=5,
                                  primitive_cell=True, scale=True,
                                  attempt_supercell=False, allow_subset=False,
                                  comparator=ElementComparator())

            if sm.fit(m_struct, t_struct):
                return m["material_id"]

        return None

    def test_ordering_enumeration(self):

        # simple afm
        structure = Structure.from_file(os.path.join(ref_dir, "ordering/LaMnO3.json"))
        wf = MagneticOrderingsWF(structure)
        self.assertEqual(wf.input_origin, "afm")

        # ferrimagnetic (Cr produces net spin)
        structure = Structure.from_file(os.path.join(ref_dir, "ordering/Cr2NiO4.json"))
        wf = MagneticOrderingsWF(structure)
        self.assertEqual(wf.input_origin, "ferri_by_Cr")

        # antiferromagnetic on single magnetic site
        structure = Structure.from_file(os.path.join(ref_dir, "ordering/Cr2WO6.json"))
        wf = MagneticOrderingsWF(structure)
        self.assertEqual(wf.input_origin, "afm_by_Cr")

        # afm requiring large cell size
        # (enable for further development of workflow, too slow for CI)

        # structure = Structure.from_file(os.path.join(ref_dir, "CuO.json"))
        # wf = MagneticOrderingsWF(structure, default_magmoms={'Cu': 1.73},
        #                         transformation_kwargs={'max_cell_size': 4})
        # self.assertEqual(wf.input_origin, "afm")

        # antiferromagnetic by structural motif
        structure = Structure.from_file(os.path.join(ref_dir, "ordering/Ca3Co2O6.json"))
        wf = MagneticOrderingsWF(
            structure,
            strategies=("antiferromagnetic_by_motif",),
            # this example just misses default cut-off, so do not truncate
            truncate_by_symmetry=False,
            transformation_kwargs={"max_cell_size": 2},
        )
        self.assertEqual(wf.input_origin, "afm_by_motif_2a")

    def setUp(self):
        self.structure = Structure.from_spacegroup("Fm-3m", Lattice.cubic(3.5),
                                                   ["Ni"], [[0, 0, 0]])
        lattice = Lattice.cubic(3.010)
        frac_coords = [[0.00000, 0.00000, 0.00000],
                       [0.00000, 0.50000, 0.50000],
                       [0.50000, 0.00000, 0.50000],
                       [0.50000, 0.50000, 0.00000],
                       [0.50000, 0.00000, 0.00000],
                       [0.50000, 0.50000, 0.50000],
                       [0.00000, 0.00000, 0.50000],
                       [0.00000, 0.50000, 0.00000]]
        species = ['Mg', 'Mg', 'Mg', 'Mg', 'O', 'O', 'O', 'O']
        self.MgO = Structure(lattice, species, frac_coords)

        slabs = generate_all_slabs(self.structure, max_index=2,
                                   min_slab_size=6.0, min_vacuum_size=15.0,
                                   max_normal_search=1, center_slab=True)
        self.slab_dict = {''.join([str(i) for i in slab.miller_index]):
                              slab for slab in slabs}
        self.asf_211 = AdsorbateSiteFinder(self.slab_dict["211"])
        self.asf_100 = AdsorbateSiteFinder(self.slab_dict["100"])
        self.asf_111 = AdsorbateSiteFinder(self.slab_dict["111"])
        self.asf_110 = AdsorbateSiteFinder(self.slab_dict["110"])
        self.asf_struct = AdsorbateSiteFinder(
            Structure.from_sites(self.slab_dict["111"].sites))

 def test_unique_structure_substitutions( self ):
     # integration test
     # Create a pymatgen structure with 16 sites in a 4x4 square grid
     coords = np.array( [ [ 0.0, 0.0, 0.0 ],
                          [ 0.25, 0.0, 0.0 ],
                          [ 0.5, 0., 0.0 ],
                          [ 0.75, 0.0, 0.0 ],
                          [ 0.0, 0.25, 0.0 ],
                          [ 0.25, 0.25, 0.0 ],
                          [ 0.5, 0.25, 0.0 ],
                          [ 0.75, 0.25, 0.0 ],
                          [ 0.0, 0.5, 0.0 ],
                          [ 0.25, 0.5, 0.0 ],
                          [ 0.5, 0.5, 0.0 ],
                          [ 0.75, 0.5, 0.0 ],
                          [ 0.0, 0.75, 0.0 ],
                          [ 0.25, 0.75, 0.0 ],
                          [ 0.5, 0.75, 0.0 ],
                          [ 0.75, 0.75, 0.0 ] ] )
     atom_list = [ 'Li' ] * len( coords )
     lattice = Lattice.from_parameters( a = 3.0, b=3.0, c=3.0, alpha=90, beta=90, gamma=90 )
     parent_structure = Structure( lattice, atom_list, coords )
     parent_structure.replace( 0, 'O' ) # substitute one site with 'O'
     ns = unique_structure_substitutions( parent_structure, 'Li', { 'Na':1, 'Li':14 } )
     self.assertEqual( len( ns ), 5 )
     distances = np.array( sorted( [ s.get_distance( s.indices_from_symbol('O')[0], s.indices_from_symbol('Na')[0] ) for s in ns ] ) )
     np.testing.assert_array_almost_equal( distances, np.array( [ 0.75    ,  1.06066 ,  1.5     ,  1.677051,  2.12132 ] ) )
     np.testing.assert_array_equal( np.array( sorted( [ s.number_of_equivalent_configurations for s in ns ] ) ), np.array( [ 1, 2, 4, 4, 4 ] ) )
     np.testing.assert_array_equal( np.array( sorted( [ s.full_configuration_degeneracy for s in ns ] ) ), np.array( [ 1, 2, 4, 4, 4 ] ) )

    def setUp(self):

        with open(os.path.join(test_dir, 'LobsterCompleteDos_spin.json'), 'r') as f:
            data_spin = json.load(f)
        self.LobsterCompleteDOS_spin = LobsterCompleteDos.from_dict(data_spin)

        with open(os.path.join(test_dir, 'LobsterCompleteDos_nonspin.json'), 'r') as f:
            data_nonspin = json.load(f)
        self.LobsterCompleteDOS_nonspin = LobsterCompleteDos.from_dict(data_nonspin)

        with open(os.path.join(test_dir, 'structure_KF.json'), 'r') as f:
            data_structure = json.load(f)
        self.structure = Structure.from_dict(data_structure)

        with open(os.path.join(test_dir, 'LobsterCompleteDos_MnO.json'), 'r') as f:
            data_MnO = json.load(f)
        self.LobsterCompleteDOS_MnO = LobsterCompleteDos.from_dict(data_MnO)

        with open(os.path.join(test_dir, 'LobsterCompleteDos_MnO_nonspin.json'), 'r') as f:
            data_MnO_nonspin = json.load(f)
        self.LobsterCompleteDOS_MnO_nonspin = LobsterCompleteDos.from_dict(data_MnO_nonspin)

        with open(os.path.join(test_dir, 'structure_MnO.json'), 'r') as f:
            data_MnO = json.load(f)
        self.structure_MnO = Structure.from_dict(data_MnO)

    def test_disordered_primitive_to_ordered_supercell(self):
        sm_atoms = StructureMatcher(ltol=0.2, stol=0.3, angle_tol=5,
                                    primitive_cell=False, scale=True,
                                    attempt_supercell=True,
                                    allow_subset=True,
                                    supercell_size='num_atoms',
                                    comparator=OrderDisorderElementComparator())
        sm_sites = StructureMatcher(ltol=0.2, stol=0.3, angle_tol=5,
                                    primitive_cell=False, scale=True,
                                    attempt_supercell=True,
                                    allow_subset=True,
                                    supercell_size='num_sites',
                                    comparator=OrderDisorderElementComparator())
        lp = Lattice.orthorhombic(10, 20, 30)
        pcoords = [[0, 0, 0],
                   [0.5, 0.5, 0.5]]
        ls = Lattice.orthorhombic(20, 20, 30)
        scoords = [[0, 0, 0],
                   [0.75, 0.5, 0.5]]
        prim = Structure(lp, [{'Na': 0.5}, {'Cl': 0.5}], pcoords)
        supercell = Structure(ls, ['Na', 'Cl'], scoords)
        supercell.make_supercell([[-1, 1, 0], [0, 1, 1], [1, 0, 0]])

        self.assertFalse(sm_sites.fit(prim, supercell))
        self.assertTrue(sm_atoms.fit(prim, supercell))

        self.assertRaises(ValueError, sm_atoms.get_s2_like_s1, prim, supercell)
        self.assertEqual(len(sm_atoms.get_s2_like_s1(supercell, prim)), 4)

    def run_task(self, fw_spec):

        prev_dir = fw_spec.get('PREV_DIR', None)
        self.custom_params = self.get('custom_params', None)   

        if isinstance(self["structure"], Structure):
            s = self["structure"]
        elif isinstance(self["structure"], dict):
            s = Structure.from_dict(self["structure"])
        else:
            s = Structure.from_file(os.path.join(prev_dir, self["structure"]))


        vis = load_class("pymatgen.io.vasp.sets", self["vasp_input_set"])(
                         **self.get("input_set_params", {}))
        vis.write_input(s, ".")


        # Write Custom KPOINTS settings if necessary
        ksettings = self.custom_params.get('user_kpts_settings', None) if isinstance(
                self.custom_params, dict) else None
        if ksettings:
            style = ksettings.get('kpts_style', 'Gamma')
            kpoints = ksettings.get('kpts', [16,16,16])
            shift = ksettings.get('kpts_shift', [0,0,0])
            k = Kpoints(kpts=[kpoints], kpts_shift=shift)
            k.style = style
            k.write_file("KPOINTS")

def poscar_from_sitelist( configs, labels, sitelists, structure, subset = None ):
    """
    Uses pymatgen Structure.to() method to generates POSCAR files for a set of 
    configurations within a parent structure.
    Args:
        configs   (list):               site configurations
        labels    (list):               atom labels to output
        sitelist  (list):               list of sites in fractional coords
        structure (pymatgen Structure): Parent structure
        subset    (Optional [list]):    list of atom indices to output 
    """

    if subset:
        species_clean = [ spec for i,spec in enumerate( structure.species ) if i not in subset ]
        species_config = [ spec for i,spec in enumerate( structure.species ) if i in subset ]
        frac_coords_clean = [ coord for i, coord in enumerate( structure.frac_coords ) if i not in subset ]
        clean_structure = Structure( structure.lattice, species_clean, frac_coords_clean )

    else:
        clean_structure = Structure( structure.lattice, [], [] )
        species_config = structure.species

    for idx, config in enumerate( configs, start=1 ):
       structure_config = clean_structure.copy()
       for label in labels:
           for pos in config.position( label ):
               for sitelist in sitelists:
                   structure_config.append( species_config[ pos ], sitelist[ pos ] )
       structure_config.to( filename="POSCAR_{}.vasp".format( idx ) )

    def _match_material(self, taskdoc):
        """
        Returns the material_id that has the same structure as this task as
         determined by the structure matcher. Returns None if no match.

        Args:
            taskdoc (dict): a JSON-like task document

        Returns:
            (int) matching material_id or None
        """
        formula = taskdoc["formula_reduced_abc"]

        if "parent_structure" in taskdoc:  # this is used to intentionally combine multiple data w/same formula but slightly different structure, e.g. from an ordering scheme
            t_struct = Structure.from_dict(taskdoc["parent_structure"]["structure"])
            q = {"formula_reduced_abc": formula, "parent_structure.spacegroup.number": taskdoc["parent_structure"]["spacegroup"]["number"]}
        else:
            sgnum = taskdoc["output"]["spacegroup"]["number"]
            t_struct = Structure.from_dict(taskdoc["output"]["structure"])
            q = {"formula_reduced_abc": formula, "sg_number": sgnum}

        for m in self._materials.find(q, {"parent_structure": 1, "structure": 1, "material_id": 1}):
            s_dict = m["parent_structure"]["structure"] if "parent_structure" in m else m["structure"]
            m_struct = Structure.from_dict(s_dict)
            sm = StructureMatcher(ltol=0.2, stol=0.3, angle_tol=5,
                                  primitive_cell=True, scale=True,
                                  attempt_supercell=False, allow_subset=False,
                                  comparator=ElementComparator())

            if sm.fit(m_struct, t_struct):
                return m["material_id"]

        return None

    def test_structure(self):
        quartz = self.quartz.structure
        np.testing.assert_array_equal(quartz.lattice.matrix,
                                      [[4.913400, 0, 0],
                                       [-2.456700, 4.255129, 0],
                                       [0, 0, 5.405200]])
        self.assertEqual(quartz.formula, "Si3 O6")
        self.assertNotIn("molecule-ID", self.quartz.atoms.columns)

        ethane = self.ethane.structure
        np.testing.assert_array_equal(ethane.lattice.matrix,
                                      np.diag([10.0] * 3))
        lbounds = np.array(self.ethane.box.bounds)[:, 0]
        coords = self.ethane.atoms[["x", "y", "z"]].values - lbounds
        np.testing.assert_array_equal(ethane.cart_coords, coords)
        np.testing.assert_array_equal(ethane.site_properties["charge"],
                                      self.ethane.atoms["q"])
        tatb = self.tatb.structure
        frac_coords = tatb.frac_coords[381]
        real_frac_coords = frac_coords - np.floor(frac_coords)
        np.testing.assert_array_almost_equal(real_frac_coords,
                                             [0.01553397,
                                              0.71487872,
                                              0.14134139])

        co = Structure.from_spacegroup(194,
                                       Lattice.hexagonal(2.50078, 4.03333),
                                       ["Co"], [[1/3, 2/3, 1/4]])
        ld_co = LammpsData.from_structure(co)
        self.assertEqual(ld_co.structure.composition.reduced_formula, "Co")
        ni = Structure.from_spacegroup(225, Lattice.cubic(3.50804),
                                       ["Ni"], [[0, 0, 0]])
        ld_ni = LammpsData.from_structure(ni)
        self.assertEqual(ld_ni.structure.composition.reduced_formula, "Ni")

    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 setUp(self):
     c1 = [[0.5] * 3, [0.9] * 3]
     c2 = [[0.5] * 3, [0.9, 0.1, 0.1]]
     s1 = Structure(Lattice.cubic(5), ['Si', 'Si'], c1)
     s2 = Structure(Lattice.cubic(5), ['Si', 'Si'], c2)
     structs = []
     for s in s1.interpolate(s2, 3, pbc=True):
         structs.append(Structure.from_sites(s.sites, to_unit_cell=True))
     self.structures = structs
     self.vis = MITNEBSet(self.structures)

    def test_apply_transformation_mult(self):
        # Test returning multiple structures from each transformation.
        disord = Structure(np.eye(3) * 4.209, [{"Cs+": 0.5, "K+": 0.5}, "Cl-"], [[0, 0, 0], [0.5, 0.5, 0.5]])
        disord.make_supercell([2, 2, 1])

        tl = [EnumerateStructureTransformation(), OrderDisorderedStructureTransformation()]
        t = SuperTransformation(tl, nstructures_per_trans=10)
        self.assertEqual(len(t.apply_transformation(disord, return_ranked_list=20)), 8)
        t = SuperTransformation(tl)
        self.assertEqual(len(t.apply_transformation(disord, return_ranked_list=20)), 2)

 def run(self):
     if os.path.exists("CONTCAR"):
         structure = Structure.from_file("CONTCAR")
     elif (not self.contcar_only) and os.path.exists("POSCAR"):
         structure = Structure.from_file("POSCAR")
     else:
         raise RuntimeError("No CONTCAR/POSCAR detected to generate input!")
     modname, classname = self.input_set.rsplit(".", 1)
     mod = __import__(modname, globals(), locals(), [classname], 0)
     vis = getattr(mod, classname)(structure, **self.kwargs)
     vis.write_input(".")