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)