enh: enabled sorting of geometries before writing dcc-2019-jun

Also added that as a staticmethod to vasp objects.

Signed-off-by: Nick Papior <[email protected]>

Author: zerothi

sisl/io/vasp/car.py

@@ -23,7 +23,7 @@ def _setup(self, *args, **kwargs):
         self._scale = 1.
 
     @sile_fh_open()
-    def write_geometry(self, geometry, dynamic=None):
+    def write_geometry(self, geometry, dynamic=None, sort=False):
         r""" Writes the geometry to the contained file
 
         Parameters
@@ -34,6 +34,9 @@ def write_geometry(self, geometry, dynamic=None):
            define which atoms are dynamic in the VASP run (default is True,
            which means all atoms are dynamic).
            If None, the resulting file will not contain any dynamic flags
+        sort : bool, optional
+           before writing `geometry` first re-order species to
+           have species in consecutive blocks (see `geometry_sort`)
 
         Examples
         --------
@@ -42,7 +45,14 @@ def write_geometry(self, geometry, dynamic=None):
         >>> geom.write(car) # regular car without Selective Dynamics
         >>> geom.write(car, dynamic=False) # fix all atoms
         >>> geom.write(car, dynamic=[False, (True, False, True)]) # fix 1st and y coordinate of 2nd
+
+        See Also
+        --------
+        geometry_sort: method used to sort atoms in geometry
         """
+        if sort:
+            geometry = self.geometry_sort(geometry)
+
         # Check that we can write to the file
         sile_raise_write(self)
 

sisl/io/vasp/sile.py

@@ -1,19 +1,58 @@
 """
 Define a common VASP Sile
 """
+import sisl._array as _a
 
 from ..sile import Sile, SileCDF, SileBin
 
 __all__ = ['SileVASP', 'SileCDFVASP', 'SileBinVASP']
 
 
+def _geometry_sort(geometry):
+    r""" Order atoms in geometry according to species such that all of one specie is consecutive
+
+    When creating VASP input files (`poscarSileVASP` for instance) the equivalent
+    ``POTCAR`` file needs to contain the pseudos for each specie as they are provided
+    in blocks.
+
+    I.e. for a geometry like this:
+    .. code::
+
+        [Atom(6), Atom(4), Atom(6)]
+
+    the resulting ``POTCAR`` needs to contain the pseudo for Carbon twice.
+
+    This method will re-order atoms according to the species"
+
+    Parameters
+    ----------
+    geometry : Geometry
+       geometry to be re-ordered
+
+    Returns
+    -------
+    geometry: reordered geometry
+    """
+    na = len(geometry)
+    idx = _a.emptyi(na)
+
+    ia = 0
+    for _, idx_s in geometry.atoms.iter(species=True):
+        idx[ia:ia + len(idx_s)] = idx_s
+        ia += len(idx_s)
+
+    assert ia == na
+
+    return geometry.sub(idx)
+
+
 class SileVASP(Sile):
-    pass
+    geometry_sort = staticmethod(_geometry_sort)
 
 
 class SileCDFVASP(SileCDF):
-    pass
+    geometry_sort = staticmethod(_geometry_sort)
 
 
 class SileBinVASP(SileBin):
-    pass
+    geometry_sort = staticmethod(_geometry_sort)

sisl/io/vasp/tests/test_car.py

@@ -27,6 +27,26 @@ def test_geometry_car_mixed(sisl_tmp):
     assert carSileVASP(f).read_geometry() == geom
 
 
+def test_geometry_car_sort(sisl_tmp):
+    f = sisl_tmp('test_sort.POSCAR', _dir)
+
+    atoms = [Atom[1],
+             Atom[2],
+             Atom[2],
+             Atom[1],
+             Atom[1],
+             Atom[2],
+             Atom[3]]
+    xyz = np.random.rand(len(atoms), 3)
+    geom = Geometry(xyz, atoms, 100)
+
+    geom.write(carSileVASP(f, 'w'), sort=True)
+
+    assert carSileVASP(f).read_geometry() != geom
+    geom = carSileVASP(f).geometry_sort(geom)
+    assert carSileVASP(f).read_geometry() == geom
+
+
 def test_geometry_car_allsame(sisl_tmp):
     f = sisl_tmp('test_read_write.POSCAR', _dir)