enh: improved PR by Thomas now one can query/set fixed for geometries

Now one can write/read the fixed status of atoms.
When using lists for `fixed` argument one can easily define a simple
scheme to limit some atoms vs. others.

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

Author: zerothi

CHANGELOG

@@ -1,6 +1,8 @@
 0.9.9
 =====
 
+- Enabled VASP *CAR files to write/read dynamic specifications
+
 - Enabled xarray.DataArray returning from BrillouinZone objects #182
 
 - Several improvements to outSileSiesta.read_scf (thanks to Pol Febrer)

sisl/io/vasp/car.py

@@ -5,6 +5,7 @@
 from ..sile import *
 
 # Import the geometry object
+import sisl._array as _a
 from sisl.messages import warn
 from sisl import Geometry, PeriodicTable, Atom, SuperCell
 
@@ -22,13 +23,22 @@ def _setup(self, *args, **kwargs):
         self._scale = 1.
 
     @sile_fh_open()
-    def write_geometry(self, geometry):
-        """ Writes the geometry to the contained file
+    def write_geometry(self, geometry, fixed=False):
+        r""" Writes the geometry to the contained file
 
         Parameters
         ----------
         geometry : Geometry
            geometry to be written to the file
+        fixed : bool or list, optional
+           define which atoms to be fixed in the VASP run
+
+        Examples
+        --------
+        >>> car = carSileVASP('POSCAR', 'w')
+        >>> geom = geom.graphene()
+        >>> geom.write(car, fixed=False) # fix none
+        >>> geom.write(car, fixed=[True, (False, True, False)]) # fix 1st and y coordinate of 2nd
         """
         # Check that we can write to the file
         sile_raise_write(self)
@@ -70,9 +80,18 @@ def write_geometry(self, geometry):
         self._write('Selective dynamics\n')
         self._write('Cartesian\n')
 
-        fmt = '{:18.9f}' * 3 + ' F F F\n'
+        if isinstance(fixed, bool):
+            fixed = [fixed] * 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
         for ia in geometry:
-            self._write(fmt.format(*geometry.xyz[ia, :]))
+            self._write(fmt.format(*geometry.xyz[ia, :]) + todyn(fixed[ia]))
 
     @sile_fh_open(True)
     def read_supercell(self):
@@ -92,8 +111,15 @@ def read_supercell(self):
         return SuperCell(cell)
 
     @sile_fh_open()
-    def read_geometry(self):
-        """ Returns Geometry object from the CONTCAR/POSCAR file
+    def read_geometry(self, ret_fixed=False):
+        r""" Returns Geometry object from the CONTCAR/POSCAR file
+
+        Possibly also return the dynamics (if present)
+
+        Parameters
+        ----------
+        ret_fixed : bool, optional
+           also read selective dynamics (if present), if not, a list of False will be returned
         """
         sc = self.read_supercell()
 
@@ -124,9 +150,9 @@ def read_geometry(self):
         # Read whether this is selective or direct
         # Currently direct is not used
         opt = self.readline()
-        #direct = True
+        dynamics = False
         if opt[0] in 'Ss':
-            #direct = False
+            dynamics = True
             opt = self.readline()
 
         # Check whether this is in fractional or direct
@@ -137,18 +163,27 @@ def read_geometry(self):
 
         # Number of atoms
         na = len(atom)
+        # pre-create the fixed list
+        fixed = [[False] * 3] * na
 
-        xyz = np.empty([na, 3], np.float64)
+        xyz = _a.emptyd([na, 3])
         for ia in range(na):
-            xyz[ia, :] = list(map(float, self.readline().split()[:3]))
+            line = self.readline().split()
+            xyz[ia, :] = list(map(float, line[:3]))
+            if dynamics:
+                fixed[ia] = list(map(lambda x: x.lower() == 't', line[3:6]))
+
         if cart:
             # The unit of the coordinates are cartesian
             xyz *= self._scale
         else:
-            xyz = np.dot(xyz, sc.cell)
+            xyz = xyz.dot(sc.cell)
 
         # The POT/CONT-CAR does not contain information on the atomic species
-        return Geometry(xyz=xyz, atom=atom, sc=sc)
+        geom = Geometry(xyz=xyz, atom=atom, sc=sc)
+        if ret_fixed:
+            return geom, np.array(fixed, dtype=np.bool_)
+        return geom
 
     def ArgumentParser(self, p=None, *args, **kwargs):
         """ Returns the arguments that is available for this Sile """

sisl/io/vasp/tests/test_car.py

@@ -37,3 +37,40 @@ def test_geometry_car_allsame(sisl_tmp):
     geom.write(carSileVASP(f, 'w'))
 
     assert carSileVASP(f).read_geometry() == geom
+
+
+def test_geometry_car_allsame(sisl_tmp):
+    f = sisl_tmp('test_read_write.POSCAR', _dir)
+
+    atoms = Atom[1]
+    xyz = np.random.rand(10, 3)
+    geom = Geometry(xyz, atoms, 100)
+
+    geom.write(carSileVASP(f, 'w'))
+
+    assert carSileVASP(f).read_geometry() == geom
+
+
+def test_geometry_car_fixed(sisl_tmp):
+    f = sisl_tmp('test_fixed.POSCAR', _dir)
+
+    atoms = Atom[1]
+    xyz = np.random.rand(10, 3)
+    geom = Geometry(xyz, atoms, 100)
+
+    read = carSileVASP(f)
+
+    geom.write(carSileVASP(f, 'w'))
+    g, fix = read.read_geometry(ret_fixed=True)
+    assert not np.any(fix)
+
+    geom.write(carSileVASP(f, 'w'), fixed=True)
+    g, fix = read.read_geometry(ret_fixed=True)
+    assert np.all(fix)
+
+    fixed = [False] * len(geom)
+    fixed[0] = [True, False, True]
+    geom.write(carSileVASP(f, 'w'), fixed=fixed)
+    g, fix = read.read_geometry(ret_fixed=True)
+    assert np.array_equal(fixed[0], fix[0])
+    assert not np.any(fix[1:])