Add sort_structure option to Poscar. Fixes #1633.

materialsproject · Oct 18, 2019 · b259327 · b259327
1 parent 8ccaa02
commit b259327
Showing 1 changed file with 30 additions and 18 deletions.
diff --git a/pymatgen/io/vasp/inputs.py b/pymatgen/io/vasp/inputs.py
@@ -96,24 +96,34 @@ class Poscar(MSONable):
         to -1 (MB hasn"t been performed).
     """
 
-    def __init__(self, structure, comment=None, selective_dynamics=None,
-                 true_names=True, velocities=None, predictor_corrector=None,
-                 predictor_corrector_preamble=None):
-        """
-        Args:
-            structure (Structure):  Structure object.
-            comment (str): Optional comment line for POSCAR. Defaults to unit
-                cell formula of structure. Defaults to None.
-            selective_dynamics (Nx3 array): bool values for selective dynamics,
-                where N is number of sites. Defaults to None.
-            true_names (bool): Set to False is the names in the POSCAR are not
-                well-defined and ambiguous. This situation arises commonly in
-                vasp < 5 where the POSCAR sometimes does not contain element
-                symbols. Defaults to True.
-            velocities (Nx3 array): Velocities for the POSCAR. Typically parsed
-                in MD runs or can be used to initialize velocities.
-            predictor_corrector (Nx3 array): Predictor corrector for the POSCAR.
-                Typically parsed in MD runs.
+    def __init__(self,
+                 structure: Structure,
+                 comment: str = None,
+                 selective_dynamics=None,
+                 true_names: bool = True,
+                 velocities=None,
+                 predictor_corrector=None,
+                 predictor_corrector_preamble=None,
+                 sort_structure: bool = False):
+        """
+
+        :param structure: Structure object.
+        :param comment: Optional comment line for POSCAR. Defaults to unit
+            cell formula of structure. Defaults to None.
+        :param selective_dynamics: bool values for selective dynamics,
+            where N is number of sites. Defaults to None.
+        :param true_names: Set to False if the names in the POSCAR are not
+            well-defined and ambiguous. This situation arises commonly in
+            vasp < 5 where the POSCAR sometimes does not contain element
+            symbols. Defaults to True.
+        :param velocities: Velocities for the POSCAR. Typically parsed
+            in MD runs or can be used to initialize velocities.
+        :param predictor_corrector: Predictor corrector for the POSCAR.
+            Typically parsed in MD runs.
+        :param predictor_corrector_preamble: Preamble to the predictor
+            corrector.
+        :param sort_structure: Whether to sort structure. Useful if species
+            are not grouped properly together.
         """
         if structure.is_ordered:
             site_properties = {}
@@ -125,6 +135,8 @@ def __init__(self, structure, comment=None, selective_dynamics=None,
                 site_properties["predictor_corrector"] = predictor_corrector
             structure = Structure.from_sites(structure)
             self.structure = structure.copy(site_properties=site_properties)
+            if sort_structure:
+                self.structure = self.structure.get_sorted_structure()
             self.true_names = true_names
             self.comment = structure.formula if comment is None else comment
             self.predictor_corrector_preamble = predictor_corrector_preamble