Merge branch 'bandsXarray'

Author: zerothi

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/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