enh: fixed #182

Now BrillouinZone objects can return DataArray objects from
xarray (if it can be imported).

It is called asdatarray.

This makes things much easier for certain users.

By default the dimensions are named 'k', 'v1', ...
but users may provide v1, ... names by an argument.
Also, the name of the DataArray will be the function name

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

Author: zerothi

CHANGELOG

@@ -1,3 +1,18 @@
+0.9.9
+=====
+
+- Enabled xarray.DataArray returning from BrillouinZone objects #182
+
+- Several improvements to outSileSiesta.read_scf (thanks to Pol Febrer)
+
+- A huge performance increase for data extraction in tbtncSileTbtrans
+	(thanks to Gaetano Calogero for finding the bottleneck)
+
+- Added preliminary usage of Mixers, primarily intented for extending
+	sisl operations where SCF are used (may heavily change).
+
+- Now sisl is Python >=3.6 only
+
 0.9.8
 =====
 

sisl/io/siesta/bands.py

@@ -16,7 +16,7 @@ class bandsSileSiesta(SileSiesta):
     @sile_fh_open()
     def read_data(self, as_dataarray=False):
         """ Returns data associated with the bands file
-        
+
         Parameters
         --------
         as_dataarray: boolean, optional
@@ -84,12 +84,12 @@ def read_data(self, as_dataarray=False):
                 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])),

sisl/io/siesta/out.py

@@ -346,7 +346,7 @@ def return_forces(Fs):
 
             if last:
                 return return_forces(Fs[-1])
- 
+
             return return_forces(Fs)
 
         return return_forces(next_force())

sisl/io/siesta/tests/test_bands.py

@@ -16,6 +16,7 @@ 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

sisl/physics/brillouinzone.py

@@ -101,6 +101,12 @@
 from sisl.supercell import SuperCell
 from sisl.grid import Grid
 
+try:
+    import xarray
+    _has_xarray = True
+except ImportError:
+    _has_xarray = False
+
 
 __all__ = ['BrillouinZone', 'MonkhorstPack', 'BandStructure']
 
@@ -477,7 +483,7 @@ def asarray(self):
         ...    spin_moment = (es.spin_moment(E, distribution=dist) * occ.reshape(-1, 1)).sum(0)
         ...    return oplist([DOS, PDOS, spin_moment])
         >>> bz = BrillouinZone(hamiltonian)
-        >>> DOS, PDOS, spin_moment = bz.asaverage().eigenstate(wrap=wrap)
+        >>> DOS, PDOS, spin_moment = bz.asarray().eigenstate(wrap=wrap)
 
         See Also
         --------
@@ -567,6 +573,100 @@ def _call(self, *args, **kwargs):
         setattr(self, '_bz_call', types.MethodType(_call, self))
         return self
 
+    if _has_xarray:
+        def asdataarray(self):
+            r""" Return `self` with `xarray.DataArray` returned quantities
+
+            This forces the `__call__` routine to return a single `xarray.DataArray`.
+
+            Notes
+            -----
+            If you wrap the sub-method to return multiple data-sets, you should use `asdataset`
+            instead which returns a combination of data-arrays (so-called `xarray.Dataset`).
+
+            All invocations of sub-methods are added these keyword-only arguments:
+
+            eta : bool, optional
+                if true a progress-bar is created, default false.
+            wrap : callable, optional
+                a function that accepts the output of the given routine and post-process
+                it. Defaults to ``lambda x: x``.
+            coords : list of str or list of (str, array), optional
+                a list of coordinates used in ``xarray.DataArray(..., coords=coords)``.
+                By default the coordinates are named ``['k', 'v1', 'v2', ...]``
+                depending on the shape of the returned quantity.
+                These may optionally be a list of tuples (not a dictionary)!
+
+            Examples
+            --------
+            >>> obj = BrillouinZone(...)
+            >>> obj.asdataarray().eigh(eta=True)
+
+            See Also
+            --------
+            asyield : all output returned through an iterator
+            asaverage : take the average (with k-weights) of the Brillouin zone
+            assum : return the sum of values in the Brillouin zone
+            aslist : all output returned as a Python list
+            """
+
+            def _call(self, *args, **kwargs):
+                func = self._bz_get_func()
+
+                # xarray specific data (default to function name)
+                name = kwargs.pop('name', func.__name__)
+                coords = kwargs.pop('coords', None)
+
+                has_wrap = 'wrap' in kwargs
+                if has_wrap:
+                    wrap = allow_kwargs('parent', 'k', 'weight')(kwargs.pop('wrap'))
+                eta = tqdm_eta(len(self), self.__class__.__name__ + '.asarray',
+                               'k', kwargs.pop('eta', False))
+                parent = self.parent
+                k = self.k
+                w = self.weight
+                if has_wrap:
+                    v = wrap(func(*args, k=k[0], **kwargs), parent=parent, k=k[0], weight=w[0])
+                else:
+                    v = func(*args, k=k[0], **kwargs)
+                if v.ndim == 0:
+                    a = np.empty([len(self)], dtype=v.dtype)
+                else:
+                    a = np.empty((len(self), ) + v.shape, dtype=v.dtype)
+                a[0] = v
+                del v
+                eta.update()
+                if has_wrap:
+                    for i in range(1, len(k)):
+                        a[i] = wrap(func(*args, k=k[i], **kwargs), parent=parent, k=k[i], weight=w[i])
+                        eta.update()
+                else:
+                    for i in range(1, len(k)):
+                        a[i] = func(*args, k=k[i], **kwargs)
+                        eta.update()
+                eta.close()
+
+                # Create coords
+                if coords is None:
+                    coords = [('k', _a.arangei(len(self)))]
+                    for i, v in enumerate(a.shape[1:]):
+                        coords.append((f"v{i+1}", _a.arangei(v)))
+                else:
+                    coords = list(coords)
+                    coords.insert(0, ('k', _a.arangei(len(self))))
+                    for i in range(1, len(coords)):
+                        if isinstance(coords[i], str):
+                            coords[i] = (coords[i], _a.arangei(a.shape[i]))
+                attrs = {'bz': self,
+                         'parent': self.parent,
+                }
+
+                return xarray.DataArray(a, coords=coords, name=name, attrs=attrs)
+
+            # Set instance __bz_call
+            setattr(self, '_bz_call', types.MethodType(_call, self))
+            return self
+
     def aslist(self):
         """ Return `self` with `list` returned quantities
 

sisl/physics/tests/test_brillouinzone.py

@@ -248,6 +248,30 @@ def test_as_simple(self):
         assert np.allclose((asarray / len(bz)).sum(0), asaverage)
         bz.asnone().eigh()
 
+    def test_as_dataarray(self):
+        try:
+            import xarray
+        except ImportError:
+            pytest.skip('xarray not available')
+
+        from sisl import geom, Hamiltonian
+        g = geom.graphene()
+        H = Hamiltonian(g)
+        H.construct([[0.1, 1.44], [0, -2.7]])
+
+        bz = MonkhorstPack(H, [2, 2, 2], trs=False)
+
+        # Assert that as* all does the same
+        asarray = bz.asarray().eigh()
+        asdarray = bz.asdataarray().eigh()
+        assert np.allclose(asarray, asdarray.values)
+        assert isinstance(asdarray.bz, MonkhorstPack)
+        assert isinstance(asdarray.parent, Hamiltonian)
+        assert asdarray.dims == ('k', 'v1')
+
+        asdarray = bz.asdataarray().eigh(coords=['orb'])
+        assert asdarray.dims == ('k', 'orb')
+
     def test_as_single(self):
         from sisl import geom, Hamiltonian
         g = geom.graphene()