Commit 5486af00 authored by Kloss's avatar Kloss Committed by Joseph Weston

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

parent d358f668
......@@ -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
......@@ -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)
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment