change _crossover_size to class attribute, clean up docstrings and comments, __call__ restructured

kwant/physics/dispersion.py

@@ -38,7 +38,7 @@ class Bands:
     (currently this must be a scalar as all infinite systems are quasi-1-d), it
     returns a NumPy array containing the eigenenergies of all modes at this
     momentum. If `derivative_order > 0` or `return_eigenvectors = True`,
-    additional arrays are returned returned.
+    additional arrays are returned.
 
     Examples
     --------
@@ -48,6 +48,7 @@ class Bands:
     >>> pyplot.plot(momenta, energies)
     >>> pyplot.show()
     """
+    _crossover_size = 8
 
     def __init__(self, sys, args=(), *, params=None):
         syst = sys
@@ -60,45 +61,36 @@ class Bands:
         self.hop[:, : hop.shape[1]] = hop
         self.hop[:, hop.shape[1]:] = 0
 
-        # Switch computation from diag(a^\dagger @ b) to sum(a.T * b, axis=0),
-        # which is computationally more efficient for larger matrices.
-        # a and b are NxN matrices, and we switch for N > _crossover_size.
-        # Both ways are numerically (but not necessarily bitwise) equivalent.
-        self._crossover_size = 8
-
     def __call__(self, k, derivative_order=0, return_eigenvectors=False):
-        """Calculate all energies :math:`E` (optionally higher derivatives
-        :math:`E', E''` and eigenvectors) for a given momentum :math:`k`
+        r"""Calculate band energies at a given momentum.
 
         :math:`E_n, \quad E'_n = dE_n / dk, \quad E''_n = d^2E_n / dk^2,
         \quad n \in \{0, nbands - 1\}`
 
-        :math:`nbands` is the number of open modes
-
+        :math:`nbands` is the number of open modes.
         Eigenvectors are orthonormal.
 
-
         Parameters
         ----------
         k : float
             momentum
 
         derivative_order : {0, 1, 2}, optional
-            Maximal derivative order to calculate. Default is zero
+            Maximal derivative order to calculate. Default is zero.
 
         return_eigenvectors : bool, optional
-            if set to `True` return the eigenvectors as last tuple element
-
+            if set to `True` return the eigenvectors as last tuple element.
+            By default, no eigenvectors are returned.
 
         Returns
         ----------
-        energies : numpy float array, size `nbands`
+        energies : numpy float array, size ``nbands``
             energies :math:`E`
-        velocities : numpy float array, size `nbands`
+        velocities : numpy float array, size ``nbands``
             velocities (first order energy derivatives :math:`E'`)
-        curvatures : numpy float array, size `nbands`
+        curvatures : numpy float array, size ``nbands``
             curvatures (second energy derivatives :math:`E''`)
-        eigenvectors : numpy float array, size `nbands`
+        eigenvectors : numpy float array, size ``nbands x nbands``
             eigenvectors
 
         The number of retured elements varies. If `derivative_order > 0`
@@ -110,11 +102,9 @@ class Bands:
 
         Notes
         -----
-        *   The ordering of the energies and velocities is the same and
-            according to the magnitude of the energy eigenvalues.
-            Therefore, the bands are not continously ordered.
-        *   The curvature `E''` can be only calculated
-            for non-degenerate bands.
+        *   All the output arrays are sorted from the lowest energy band
+            to the highest.
+        *   The curvature `E''` can be only calculated for non-degenerate bands.
         """
         # Equation to solve is
         # (V^\dagger e^{ik} + H + V e^{-ik}) \psi = E \psi
@@ -125,37 +115,46 @@ class Bands:
 
         mat = self.hop * complex(math.cos(k), -math.sin(k))
         h0 = mat + mat.conjugate().transpose() + self.ham
-        if derivative_order == 0:
-            if return_eigenvectors:
-                energies, eigenvectors = np.linalg.eigh(h0)
-                return energies, eigenvectors
-            return np.sort(np.linalg.eigvalsh(h0).real)
-
-        energies, eigenvectors = np.linalg.eigh(h0)
-        h1 = 1j*(- mat + mat.conjugate().transpose())
-        if derivative_order == 1 and len(energies) > self._crossover_size:
-            velocities = np.sum(eigenvectors.conjugate() * (h1 @ eigenvectors),
-                                axis=0).real
+
+        # compute energies and if required eigenvectors
+        if return_eigenvectors or derivative_order > 0:
+            energies, eigenvectors = np.linalg.eigh(h0)
         else:
-            ph1p = eigenvectors.conjugate().transpose() @ h1 @ eigenvectors
-            velocities = np.diag(ph1p).real
-        if derivative_order == 1:
-            if return_eigenvectors:
-                return energies, velocities, eigenvectors
-            return energies, velocities
-
-        # ediff_{i,j} =  1 / (E_i - E_j) if i != j, else 0
-        ediff = energies.reshape(-1, 1) - energies.reshape(1, -1)
-        ediff = np.divide(1, ediff, out=np.zeros_like(ediff), where=ediff != 0)
-        h2 = - (mat + mat.conjugate().transpose())
-        curvatures = np.sum(eigenvectors.conjugate() * (h2 @ eigenvectors)
-                            - 2 * ediff * np.abs(ph1p)**2, axis=0).real
-        # above expression is similar to: curvatures =
-        # np.diag(eigenvectors.conjugate().transpose() @ h2 @ eigenvectors
-        #         + 2 * ediff @ np.abs(ph1p)**2).real
-        if derivative_order == 2:
-            if return_eigenvectors:
-                return energies, velocities, curvatures, eigenvectors
-            return energies, velocities, curvatures
-        raise NotImplementedError('derivative_order= {} not implemented'
-                                  .format(derivative_order))
+            energies = np.sort(np.linalg.eigvalsh(h0).real)
+        output = (energies,)
+
+        if derivative_order >= 1:  # compute velocities
+            h1 = 1j*(- mat + mat.conjugate().transpose())
+            # Switch computation from diag(a^\dagger @ b) to sum(a.T * b, axis=0),
+            # which is computationally more efficient for larger matrices.
+            # a and b are NxN matrices, and we switch for N > _crossover_size.
+            # Both ways are numerically (but not necessarily bitwise) equivalent.
+            if derivative_order == 1 and len(energies) > self._crossover_size:
+                velocities = np.sum(eigenvectors.conjugate() * (h1 @ eigenvectors),
+                                    axis=0).real
+            else:
+                ph1p = eigenvectors.conjugate().transpose() @ h1 @ eigenvectors
+                velocities = np.diag(ph1p).real
+            output += (velocities,)
+        if derivative_order >= 2:  # compute curvatures
+            # ediff_{i,j} =  1 / (E_i - E_j) if i != j, else 0
+            ediff = energies.reshape((-1, 1)) - energies.reshape((1, -1))
+            ediff = np.divide(1, ediff, out=np.zeros_like(ediff),
+                              where=(ediff != 0))
+            h2 = - (mat + mat.conjugate().transpose())
+            curvatures = np.sum(eigenvectors.conjugate() * (h2 @ eigenvectors)
+                                - 2 * ediff * np.abs(ph1p)**2, axis=0).real
+            # above expression is similar to: curvatures =
+            # np.diag(eigenvectors.conjugate().transpose() @ h2 @ eigenvectors
+            #         + 2 * ediff @ np.abs(ph1p)**2).real
+            output += (curvatures,)
+        if derivative_order > 2:
+            raise NotImplementedError('Derivatives of the energy dispersion ' +
+                                      'only implemented up to second order.')
+        if return_eigenvectors:
+            output += (eigenvectors,)
+        if len(output) == 1:
+            # Backwards compatibility: if returning only energies,
+            # don't make it a length-1 tuple.
+            return output[0]
+        return output

kwant/physics/tests/test_dispersion.py

@@ -97,10 +97,10 @@ def test_band_velocity_derivative():
         assert_array_almost_equal(dvel, num_dvel)
 
 
-def test_eigenvector_calculation(k=1.2, lead=make_lead()):  # k-point arbitrary
-
-    lead = lead.finalized()
+def test_eigenvector_calculation():
+    lead = make_lead().finalized()
     bands = kwant.physics.Bands(lead)
+    k = 1.2   # k-point arbitrary
 
     # check if eigenvalues are sorted and always sorted in the same way
     energies = bands(k)
@@ -144,7 +144,7 @@ def test_eigenvector_calculation(k=1.2, lead=make_lead()):  # k-point arbitrary
         assert_array_almost_equal(v_1, v_2)
 
 
-def test_raise_implemented(k=1, lead=make_lead()):  # k-point arbitrary
-    bands = kwant.physics.Bands(lead.finalized())
-    raises(NotImplementedError, bands, k, derivative_order=-1)
+def test_raise_implemented():
+    k = 1  # k-point arbitrary
+    bands = kwant.physics.Bands(make_lead().finalized())
     raises(NotImplementedError, bands, k, derivative_order=3)