enh: allowed bool arrays for state objects

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

Author: zerothi

sisl/physics/state.py

@@ -1,5 +1,6 @@
 import numpy as np
 from numpy import einsum
+from numpy import ndarray, bool_
 
 import sisl._array as _a
 from sisl.messages import warn
@@ -31,6 +32,14 @@ def __init__(self, parent, **info):
         self.parent = parent
         self.info = info
 
+    def _sanitize_index(self, idx):
+        r""" Ensure indices are transferred to acceptable integers """
+        if isinstance(idx, ndarray) and idx.dtype == bool_:
+            return np.flatnonzero(idx)
+        elif isinstance(idx, (list, tuple)) and isinstance(idx[0], bool):
+            return np.flatnonzero(idx)
+        return _a.asarrayi(idx).ravel()
+
 
 class Coefficient(ParentContainer):
     """ An object holding coefficients for a parent with info
@@ -136,7 +145,7 @@ def sub(self, idx):
         Coefficient
             a new coefficient only containing the requested elements
         """
-        idx = _a.asarrayi(idx).ravel() # this ensures that the first dimension is preserved
+        idx = self._sanitize_index(idx)
         sub = self.__class__(self.c[idx].copy(), self.parent)
         sub.info = self.info
         return sub
@@ -262,7 +271,7 @@ def sub(self, idx):
         State
            a new state only containing the requested elements
         """
-        idx = _a.asarrayi(idx).ravel() # this ensures that the first dimension is preserved
+        idx = self._sanitize_index(idx)
         sub = self.__class__(self.state[idx].copy(), self.parent)
         sub.info = self.info
         return sub
@@ -705,7 +714,7 @@ def outer(self, idx=None):
         """
         if idx is None:
             return einsum('k,ki,kj->ij', self.c, self.state, _conj(self.state))
-        idx = np.asarray(idx).ravel()
+        idx = self._sanitize_index(idx)
         return einsum('k,ki,kj->ij', self.c[idx], self.state[idx], _conj(self.state[idx]))
 
     def sort(self, ascending=True):
@@ -763,7 +772,7 @@ def sub(self, idx):
         -------
         StateC : a new object with a subset of the states
         """
-        idx = _a.asarrayi(idx).ravel()
+        idx = self._sanitize_index(idx)
         sub = self.__class__(self.state[idx, ...], self.c[idx], self.parent)
         sub.info = self.info
         return sub

sisl/physics/tests/test_state.py

@@ -116,6 +116,8 @@ def test_state_sub1():
     for i, sub in enumerate(state.iter(True)):
         assert (sub ** 2).sum() == norm2[i]
 
+    assert np.allclose(state.sub([False, True, False, True]).state, state.sub([1, 3]).state)
+
 
 def test_state_outer1():
     state = ar(10, 10)