ex2.py

"""
1. Generate all reciprocal lattice vectors G up to a given length.
2.In units of Electronvolts construct the (square) characteristic matrix in the basis of all G by
    * Implementing a function V(g) as defined in class
    * Off-diagonal matrix elements are V(G1-G2) where G1 and G2 are the rec.
      lattice vectors corresponding to row and column of the matrix
    * The diagonal matrix elements are the energies of a free electron with wavevector k+G
3. Calculate the eigenvalues of this matrix for a number of k-points in the irreducible wedge
4. Plot these eigenvalues, this is the band structure
5. Calculate the fundamental energy gap
"""

import matplotlib.pyplot as plt
import numpy as np
from numpy.typing import NDArray
from ex1 import fcc
import scipy.constants as c
from typing import Optional

class bandstructure:
    """
    Utilities for bandstructure calculation

    Uses pseudopotential approximation
    """
    def __init__(self, material: fcc) -> None:
        """
        Args:
            material: Instance of fcc class
        """
        self.mat = material

        max_G = np.sqrt(11) * 2*np.pi/self.mat.a
        self.lat_vecs = self.mat.gen_lat_vecs(max_G)
        self.assemble_hamiltonian_wo_diag()

        self.E_vals = None
        self.k_vecs = None

    def pseudo_potential(self, dG: NDArray) -> complex: 
        """
        Returns the fourier coefficient of the pseudo potential
        belonging to the reciprocal lattice vector dG = Gi-Gj
        """
        b = self.mat.b
        l2 = round(np.sum((dG/b)**2))
        arg = np.sum(dG*self.mat.s)

        if l2 == 3:
            Vg = self.mat.Vs3*np.cos(arg) - 1j * self.mat.Va3*np.sin(arg)
        elif l2 == 4:
            Vg = 1j * self.mat.Va4*np.sin(arg)
        elif l2 == 8:
            Vg = self.mat.Vs8*np.cos(arg)
        elif l2 == 11:
            Vg = self.mat.Vs11*np.cos(arg) - 1j * self.mat.Va11*np.sin(arg)
        else:
            Vg = 0.0 + 0.0j

        return Vg

    def assemble_hamiltonian_wo_diag(self) -> NDArray:
        """
        Generates all off-diagonal elements of the hamiltonian matrix
        """
        num_vecs = len(self.lat_vecs)
        self.ham_wo_diag = np.zeros((num_vecs, num_vecs), dtype=complex)

        for i in range(num_vecs):
            Gi = self.lat_vecs[i]
            for j in range(num_vecs):
                if i != j:
                    Gj = self.lat_vecs[j]
                    self.ham_wo_diag[i, j] = self.pseudo_potential(Gi-Gj)

    def hamiltonian(self, k: NDArray) -> NDArray:
        """
        Returns the hamiltonian matrix for a specific k-vector
        """
        diag = c.hbar**2/(2*c.m_e) * np.sum((self.lat_vecs + k)**2, axis=1)
        return self.ham_wo_diag + np.diag(diag)

    def get_E(self, k: NDArray) -> NDArray:
        """
        Computes the Energy Eigenvalues belonging to a k-vector.
        """
        ham = self.hamiltonian(k)
        return np.linalg.eigvalsh(ham)

    def compute_bandstructure(self, num_points: int = 20) -> None:
        """
        Generates all Energy Eigenvalues along the L-Gamma-X-U-Gamma
        path of the irreducible wedge

        Args:
            num_points (int): the number of k-vectors between two high-symmetry points
        """
        E_vals = []

        self.tick_positions = []
        self.tick_labels = []
        i = 0

        k0 = None
        high_symmetry_points = ["L", "Gamma", "X", "U", "Gamma"]

        for point in high_symmetry_points:
            k1 = self.mat.hsp[point]

            if k0 is not None:
                vecs = self.mat.gen_k_vecs(k0, k1, num_points)
                for k in vecs:
                    E_vals.append(self.get_E(k))
                    i += 1

            E_vals.append(self.get_E(k1))

            # One tick for every high symmetry point
            self.tick_positions.append(i)
            if point == "Gamma":
                point = r"$\Gamma$"
            self.tick_labels.append(point) 

            k0 = np.copy(k1)

        self.E_vals = np.array(E_vals)

    def plot_bandstructure(self, filename: Optional[str] = None, num_bands: int = 6) -> None:
        """
        Plots the num_bands lowest bands
        and saves the plot to filename
        """
        if filename is None:
            filename = self.mat.name + "_bandstructure"

        if self.E_vals is None:
            print("No Energy Eigenvalues found, call compute_bandstructure() before plotting!")
            return

        fig, ax = plt.subplots(1,1,
                               figsize=(10,10),
                               layout="constrained",
                               )

        for band in range(num_bands):
            ax.plot(self.E_vals[:,band] / c.e)

        ax.set_xticks(self.tick_positions)
        ax.set_xticklabels(self.tick_labels)
        ax.grid()

        ax.set_ylabel("Energy in eV", fontsize=12)

        fig.savefig(filename)


if __name__ == "__main__":
    for mat_name in ["Si", "Ge", "ZnSe"]: 
        mat = fcc(mat_name)

        band = bandstructure(mat)
        band.compute_bandstructure()
        band.plot_bandstructure()