change matrix computation depending on size, change variable names to be more explicit

kwant/physics/dispersion.py

@@ -6,6 +6,8 @@
 # the file AUTHORS.rst at the top-level directory of this distribution and at
 # http://kwant-project.org/authors.
 
+"""Band structure calculation for the leads"""
+
 import math
 import numpy as np
 from .. import system
@@ -35,8 +37,8 @@ class Bands:
     An instance of this class can be called like a function.  Given a momentum
     (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. Velocities and velocity derivatives are calculated if the flag
-    ``derivative_order`` is set.
+    momentum. If `derivative_order > 0` or `return_eigenvectors = True`,
+    additional arrays are returned returned.
 
     Examples
     --------
@@ -58,15 +60,24 @@ class Bands:
         self.hop[:, : hop.shape[1]] = hop
         self.hop[:, hop.shape[1]:] = 0
 
-    def __call__(self, k, *, derivative_order=0, return_eigenvectors=False):
-        """Calculate all energies :math:`E`, velocities :math:`v`
-        and velocity derivatives `v'` for a given momentum :math:`k`
+        # 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`
 
-        :math:`E_n, \quad v_n = dE_n / dk, \quad v'_n = d^2E_n / dk^2,
+        :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
 
+        Eigenvectors are orthonormal.
+
+
         Parameters
         ----------
         k : float
@@ -75,17 +86,27 @@ class Bands:
         derivative_order : {0, 1, 2}, optional
             Maximal derivative order to calculate. Default is zero
 
+        return_eigenvectors : bool, optional
+            if set to `True` return the eigenvectors as last tuple element
+
 
         Returns
         ----------
-        ener : numpy float array
-            energies (and optionally also higher momentum derivatives)
-
-            if derivative_order = 0
-                numpy float array of the energies :math:`E`, shape (nbands,)
-            if derivative_order > 0
-                numpy float array, shape (derivative_order + 1, nbands) of
-                energies and derivatives :math:`(E, E', E'')`
+        energies : numpy float array, size `nbands`
+            energies :math:`E`
+        velocities : numpy float array, size `nbands`
+            velocities (first order energy derivatives :math:`E'`)
+        curvatures : numpy float array, size `nbands`
+            curvatures (second energy derivatives :math:`E''`)
+        eigenvectors : numpy float array, size `nbands`
+            eigenvectors
+
+        The number of retured elements varies. If `derivative_order > 0`
+        we return all derivatives up to the order `derivative_order`,
+        and total the number of returned elements is `derivative_order + 1`
+        If 'return_eigenvectors=True', in addition the eigenvectors are
+        returned as the last element. In that case, the total number of
+        returned elements is `derivative_order + 2`.
 
         Notes
         -----
@@ -98,33 +119,43 @@ class Bands:
         # Equation to solve is
         # (V^\dagger e^{ik} + H + V e^{-ik}) \psi = E \psi
         # we obtain the derivatives by perturbing around momentum k
+        # The perturbed Hamiltonian H(k) is
+        # H(k+dk) = H_0(k) + H_1(k) dk + H_2(k) dk^2 + O(dk^3)
+        # and we use h0, h1 and h2 in the code to refer to above terms
 
         mat = self.hop * complex(math.cos(k), -math.sin(k))
-        ham = mat + mat.conjugate().transpose() + self.ham
+        h0 = mat + mat.conjugate().transpose() + self.ham
         if derivative_order == 0:
             if return_eigenvectors:
-                ener, psis = np.linalg.eigh(ham)
-                return ener, psis
-            return np.sort(np.linalg.eigvalsh(ham).real)
+                energies, eigenvectors = np.linalg.eigh(h0)
+                return energies, eigenvectors
+            return np.sort(np.linalg.eigvalsh(h0).real)
 
-        ener, psis = np.linalg.eigh(ham)
+        energies, eigenvectors = np.linalg.eigh(h0)
         h1 = 1j*(- mat + mat.conjugate().transpose())
-        ph1p = np.dot(psis.conjugate().transpose(), np.dot(h1, psis))
-        velo = np.real(np.diag(ph1p))
+        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
         if derivative_order == 1:
             if return_eigenvectors:
-                return ener, velo, psis
-            return ener, velo
+                return energies, velocities, eigenvectors
+            return energies, velocities
 
-        ediff = ener.reshape(-1, 1) - ener.reshape(1, -1)
+        # 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())
-        curv = (np.real(np.diag(
-                np.dot(psis.conjugate().transpose(), np.dot(h2, psis)))) +
-                2 * np.sum(ediff * np.abs(ph1p)**2, axis=1))
+        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 ener, velo, curv, psis
-            return ener, velo, curv
+                return energies, velocities, curvatures, eigenvectors
+            return energies, velocities, curvatures
         raise NotImplementedError('derivative_order= {} not implemented'
                                   .format(derivative_order))

kwant/physics/tests/test_dispersion.py

@@ -99,7 +99,8 @@ def test_band_velocity_derivative():
 
 def test_eigenvector_calculation(k=1.2, lead=make_lead()):  # k-point arbitrary
 
-    bands = kwant.physics.Bands(lead.finalized())
+    lead = lead.finalized()
+    bands = kwant.physics.Bands(lead)
 
     # check if eigenvalues are sorted and always sorted in the same way
     energies = bands(k)
@@ -134,6 +135,14 @@ def test_eigenvector_calculation(k=1.2, lead=make_lead()):  # k-point arbitrary
         energies = np.asarray([energies])
         assert_array_almost_equal(hamiltonian @ psi, energies * psi)
 
+    # check that algorithm switch is correct
+    bands_switch = kwant.physics.Bands(lead)
+    bands_switch._crossover_size = 0
+    for order in [1, 2]:
+        e_1, v_1, *_ = bands(k, order)
+        e_2, v_2, *_ = bands_switch(k, order)
+        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())