add derivatives for dispersion

kwant/physics/dispersion.py

@@ -35,7 +35,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
+    momentum. Velocities and velocity derivatives are calculated if the flag
+    ``deriv`` is set.
 
     Examples
     --------
@@ -57,9 +58,65 @@ class Bands:
         self.hop[:, : hop.shape[1]] = hop
         self.hop[:, hop.shape[1]:] = 0
 
-    def __call__(self, k):
-        # Note: Equation to solve is
-        #       (V^\dagger e^{ik} + H + V e^{-ik}) \psi = E \psi
+    def __call__(self, k, deriv=0):
+        """Calculate all energies :math:`E`, velocities :math:`v`
+        and velocity derivatives `v'` for a given momentum :math:`k`
+
+        :math:`E_n, \quad v_n = dE_n / dk, \quad v'_n = d^2E_n / dk^2,
+        \quad n \in \{0, nbands - 1\}`
+
+        :math:`nbands` is the number of open modes
+
+        Parameters
+        ----------
+        k : float
+            momentum
+
+        deriv : {0, 1, 2}, optional
+            Maximal derivative order to calculate. Default is zero
+
+
+        Returns
+        ----------
+        ener : numpy float array
+            energies (and optionally also higher momentum derivatives)
+
+            if deriv = 0
+                numpy float array of the energies :math:`E`, shape (nbands,)
+            if deriv > 0
+                numpy float array, shape (deriv + 1, nbands) of
+                energies and derivatives :math:`(E, E', E'')`
+
+        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.
+        """
+        # Equation to solve is
+        # (V^\dagger e^{ik} + H + V e^{-ik}) \psi = E \psi
+        # we obtain the derivatives by perturbing around momentum k
+
         mat = self.hop * complex(math.cos(k), -math.sin(k))
-        mat += mat.conjugate().transpose() + self.ham
-        return np.sort(np.linalg.eigvalsh(mat).real)
+        ham = mat + mat.conjugate().transpose() + self.ham
+        if deriv == 0:
+            return np.sort(np.linalg.eigvalsh(ham).real)
+
+        ener, psis = np.linalg.eigh(ham)
+        h1 = 1j*(- mat + mat.conjugate().transpose())
+        ph1p = np.dot(psis.conjugate().transpose(), np.dot(h1, psis))
+        velo = np.real(np.diag(ph1p))
+        if deriv == 1:
+            return np.array([ener, velo])
+
+        ediff = ener.reshape(-1, 1) - ener.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))
+        if deriv == 2:
+            return np.array([ener, velo, curv])
+        raise NotImplementedError('deriv= {} not implemented'.format(deriv))

kwant/physics/tests/test_dispersion.py

@@ -8,10 +8,12 @@
 
 from numpy.testing import assert_array_almost_equal, assert_almost_equal
 from pytest import raises
+from numpy import linspace
 
 import kwant
 from math import pi, cos, sin
 
+
 def test_band_energies(N=5):
     syst = kwant.Builder(kwant.TranslationalSymmetry((-1, 0)))
     lat = kwant.lattice.square()
@@ -26,6 +28,7 @@ def test_band_energies(N=5):
         assert_array_almost_equal(sorted(energies),
                                   sorted([2 - 2 * cos(k), 4 - 2 * cos(k)]))
 
+
 def test_same_as_lead():
     syst = kwant.Builder(kwant.TranslationalSymmetry((-1,)))
     lat = kwant.lattice.chain()
@@ -39,6 +42,7 @@ def test_same_as_lead():
     for momentum in momenta:
         assert_almost_equal(bands(momentum)[0], 0)
 
+
 def test_raise_nonhermitian():
     syst = kwant.Builder(kwant.TranslationalSymmetry((-1,)))
     lat = kwant.lattice.chain()
@@ -46,3 +50,59 @@ def test_raise_nonhermitian():
     syst[lat(0), lat(1)] = complex(cos(0.2), sin(0.2))
     syst = syst.finalized()
     raises(ValueError, kwant.physics.Bands, syst)
+
+
+def test_band_velocities():
+    syst = kwant.Builder(kwant.TranslationalSymmetry((-1, 0)))
+    lat = kwant.lattice.square()
+    syst[lat(0, 0)] = 1
+    syst[lat(0, 1)] = 3
+    syst[lat(1, 0), lat(0, 0)] = -1
+    syst[lat(1, 1), lat(0, 1)] = 2
+    bands = kwant.physics.Bands(syst.finalized())
+    eps = 1E-4
+    for k in linspace(-pi, pi, 200):
+        vel = bands(k, deriv=1)[1]
+        # higher order formula for first derivative to get required accuracy
+        num_vel = (- bands(k+2*eps) + bands(k-2*eps) +
+                   8*(bands(k+eps) - bands(k-eps))) / (12 * eps)
+        assert_array_almost_equal(vel, num_vel)
+
+
+def test_band_velocity_derivative():
+    syst = kwant.Builder(kwant.TranslationalSymmetry((-1, 0)))
+    lat = kwant.lattice.square()
+    syst[lat(0, 0)] = 1
+    syst[lat(0, 1)] = 3
+    syst[lat(1, 0), lat(0, 0)] = -1
+    syst[lat(1, 1), lat(0, 1)] = 2
+    bands = kwant.physics.Bands(syst.finalized())
+    eps = 1E-4
+    eps2 = eps * eps
+    c3 = 1 / 90
+    c2 = - 3 / 20
+    c1 = 3 / 2
+    c0 = - 49 / 18
+    for k in linspace(-pi, pi, 200):
+        dvel = bands(k, deriv=2)[2]
+        # higher order formula for second derivative to get required accuracy
+        num_dvel = (c3 * (bands(k+3*eps) + bands(k-3*eps)) +
+                    c2 * (bands(k+2*eps) + bands(k-2*eps)) +
+                    c1 * (bands(k+eps) + bands(k-eps)) +
+                    c0 * bands(k)) / eps2
+        assert_array_almost_equal(dvel, num_dvel)
+
+
+def test_raise_implemented():
+    syst = kwant.Builder(kwant.TranslationalSymmetry((-1, 0)))
+    lat = kwant.lattice.square()
+    syst[[lat(0, 0), lat(0, 1)]] = 3
+    syst[lat(0, 1), lat(0, 0)] = -1
+    syst[((lat(1, y), lat(0, y)) for y in range(2))] = -1
+    bands = kwant.physics.Bands(syst.finalized())
+    assert bands(1.).shape == (2,)
+    assert bands(1., deriv=0).shape == (2,)
+    assert bands(1., deriv=1).shape == (2, 2)
+    assert bands(1., deriv=2).shape == (3, 2)
+    raises(NotImplementedError, bands, 1., -1)
+    raises(NotImplementedError, bands, 1., 3)