Commit 5486af00 by 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