#!/usr/bin/env python
# coding: utf-8

# (graphene-nb)=
# # Graphene band structure
# 
# Graphene is a two-dimensional material with a honeycomb lattice structure. It has two atoms per unit cell, which we represent as two orbitals in our tight-binding model. We start by defining the lattice vectors and the coordinates of the orbitals in fractional units. These are passed to the `Lattice` class to create a lattice object, along with a list of periodic directions which will be treated with periodic boundary conditions.
# 
# :::{note}
# 
# We specify that all the lattice directions are periodic by passing `periodic_dirs=...`. This is a convenient shorthand for specifying all directions as periodic. Alternatively, we could have passed `periodic_dirs='all'`, or explicitly listed `periodic_dirs=[0, 1]`.
# 
# :::

# In[ ]:


from pythtb import TBModel, Lattice
import numpy as np
import matplotlib.pyplot as plt


# In[ ]:


# define lattice vectors
lat_vecs = [[1, 0], [1 / 2, np.sqrt(3) / 2]]
# define coordinates of orbitals
orb_vecs = [[1 / 3, 1 / 3], [2 / 3, 2 / 3]]
lat = Lattice(lat_vecs, orb_vecs, periodic_dirs=...)

# make two dimensional tight-binding graphene model
my_model = TBModel(lat)

# set model parameters
delta = 0.0
t = -1.0

# set on-site energies
my_model.set_onsite([-delta, delta])
# set hoppings (one for each connected pair of orbitals)
# (amplitude, i, j, [lattice vector to cell containing j])
my_model.set_hop(t, 0, 1, [0, 0])
my_model.set_hop(t, 1, 0, [1, 0])
my_model.set_hop(t, 1, 0, [0, 1])

print(my_model)


# ## Path in Brillouin zone from `TBModel`
# 
# We generate list of k-points following a segmented path in the BZ list of nodes (high-symmetry points) using `TBModel.k_path`.
# These high-symmetry points are specified in reduced coordinates: each entry is a list of fractional coordinates with respect to the reciprocal lattice vectors. We choose the nodes:
# 
# - $\Gamma$  = `[0, 0]`
# - $K$       = `[2/3, 1/3]`
# - $M$       = `[1/2, 1/2]`
# - $\Gamma$  = `[0, 0]`
# 
# Outputs:
# - `k_path`: list of interpolated k-points
# - `k_dist`: horizontal axis position of each k-point in the list
# - `k_node_dist`: horizontal axis position of each original node

# In[ ]:


k_nodes = [[0, 0], [2 / 3, 1 / 3], [1 / 2, 1 / 2], [0, 0]]
k_node_labels = (r"$\Gamma $", r"$K$", r"$M$", r"$\Gamma $")
nk = 121

k_path, k_dist, k_node_dist = my_model.k_path(k_nodes, nk)


# ## Band structure
# 
# Graphene has a characteristic linear band crossing at the K point in the Brillouin zone, known as a Dirac cone. This results in unique electronic properties, such as high electron mobility and massless charge carriers. The band structure can be calculated using the tight-binding model and visualized along high-symmetry paths in the Brillouin zone.
# 
# We diagonalize the Hamiltonian at each k-point to obtain the energy eigenvalues, which represent the allowed energy levels for electrons in the material. Plotting these energy levels against the k-points gives us the band structure of graphene, revealing its electronic properties.

# In[ ]:


evals = my_model.solve_ham(k_path)


# In[ ]:


fig, ax = plt.subplots()

ax.set_xlim(k_node_dist[0], k_node_dist[-1])
ax.set_xticks(k_node_dist)
ax.set_xticklabels(k_node_labels)

for n in range(len(k_node_dist)):
    ax.axvline(x=k_node_dist[n], linewidth=0.5, color="k")

ax.set_title("Graphene band structure")
ax.set_xlabel("Path in k-space")
ax.set_ylabel("Band energy")

# plot bands
ax.plot(k_dist, evals, c="b")
plt.show()


# Alternatively, we can use the `TBModel.band_structure` method to compute the band structure directly along the specified path without having to generate the k-points separately or specify the matplotlib details. This method takes care of computing the k-points, solving the Hamiltonian, and plotting the results in one step.

# In[ ]:


my_model.plot_bands(k_nodes, k_node_labels)
plt.show()