enh: enabled writing without selective dynamics (#190)

* enh: enabled writing without selective dynamics

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

* 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]>

* bug: fixed sort and dynamic together

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

* enh: finalized TF suggestions and clarified intents

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

Author: zerothi

sisl/io/vasp/car.py

@@ -23,27 +23,42 @@ def _setup(self, *args, **kwargs):
         self._scale = 1.
 
     @sile_fh_open()
-    def write_geometry(self, geometry, dynamic=True):
+    def write_geometry(self, geometry, dynamic=True, group_species=False):
         r""" Writes the geometry to the contained file
 
         Parameters
         ----------
         geometry : Geometry
            geometry to be written to the file
-        dynamic : bool or list, optional
+        dynamic : None, bool or list, optional
            define which atoms are dynamic in the VASP run (default is True,
-           which means all atoms are dynamic)
+           which means all atoms are dynamic).
+           If None, the resulting file will not contain any dynamic flags
+        group_species: bool, optional
+           before writing `geometry` first re-order species to
+           have species in consecutive blocks (see `geometry_group`)
 
         Examples
         --------
         >>> car = carSileVASP('POSCAR', 'w')
         >>> geom = geom.graphene()
+        >>> 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_group: method used to group atoms together according to their species
         """
         # Check that we can write to the file
         sile_raise_write(self)
 
+        if group_species:
+            geometry, idx = self.geometry_group(geometry, ret_index=True)
+        else:
+            # small hack to allow dynamic
+            idx = _a.arangei(len(geometry))
+
         # LABEL
         self._write('sisl output\n')
 
@@ -52,10 +67,8 @@ def write_geometry(self, geometry, dynamic=True):
 
         # Write unit-cell
         fmt = ('   ' + '{:18.9f}' * 3) + '\n'
-        tmp = np.zeros([3], np.float64)
         for i in range(3):
-            tmp[:3] = geometry.cell[i, :]
-            self._write(fmt.format(*tmp))
+            self._write(fmt.format(*geometry.cell[i]))
 
         # Figure out how many species
         pt = PeriodicTable()
@@ -78,21 +91,26 @@ def write_geometry(self, geometry, dynamic=True):
         self._write(fmt.format(*s))
         fmt = ' {:d}' * len(d) + '\n'
         self._write(fmt.format(*d))
-        self._write('Selective dynamics\n')
+        if dynamic is None:
+            # We write in direct mode
+            dynamic = [None] * len(geometry)
+            def todyn(fix):
+                return '\n'
+        else:
+            self._write('Selective dynamics\n')
+            b2s = {True: 'T', False: 'F'}
+            def todyn(fix):
+                if isinstance(fix, bool):
+                    return ' {0} {0} {0}\n'.format(b2s[fix])
+                return ' {} {} {}\n'.format(b2s[fix[0]], b2s[fix[1]], b2s[fix[2]])
         self._write('Cartesian\n')
 
         if isinstance(dynamic, bool):
             dynamic = [dynamic] * len(geometry)
 
-        b2s = {True: 'T', False: 'F'}
-        def todyn(fix):
-            if isinstance(fix, bool):
-                return '{0} {0} {0}\n'.format(b2s[fix])
-            return '{} {} {}\n'.format(b2s[fix[0]], b2s[fix[1]], b2s[fix[2]])
-
-        fmt = '{:18.9f} ' * 3
+        fmt = '{:18.9f}' * 3
         for ia in geometry:
-            self._write(fmt.format(*geometry.xyz[ia, :]) + todyn(dynamic[ia]))
+            self._write(fmt.format(*geometry.xyz[ia, :]) + todyn(dynamic[idx[ia]]))
 
     @sile_fh_open(True)
     def read_supercell(self):
@@ -115,12 +133,13 @@ def read_supercell(self):
     def read_geometry(self, ret_dynamic=False):
         r""" Returns Geometry object from the CONTCAR/POSCAR file
 
-        Possibly also return the dynamics (if present)
+        Possibly also return the dynamics (if present).
 
         Parameters
         ----------
         ret_dynamic : bool, optional
-           also read selective dynamics (if present), if not, a list of True will be returned
+           also return selective dynamics (if present), if not, None will
+           be returned.
         """
         sc = self.read_supercell()
 
@@ -148,25 +167,26 @@ def read_geometry(self, ret_dynamic=False):
                 for spec, nsp in zip(species, species_count)
                 for i in range(nsp)]
 
-        # Read whether this is selective or direct
-        # Currently direct is not used
+        # Number of atoms
+        na = len(atom)
+
+        # check whether this is Selective Dynamics
         opt = self.readline()
-        dynamics = False
         if opt[0] in 'Ss':
             dynamics = True
+            # pre-create the dynamic list
+            dynamic = np.empty([na, 3], dtype=np.bool_)
             opt = self.readline()
+        else:
+            dynamics = False
+            dynamic = None
 
         # Check whether this is in fractional or direct
-        # coordinates
+        # coordinates (Direct == fractional)
         cart = False
         if opt[0] in 'CcKk':
             cart = True
 
-        # Number of atoms
-        na = len(atom)
-        # pre-create the dynamic list
-        dynamic = [[False] * 3] * na
-
         xyz = _a.emptyd([na, 3])
         for ia in range(na):
             line = self.readline().split()
@@ -183,7 +203,7 @@ def read_geometry(self, ret_dynamic=False):
         # The POT/CONT-CAR does not contain information on the atomic species
         geom = Geometry(xyz=xyz, atom=atom, sc=sc)
         if ret_dynamic:
-            return geom, np.array(dynamic, dtype=np.bool_)
+            return geom, dynamic
         return geom
 
     def ArgumentParser(self, p=None, *args, **kwargs):

sisl/io/vasp/sile.py

@@ -1,19 +1,62 @@
 """
 Define a common VASP Sile
 """
+import sisl._array as _a
 
 from ..sile import Sile, SileCDF, SileBin
 
 __all__ = ['SileVASP', 'SileCDFVASP', 'SileBinVASP']
 
 
+def _geometry_group(geometry, ret_index=False):
+    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
+    ret_index : bool, optional
+       return sorted indices
+
+    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
+
+    if ret_index:
+        return geometry.sub(idx), idx
+    return geometry.sub(idx)
+
+
 class SileVASP(Sile):
-    pass
+    geometry_group = staticmethod(_geometry_group)
 
 
 class SileCDFVASP(SileCDF):
-    pass
+    geometry_group = staticmethod(_geometry_group)
 
 
 class SileBinVASP(SileBin):
-    pass
+    geometry_group = staticmethod(_geometry_group)

sisl/io/vasp/tests/test_car.py

@@ -27,15 +27,23 @@ def test_geometry_car_mixed(sisl_tmp):
     assert carSileVASP(f).read_geometry() == geom
 
 
-def test_geometry_car_allsame(sisl_tmp):
-    f = sisl_tmp('test_read_write.POSCAR', _dir)
+def test_geometry_car_group(sisl_tmp):
+    f = sisl_tmp('test_sort.POSCAR', _dir)
 
-    atoms = Atom[1]
-    xyz = np.random.rand(10, 3)
+    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'))
+    geom.write(carSileVASP(f, 'w'), group_species=True)
 
+    assert carSileVASP(f).read_geometry() != geom
+    geom = carSileVASP(f).geometry_group(geom)
     assert carSileVASP(f).read_geometry() == geom
 
 
@@ -60,14 +68,19 @@ def test_geometry_car_dynamic(sisl_tmp):
 
     read = carSileVASP(f)
 
-    geom.write(carSileVASP(f, 'w'))
+    # no dynamic (direct geometry)
+    geom.write(carSileVASP(f, 'w'), dynamic=None)
     g, dyn = read.read_geometry(ret_dynamic=True)
-    assert np.all(dyn)
+    assert dyn is None
 
     geom.write(carSileVASP(f, 'w'), dynamic=False)
     g, dyn = read.read_geometry(ret_dynamic=True)
     assert not np.any(dyn)
 
+    geom.write(carSileVASP(f, 'w'), dynamic=True)
+    g, dyn = read.read_geometry(ret_dynamic=True)
+    assert np.all(dyn)
+
     dynamic = [False] * len(geom)
     dynamic[0] = [True, False, True]
     geom.write(carSileVASP(f, 'w'), dynamic=dynamic)