diff --git a/kwant/physics/dispersion.py b/kwant/physics/dispersion.py
index f4e0e42ed52a78c58fb9f168de1f4976138c955a..e3618796ced0fa35909553843a90b51f6372dae7 100644
--- a/kwant/physics/dispersion.py
+++ b/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
diff --git a/kwant/physics/tests/test_dispersion.py b/kwant/physics/tests/test_dispersion.py
index d1a6d9a08bdb2543ac1a32f836fb4e023a7c5996..a42d6deadf5d88bb0390318cb5442f88229742dc 100644
--- a/kwant/physics/tests/test_dispersion.py
+++ b/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)