Commit 5aa6d068 authored by Kloss's avatar Kloss Committed by Joseph Weston

add eigenvector computation and return energies with higher derivatives or vectors as a tuple

parent db4c89e2
......@@ -58,7 +58,7 @@ class Bands:
self.hop[:, : hop.shape[1]] = hop
self.hop[:, hop.shape[1]:] = 0
def __call__(self, k, *, derivative_order=0):
def __call__(self, k, *, derivative_order=0, return_eigenvectors=False):
"""Calculate all energies :math:`E`, velocities :math:`v`
and velocity derivatives `v'` for a given momentum :math:`k`
......@@ -102,6 +102,9 @@ class Bands:
mat = self.hop * complex(math.cos(k), -math.sin(k))
ham = mat + mat.conjugate().transpose() + self.ham
if derivative_order == 0:
if return_eigenvectors:
ener, psis = np.linalg.eigh(ham)
return ener, psis
return np.sort(np.linalg.eigvalsh(ham).real)
ener, psis = np.linalg.eigh(ham)
......@@ -109,7 +112,9 @@ class Bands:
ph1p = np.dot(psis.conjugate().transpose(), np.dot(h1, psis))
velo = np.real(np.diag(ph1p))
if derivative_order == 1:
return np.array([ener, velo])
if return_eigenvectors:
return ener, velo, psis
return ener, velo
ediff = ener.reshape(-1, 1) - ener.reshape(1, -1)
ediff = np.divide(1, ediff, out=np.zeros_like(ediff), where=ediff != 0)
......@@ -118,6 +123,8 @@ class Bands:
np.dot(psis.conjugate().transpose(), np.dot(h2, psis)))) +
2 * np.sum(ediff * np.abs(ph1p)**2, axis=1))
if derivative_order == 2:
return np.array([ener, velo, curv])
if return_eigenvectors:
return ener, velo, curv, psis
return ener, velo, curv
raise NotImplementedError('derivative_order= {} not implemented'
.format(derivative_order))
......@@ -6,6 +6,7 @@
# the file AUTHORS.rst at the top-level directory of this distribution and at
# http://kwant-project.org/authors.
import numpy as np
from numpy.testing import assert_array_almost_equal, assert_almost_equal
from pytest import raises
from numpy import linspace
......@@ -14,14 +15,17 @@ import kwant
from math import pi, cos, sin
def test_band_energies(N=5):
def make_lead():
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
return syst
band_energies = kwant.physics.Bands(syst.finalized())
def test_band_energies(N=5):
band_energies = kwant.physics.Bands(make_lead().finalized())
for i in range(-N, N):
k = i * pi / N
energies = band_energies(k)
......@@ -93,16 +97,45 @@ def test_band_velocity_derivative():
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., derivative_order=0).shape == (2,)
assert bands(1., derivative_order=1).shape == (2, 2)
assert bands(1., derivative_order=2).shape == (3, 2)
raises(NotImplementedError, bands, 1., derivative_order=-1)
raises(NotImplementedError, bands, 1., derivative_order=3)
def test_eigenvector_calculation(k=1.2, lead=make_lead()): # k-point arbitrary
bands = kwant.physics.Bands(lead.finalized())
# check if eigenvalues are sorted and always sorted in the same way
energies = bands(k)
assert (np.sort(energies) == energies).all()
e_1 = bands(k, derivative_order=0)
e_2, *_ = bands(k, derivative_order=0, return_eigenvectors=True)
assert_array_almost_equal(energies, e_1)
assert_array_almost_equal(energies, e_2)
for order in [1, 2]:
e_1, *_ = bands(k, derivative_order=order)
e_2, *_ = bands(k, derivative_order=order, return_eigenvectors=True)
assert_array_almost_equal(energies, e_1)
assert_array_almost_equal(energies, e_2)
# check eigenvector
unity_matrix = np.identity(len(energies))
zero_matrix = 0 * unity_matrix
for order in [0, 1, 2]:
eval_evec = bands(k, derivative_order=order, return_eigenvectors=True)
assert len(eval_evec) == order + 2
energies, *_, psi = eval_evec
# completeness relation
outer_product = psi @ psi.conjugate().transpose()
assert_array_almost_equal(outer_product.real, unity_matrix)
assert_array_almost_equal(outer_product.imag, zero_matrix)
# eigenvalue equation
mat = bands.hop * np.exp(- 1j * k)
hamiltonian = mat + mat.conjugate().transpose() + bands.ham
energies = np.asarray([energies])
assert_array_almost_equal(hamiltonian @ psi, energies * psi)
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)
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