Python structure_modifier.StructureEditor类代码示例

 def apply_transformation(self, structure):
     editor = StructureEditor(structure)
     for i, sp in enumerate(self._species):
         editor.insert_site(i, sp, self._coords[i],
                            coords_are_cartesian=self._cartesian,
                            validate_proximity=self._validate_proximity)
     return editor.modified_structure.get_sorted_structure()

def _get_host(structure, species_to_remove):
    if species_to_remove:
        editor = StructureEditor(structure)
        editor.remove_species(species_to_remove)
        return editor.modified_structure
    else:
        return structure

    def complete_ordering(self, structure, num_remove_dict):
        self.logger.debug("Performing complete ordering...")
        all_structures = []
        from pymatgen.symmetry.finder import SymmetryFinder
        symprec = 0.2
        s = SymmetryFinder(structure, symprec=symprec)
        self.logger.debug("Symmetry of structure is determined to be {}."
                          .format(s.get_spacegroup_symbol()))
        sg = s.get_spacegroup()
        tested_sites = []
        starttime = time.time()
        self.logger.debug("Performing initial ewald sum...")
        ewaldsum = EwaldSummation(structure)
        self.logger.debug("Ewald sum took {} seconds."
                          .format(time.time() - starttime))
        starttime = time.time()

        allcombis = []
        for ind, num in num_remove_dict.items():
            allcombis.append(itertools.combinations(ind, num))

        count = 0
        for allindices in itertools.product(*allcombis):
            sites_to_remove = []
            indices_list = []
            for indices in allindices:
                sites_to_remove.extend([structure[i] for i in indices])
                indices_list.extend(indices)

            mod = StructureEditor(structure)
            mod.delete_sites(indices_list)
            s_new = mod.modified_structure
            energy = ewaldsum.compute_partial_energy(indices_list)
            already_tested = False
            for i, tsites in enumerate(tested_sites):
                tenergy = all_structures[i]["energy"]
                if abs((energy - tenergy) / len(s_new)) < 1e-5 and \
                        sg.are_symmetrically_equivalent(sites_to_remove, tsites,
                                                        symm_prec=symprec):
                    already_tested = True

            if not already_tested:
                tested_sites.append(sites_to_remove)
                all_structures.append({"structure": s_new, "energy": energy})

            count += 1
            if count % 10 == 0:
                timenow = time.time()
                self.logger.debug("{} structures, {:.2f} seconds."
                                  .format(count, timenow - starttime))
                self.logger.debug("Average time per combi = {} seconds"
                                  .format((timenow - starttime) / count))
                self.logger.debug("{} symmetrically distinct structures found."
                                  .format(len(all_structures)))

        self.logger.debug("Total symmetrically distinct structures found = {}"
                          .format(len(all_structures)))
        all_structures = sorted(all_structures, key=lambda s: s["energy"])
        return all_structures

 def apply_transformation(self, structure):
     species_map = {}
     for k, v in self._species_map.items():
         if isinstance(v, dict):
             value = {smart_element_or_specie(x): y for x, y in v.items()}
         else:
             value = smart_element_or_specie(v)
         species_map[smart_element_or_specie(k)] = value
     editor = StructureEditor(structure)
     editor.replace_species(species_map)
     return editor.modified_structure

    def test_add_site_property(self):
        self.modifier.add_site_property("charge", [4.1, 5])
        s = self.modifier.modified_structure
        self.assertEqual(s[0].charge, 4.1)
        self.assertEqual(s[1].charge, 5)

        #test adding multiple properties.
        mod2 = StructureEditor(s)
        mod2.add_site_property("magmom", [3, 2])
        s = mod2.modified_structure
        self.assertEqual(s[0].charge, 4.1)
        self.assertEqual(s[0].magmom, 3)

 def setUp(self):
     p = Poscar.from_file(os.path.join(test_dir, 'POSCAR'))
     self.structure = p.struct
     self.sg = SymmetryFinder(self.structure, 0.001)
     parser = CifParser(os.path.join(test_dir, 'Li10GeP2S12.cif'))
     self.disordered_structure = parser.get_structures()[0]
     self.disordered_sg = SymmetryFinder(self.disordered_structure, 0.001)
     s = p.struct
     editor = StructureEditor(p.struct)
     site = s[0]
     editor.delete_site(0)
     editor.append_site(site.species_and_occu, site.frac_coords)
     self.sg3 = SymmetryFinder(editor.modified_structure, 0.001)

    def test_init(self):
        fitter = StructureFitter(self.b, self.a)
        self.assertTrue(fitter.mapping_op != None, "No fit found!")

        #Now to try with rotated structure
        op = SymmOp.from_axis_angle_and_translation([0, 0, 1], 30, False, np.array([0, 0, 1]))
        editor = StructureEditor(self.a)
        editor.apply_operation(op)
        fitter = StructureFitter(self.b, editor.modified_structure)

        self.assertTrue(fitter.mapping_op != None, "No fit found!")

        #test with a supercell
        mod = SupercellMaker(self.a, scaling_matrix=[[2, 0, 0], [0, 1, 0], [0, 0, 1]])
        a_super = mod.modified_structure
        fitter = StructureFitter(self.b, a_super)
        self.assertTrue(fitter.mapping_op != None, "No fit found!")

        # Test with a structure with a translated point

        editor = StructureEditor(self.a)
        site = self.a[0]
        editor.delete_site(0)
        trans = np.random.randint(0, 1000, 3)
        editor.insert_site(0, site.species_and_occu, site.frac_coords + trans, False, False)
        fitter = StructureFitter(self.b, editor.modified_structure)
        self.assertTrue(fitter.mapping_op != None, "No fit found for translation {}!".format(trans))

        parser = CifParser(os.path.join(test_dir, "FePO4a.cif"))
        a = parser.get_structures()[0]
        parser = CifParser(os.path.join(test_dir, "FePO4b.cif"))
        b = parser.get_structures()[0]
        fitter = StructureFitter(b, a)
        self.assertTrue(fitter.mapping_op != None, "No fit found!")

    def fast_ordering(self, structure, num_remove_dict, num_to_return=1):
        """
        This method uses the matrix form of ewaldsum to calculate the ewald
        sums of the potential structures. This is on the order of 4 orders of
        magnitude faster when there are large numbers of permutations to
        consider. There are further optimizations possible (doing a smarter
        search of permutations for example), but this wont make a difference
        until the number of permutations is on the order of 30,000.
        """
        self.logger.debug("Performing fast ordering")
        starttime = time.time()
        self.logger.debug("Performing initial ewald sum...")

        ewaldmatrix = EwaldSummation(structure).total_energy_matrix
        self.logger.debug("Ewald sum took {} seconds."
                          .format(time.time() - starttime))
        starttime = time.time()
        m_list = []
        for indices, num in num_remove_dict.items():
            m_list.append([0, num, list(indices), None])

        self.logger.debug("Calling EwaldMinimizer...")
        minimizer = EwaldMinimizer(ewaldmatrix, m_list, num_to_return,
                                   PartialRemoveSitesTransformation.ALGO_FAST)
        self.logger.debug("Minimizing Ewald took {} seconds."
                          .format(time.time() - starttime))

        all_structures = []

        lowest_energy = minimizer.output_lists[0][0]
        num_atoms = sum(structure.composition.values())

        for output in minimizer.output_lists:
            se = StructureEditor(structure)
            del_indices = []

            for manipulation in output[1]:
                if manipulation[1] is None:
                    del_indices.append(manipulation[0])
                else:
                    se.replace_site(manipulation[0], manipulation[1])
            se.delete_sites(del_indices)
            struct = se.modified_structure.get_sorted_structure()
            all_structures.append({"energy": output[0],
                                   "energy_above_minimum": (output[0]
                                                            - lowest_energy)
                                   / num_atoms,
                                   "structure": struct})

        return all_structures

 def enumerate_ordering(self, structure):
     # Generate the disordered structure first.
     editor = StructureEditor(structure)
     for indices, fraction in zip(self._indices, self._fractions):
         for ind in indices:
             new_sp = {sp: occu * fraction
                       for sp, occu
                       in structure[ind].species_and_occu.items()}
             editor.replace_site(ind, new_sp)
     mod_s = editor.modified_structure
     # Perform enumeration
     from pymatgen.transformations.advanced_transformations import \
         EnumerateStructureTransformation
     trans = EnumerateStructureTransformation()
     return trans.apply_transformation(mod_s, 10000)

 def test_remove_oxidation_states(self):
     co_elem = Element("Co")
     o_elem = Element("O")
     co_specie = Specie("Co", 2)
     o_specie = Specie("O", -2)
     coords = list()
     coords.append([0, 0, 0])
     coords.append([0.75, 0.5, 0.75])
     lattice = Lattice.cubic(10)
     s_elem = Structure(lattice, [co_elem, o_elem], coords)
     s_specie = Structure(lattice, [co_specie, o_specie], coords)
     mod = StructureEditor(s_specie)
     mod.remove_oxidation_states()
     mod_s = mod.modified_structure
     self.assertEqual(s_elem, mod_s, "Oxidation state remover failed")

    def test_fit(self):
        """
        Take two known matched structures
            1) Ensure match
            2) Ensure match after translation and rotations
            3) Ensure no-match after large site translation
            4) Ensure match after site shuffling
            """
        sm = StructureMatcher()

        self.assertTrue(sm.fit(self.struct_list[0], self.struct_list[1]))

        # Test rotational/translational invariance
        op = SymmOp.from_axis_angle_and_translation([0, 0, 1], 30, False,
                                                    np.array([0.4, 0.7, 0.9]))
        editor = StructureEditor(self.struct_list[1])
        editor.apply_operation(op)
        self.assertTrue(sm.fit(self.struct_list[0], editor.modified_structure))

        #Test failure under large atomic translation
        editor.translate_sites([0], [.4, .4, .2], frac_coords=True)
        self.assertFalse(sm.fit(self.struct_list[0],
                                editor.modified_structure))

        editor.translate_sites([0], [-.4, -.4, -.2], frac_coords=True)
        # random.shuffle(editor._sites)
        self.assertTrue(sm.fit(self.struct_list[0], editor.modified_structure))
        #Test FrameworkComporator
        sm2 = StructureMatcher(comparator=FrameworkComparator())
        lfp = read_structure(os.path.join(test_dir, "LiFePO4.cif"))
        nfp = read_structure(os.path.join(test_dir, "NaFePO4.cif"))
        self.assertTrue(sm2.fit(lfp, nfp))
        self.assertFalse(sm.fit(lfp, nfp))

        #Test anonymous fit.
        self.assertEqual(sm.fit_anonymous(lfp, nfp),
                         {Composition("Li"): Composition("Na")})
        self.assertAlmostEqual(sm.get_minimax_rms_anonymous(lfp, nfp)[0],
                               0.096084154118549828)

        #Test partial occupancies.
        s1 = Structure([[3, 0, 0], [0, 3, 0], [0, 0, 3]],
                       [{"Fe": 0.5}, {"Fe": 0.5}, {"Fe": 0.5}, {"Fe": 0.5}],
                       [[0, 0, 0], [0.25, 0.25, 0.25],
                        [0.5, 0.5, 0.5], [0.75, 0.75, 0.75]])
        s2 = Structure([[3, 0, 0], [0, 3, 0], [0, 0, 3]],
                       [{"Fe": 0.25}, {"Fe": 0.5}, {"Fe": 0.5}, {"Fe": 0.75}],
                       [[0, 0, 0], [0.25, 0.25, 0.25],
                        [0.5, 0.5, 0.5], [0.75, 0.75, 0.75]])
        self.assertFalse(sm.fit(s1, s2))
        self.assertFalse(sm.fit(s2, s1))
        s2 = Structure([[3, 0, 0], [0, 3, 0], [0, 0, 3]],
                       [{"Fe": 0.25}, {"Fe": 0.25}, {"Fe": 0.25},
                        {"Fe": 0.25}],
                       [[0, 0, 0], [0.25, 0.25, 0.25],
                        [0.5, 0.5, 0.5], [0.75, 0.75, 0.75]])
        self.assertEqual(sm.fit_anonymous(s1, s2),
                         {Composition("Fe0.5"): Composition("Fe0.25")})

        self.assertAlmostEqual(sm.get_minimax_rms_anonymous(s1, s2)[0], 0)

 def setUp(self):
     self.si = Element("Si")
     self.fe = Element("Fe")
     self.ge = Element("Ge")
     coords = list()
     coords.append(np.array([0, 0, 0]))
     coords.append(np.array([0.75, 0.5, 0.75]))
     lattice = Lattice.cubic(10)
     s = Structure(lattice, ["Si", "Fe"], coords)
     self.modifier = StructureEditor(s)

    def get_oxi_state_decorated_structure(self, structure):
        """
        Get an oxidation state decorated structure. This currently works only
        for ordered structures only.

        Args:
            structure:
                Structure to analyze

        Returns:
            A modified structure that is oxidation state decorated.

        Raises:
            A ValueError is the valences cannot be determined.
        """
        valences = self.get_valences(structure)
        editor = StructureEditor(structure)
        editor.add_oxidation_state_by_site(valences)
        return editor.modified_structure

    def best_first_ordering(self, structure, num_remove_dict):
        self.logger.debug("Performing best first ordering")
        starttime = time.time()
        self.logger.debug("Performing initial ewald sum...")
        ewaldsum = EwaldSummation(structure)
        self.logger.debug("Ewald sum took {} seconds."
                          .format(time.time() - starttime))
        starttime = time.time()

        ematrix = ewaldsum.total_energy_matrix
        to_delete = []

        totalremovals = sum(num_remove_dict.values())
        removed = {k: 0 for k in num_remove_dict.keys()}
        for i in xrange(totalremovals):
            maxindex = None
            maxe = float("-inf")
            maxindices = None
            for indices in num_remove_dict.keys():
                if removed[indices] < num_remove_dict[indices]:
                    for ind in indices:
                        if ind not in to_delete:
                            energy = sum(ematrix[:, ind]) + \
                                sum(ematrix[:, ind]) - ematrix[ind, ind]
                            if energy > maxe:
                                maxindex = ind
                                maxe = energy
                                maxindices = indices
            removed[maxindices] += 1
            to_delete.append(maxindex)
            ematrix[:, maxindex] = 0
            ematrix[maxindex, :] = 0
        mod = StructureEditor(structure)
        mod.delete_sites(to_delete)
        self.logger.debug("Minimizing Ewald took {} seconds."
                          .format(time.time() - starttime))
        return [{"energy": sum(sum(ematrix)),
                 "structure": mod.modified_structure.get_sorted_structure()}]

 def test_add_oxidation_states(self):
     si = Element("Si")
     fe = Element("Fe")
     coords = list()
     coords.append([0, 0, 0])
     coords.append([0.75, 0.5, 0.75])
     lattice = Lattice.cubic(10)
     s = Structure(lattice, [si, fe], coords)
     oxidation_states = {"Fe": 2, "Si": -4}
     mod = StructureEditor(s)
     mod.add_oxidation_state_by_element(oxidation_states)
     mod_s = mod.modified_structure
     for site in mod_s:
         for k in site.species_and_occu.keys():
             self.assertEqual(k.oxi_state, oxidation_states[k.symbol],
                              "Wrong oxidation state assigned!")
     oxidation_states = {"Fe": 2}
     self.assertRaises(ValueError, mod.add_oxidation_state_by_element,
                       oxidation_states)
     mod.add_oxidation_state_by_site([2, -4])
     mod_s = mod.modified_structure
     self.assertEqual(mod_s[0].specie.oxi_state, 2)
     self.assertRaises(ValueError, mod.add_oxidation_state_by_site,
                       [1])