Merge branch 'master' of github.com:zerothi/sisl

Author: tfrederiksen

sisl/io/siesta/bands.py

@@ -1,5 +1,6 @@
 import numpy as np
 
+import sisl._array as _a
 from sisl.utils import strmap
 from sisl.utils.cmd import default_ArgumentParser, default_namespace
 from ..sile import add_sile, sile_fh_open
@@ -13,8 +14,16 @@ class bandsSileSiesta(SileSiesta):
     """ Bandstructure information """
 
     @sile_fh_open()
-    def read_data(self):
-        """ Returns data associated with the bands file """
+    def read_data(self, as_dataarray=False):
+        """ Returns data associated with the bands file
+        
+        Parameters
+        --------
+        as_dataarray: boolean, optional
+            if `True`, the information is returned as an `xarray.DataArray`
+            Ticks (if read) are stored as an attribute of the DataArray 
+            (under `array.ticks` and `array.ticklabels`)
+        """
         band_lines = False
 
         # Luckily the data is in eV
@@ -35,19 +44,19 @@ def read_data(self):
         no, ns, nk = map(int, l.split())
 
         # Create the data to contain all band points
-        b = np.empty([nk, ns, no], np.float64)
+        b = _a.emptyd([nk, ns, no])
 
         # for band-lines
         if band_lines:
-            k = np.empty([nk], np.float64)
+            k = _a.emptyd([nk])
             for ik in range(nk):
                 l = [float(x) for x in self.readline().split()]
                 k[ik] = l[0]
                 del l[0]
                 # Now populate the eigenvalues
                 while len(l) < ns * no:
                     l.extend([float(x) for x in self.readline().split()])
-                l = np.array(l, np.float64)
+                l = _a.arrayd(l)
                 l.shape = (ns, no)
                 b[ik, :, :] = l[:, :] - Ef
             # Now we need to read the labels for the points
@@ -61,7 +70,7 @@ def read_data(self):
             vals = (xlabels, labels), k, b
 
         else:
-            k = np.empty([nk, 3], np.float64)
+            k = _a.emptyd([nk, 3])
             for ik in range(nk):
                 l = [float(x) for x in self.readline().split()]
                 k[ik, :] = l[0:3]
@@ -71,10 +80,21 @@ def read_data(self):
                 # Now populate the eigenvalues
                 while len(l) < ns * no:
                     l.extend([float(x) for x in self.readline().split()])
-                l = np.array(l, np.float64)
+                l = _a.arrayd(l)
                 l.shape = (ns, no)
                 b[ik, :, :] = l[:, :] - Ef
             vals = k, b
+        
+        if as_dataarray:
+            from xarray import DataArray
+
+            ticks = {"ticks": xlabels, "ticklabels": labels} if band_lines else {}
+            
+            return DataArray(b, name="Energy", attrs=ticks,
+                             coords=[("k", k),
+                                     ("spin", _a.arangei(0, b.shape[1])),
+                                     ("band", _a.arangei(0, b.shape[2]))])
+
         return vals
 
     @default_ArgumentParser(description="Manipulate bands file in sisl.")

sisl/io/siesta/out.py

@@ -291,11 +291,15 @@ def next_force():
             # Now read data
             F = []
             line = self.readline()
+            if 'siesta:' in line:
+                # This is the final summary, we don't need to read it as it does not contain new information
+                # and also it make break things since max forces are not written there
+                return None
 
             # First, we encounter the atomic forces
             while '---' not in line:
                 line = line.split()
-                if not total or max:
+                if not (total or max):
                     F.append([float(x) for x in line[-3:]])
                 line = self.readline()
                 if line == '':
@@ -328,8 +332,8 @@ def return_forces(Fs):
             if max and total:
                 return (Fs[..., :-1], Fs[..., -1])
             elif max and not all:
-                # This will return a float (or actually a numpy.dtype)
-                return Fs[0]
+                # This will return a float
+                return np.atleast_1d(Fs)[0]
             return Fs
 
         if all or last:
@@ -342,11 +346,8 @@ def return_forces(Fs):
                 Fs.append(F)
 
             if last:
-                return return_forces(Fs[-2])
-                # F[-2] is really the same as F[-1], the last forces are stated twice
-                # However, the maxForce is not stated in the final summary, that's why we use F[-2]
-            if self.job_completed:
-                return return_forces(Fs[:-1])
+                return return_forces(Fs[-1])
+ 
             return return_forces(Fs)
 
         return return_forces(next_force())

sisl/io/siesta/pdos.py

@@ -106,7 +106,7 @@ def to(o, DOS):
                 return xr.DataArray(data=process(DOS).reshape(shape),
                                     dims=dims, coords=coords, name='PDOS')
 
-            D = xr.DataArray()
+            D = xr.DataArray([])
         else:
             def to(o, DOS):
                 return process(DOS)

sisl/io/siesta/tests/test_bands.py

@@ -16,7 +16,6 @@ def test_fe(sisl_files):
     assert eig.shape == (131, 2, 15)
     assert len(labels[0]) == 5
 
-
 def test_fe_ArgumentParser(sisl_files, sisl_tmp):
     try:
         import matplotlib
@@ -28,3 +27,17 @@ def test_fe_ArgumentParser(sisl_files, sisl_tmp):
     p.parse_args([], namespace=ns)
     p.parse_args(['--energy', ' -2:2'], namespace=ns)
     p.parse_args(['--energy', ' -2:2', '--plot', png], namespace=ns)
+
+
+def test_fe_xarray(sisl_files, sisl_tmp):
+    try:
+        import xarray
+    except ImportError:
+        pytest.skip('xarray not available')
+    si = sisl.get_sile(sisl_files(_dir, 'fe.bands'))
+
+    bands = si.read_data(as_dataarray=True)
+    assert len(bands['k']) == 131
+    assert len(bands['spin']) == 2
+    assert len(bands['band']) == 15
+    assert len(bands.ticks) == len(bands.ticklabels) == 5