{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Twisted bilayer graphene (in-built methods)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A common way to describe a moire system is using the 4 indices defined by Hermann.\n",
    "\n",
    "Tasks:\n",
    "- Create a twisted bilayer system using the `31 30 30 31` indices and visualize the structure.\n",
    "- Calculate the electronic bandstructure and visualize the flat band.\n",
    "- Calculate the Lancoz DOS and optimize the supercell size to get clear signatures. Note that the TB model we use includes a large number of neighbors, because of this, the memory requirements are much larger than in our previous session where we only include 3 neighbors in the GBN effective model calculations. For the sake of this session, consider small cells (i.e. `SuperCell 5` that can still be calculated easily on the mem128 machines). You can go up to `SuperCell 20` if you use one of the large memory machines.\n",
    "- Compare the previous DOS with the DOS calculated by exact diagonalization. Reminder: the Lancos DOS has been renormalized by the number of atoms. To match it with the diagonalization DOS, one needs to multiply Lanczos DOS by the number of two times the number of atoms and divide it by the first nearest neighbor hopping (check `TB.Hopping` value in your input file). The energy axis is also given in units of the first nearest neighbor hopping.\n",
    "- Add Anderson disorder to your system and see how it affects the DOS.\n",
    "- Add a magnetic field and compare the DOS in presence and absence of disorder. Can you explain what you see based on your knowledge of Landau levels?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# tBG(read from external file, i.e. use ReadXYZ)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Repeat the twisted bilayer graphene calculations from the previous section using the relaxed structure given for download here.\n",
    "- Compare the results for the previous rigid configuration with the relaxed configuration observables."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Twisted bilayer graphene on hBN (bonus)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Try to combine the knowledge from previous session with the knowledge from current session to simulate the same observables for tBG on hBN. Note that for exact diagonalization calculations, a commensurate cell is required. You can avoid this complication when considering very large supercells for Lanczos where the periodic boundary effects become increasingly negligible."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Code snippets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Some routines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "\n",
    "def latexify(fig_width=None, fig_height=None, columns=1):\n",
    "    \"\"\"Set up matplotlib's RC params for LaTeX plotting.\n",
    "    Call this before plotting a figure.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    fig_width : float, optional, inches\n",
    "    fig_height : float,  optional, inches\n",
    "    columns : {1, 2}\n",
    "    \"\"\"\n",
    "\n",
    "    # code adapted from http://www.scipy.org/Cookbook/Matplotlib/LaTeX_Examples\n",
    "\n",
    "    # Width and max height in inches for IEEE journals taken from\n",
    "    # computer.org/cms/Computer.org/Journal%20templates/transactions_art_guide.pdf\n",
    "\n",
    "    assert(columns in [1,2])\n",
    "\n",
    "    if fig_width is None:\n",
    "        fig_width = 3.39 if columns==1 else 6.9 # width in inches\n",
    "\n",
    "    if fig_height is None:\n",
    "        golden_mean = (sqrt(5)-1.0)/2.0    # Aesthetic ratio\n",
    "        fig_height = fig_width*golden_mean # height in inches\n",
    "\n",
    "    #MAX_HEIGHT_INCHES = 7.0\n",
    "    #if fig_height > MAX_HEIGHT_INCHES:\n",
    "    #    print(\"WARNING: fig_height too large:\" + fig_height + \n",
    "    #          \"so will reduce to\" + MAX_HEIGHT_INCHES + \"inches.\")\n",
    "    #    fig_height = MAX_HEIGHT_INCHES\n",
    "\n",
    "    params = {'backend': 'ps',\n",
    "              'text.latex.preamble': [r'\\usepackage{gensymb}'],\n",
    "              'axes.labelsize': 8, # fontsize for x and y labels (was 10)\n",
    "              'axes.titlesize': 8,\n",
    "              'font.size': 8, # was 10\n",
    "              'legend.fontsize': 8, # was 10\n",
    "              'xtick.labelsize': 8,\n",
    "              'ytick.labelsize': 8,\n",
    "              'text.usetex': True,\n",
    "              'figure.figsize': [fig_width,fig_height],\n",
    "              'font.family': 'serif',\n",
    "              'figure.autolayout': True,\n",
    "              'lines.linewidth' : 1.0\n",
    "    }\n",
    "\n",
    "    matplotlib.rcParams.update(params)\n",
    "\n",
    "\n",
    "def format_axes(ax):\n",
    "    \n",
    "    SPINE_COLOR = 'gray'\n",
    "\n",
    "    for spine in ['top', 'right']:\n",
    "        ax.spines[spine].set_visible(False)\n",
    "\n",
    "    for spine in ['left', 'bottom']:\n",
    "        ax.spines[spine].set_color(SPINE_COLOR)\n",
    "        ax.spines[spine].set_linewidth(0.5)\n",
    "\n",
    "    ax.xaxis.set_ticks_position('bottom')\n",
    "    ax.yaxis.set_ticks_position('left')\n",
    "\n",
    "    for axis in [ax.xaxis, ax.yaxis]:\n",
    "        axis.set_tick_params(direction='out', color=SPINE_COLOR)\n",
    "\n",
    "#q_sigma = c0 * 22.2  # [nm] hopping fitting parameter\n",
    "#q_pi = a_cc * 22.2  # [nm] hopping fitting parameter\n",
    "#r0 = 0.184 * a  # [nm] hopping fitting parameter\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
    "\n",
    "\n",
    "def getBandstructure(dir1, nPlotBands, EShift):\n",
    "\n",
    "    f = open(dir1+'/generate.bands.dat','r')\n",
    "    print(f)\n",
    "    \n",
    "    \n",
    "    i = 0\n",
    "    j = 0\n",
    "    l = 0\n",
    "    \n",
    "    k = []\n",
    "    E = []\n",
    "    \n",
    "    kVec = []\n",
    "    EVec = []\n",
    "    \n",
    "    counterPlotBands = 0\n",
    "    \n",
    "    for line in f:\n",
    "        i = i+1\n",
    "        a = line.split()\n",
    "        if i==10:\n",
    "            nbands = int(a[5])\n",
    "            print(nbands)\n",
    "            nk = int(a[7])\n",
    "            print(nbands,nk)\n",
    "        if i>16:\n",
    "            if len(a)==0:\n",
    "                j = j+0.5\n",
    "                #print(j)\n",
    "            else:\n",
    "                if (int(j)==int(nbands/2-nPlotBands/2 + counterPlotBands) and int(j) < int(nbands/2+nPlotBands/2)):\n",
    "                    if l<nk:\n",
    "                        k.append(float(a[0]))\n",
    "                        E.append((float(a[1])-EShift))\n",
    "                        l = l+1\n",
    "                    if l==nk:\n",
    "                        kVec.append(k)\n",
    "                        EVec.append(E)\n",
    "                        l = 0\n",
    "                        counterPlotBands += 1\n",
    "                        k = []\n",
    "                        E = []\n",
    "                        \n",
    "    return kVec, EVec\n",
    "\n",
    "def getSingleBand3(dir1, EShift, bandIndex):\n",
    "\n",
    "    f = open(dir1+'/generate.bands.dat','r')\n",
    "    \n",
    "    \n",
    "    i = 0\n",
    "    j = 0\n",
    "    l = 0\n",
    "    \n",
    "    k = []\n",
    "    E = []\n",
    "    \n",
    "    kVec = []\n",
    "    EVec = []\n",
    "    \n",
    "    counterPlotBands = 0\n",
    "    \n",
    "    for line in f:\n",
    "        i = i+1\n",
    "        a = line.split()\n",
    "        if i==10:\n",
    "            nbands = int(a[5])\n",
    "            nk = int(a[7])\n",
    "            #print(nbands,nk)\n",
    "        if i>14:\n",
    "            if len(a)==0:\n",
    "                j = j+0.5\n",
    "                #print(j)\n",
    "            else:\n",
    "                #print(nbands)\n",
    "                if (int(j)==int(nbands/2) + bandIndex):\n",
    "                    if l<nk:\n",
    "                        k.append(float(a[0]))\n",
    "                        E.append(float(a[1])-EShift)\n",
    "                        l = l+1\n",
    "                    if l==nk:\n",
    "                        kVec.append(k)\n",
    "                        EVec.append(E)\n",
    "                        l = 0\n",
    "                        counterPlotBands += 1\n",
    "                        k = []\n",
    "                        E = []\n",
    "                        \n",
    "    return kVec, EVec\n",
    "                \n",
    "def getGap(dir1, EShift):\n",
    "    Emin = 10000000\n",
    "    Emax = -10000000\n",
    "    bandIndexVec = [-1]\n",
    "    for bandIndex in bandIndexVec:\n",
    "        kVec, EVec = getSingleBand3(dir1, EShift, bandIndex)\n",
    "        i = 0\n",
    "        if (np.max(EVec)>Emax):\n",
    "            Emax = np.max(EVec)\n",
    "        #if (np.min(EVec)<Emin):\n",
    "        #    Emin= np.min(EVec)\n",
    "    bandIndexVec = [0]\n",
    "    for bandIndex in bandIndexVec:\n",
    "        kVec, EVec = getSingleBand3(dir1, EShift, bandIndex)\n",
    "        i = 0\n",
    "        if (np.min(EVec)<Emin):\n",
    "            Emin= np.min(EVec)\n",
    "    gap = (Emin - Emax)*1000\n",
    "    return gap\n",
    "\n",
    "def getGap2(dir1, EShift):\n",
    "    Emin = 10000000\n",
    "    Emax = -10000000\n",
    "    bandIndexVec = [-3]\n",
    "    for bandIndex in bandIndexVec:\n",
    "        kVec, EVec = getSingleBand3(dir1, EShift, bandIndex)\n",
    "        i = 0\n",
    "        if (np.max(EVec)>Emax):\n",
    "            Emax = np.max(EVec)\n",
    "        #if (np.min(EVec)<Emin):\n",
    "        #    Emin= np.min(EVec)\n",
    "    bandIndexVec = [-2]\n",
    "    for bandIndex in bandIndexVec:\n",
    "        kVec, EVec = getSingleBand3(dir1, EShift, bandIndex)\n",
    "        i = 0\n",
    "        if (np.min(EVec)<Emin):\n",
    "            Emin= np.min(EVec)\n",
    "    gap = (Emin - Emax)*1000\n",
    "    return gap\n",
    "\n",
    "def getEnergyCutPlot(dir1, chosenEnergy, Nx, Ny, cutoffMin, cutoffMax, nEnergy):\n",
    "\n",
    "    os.chdir(dir1)\n",
    "    \n",
    "\n",
    "    data = np.genfromtxt(\"generate.spectral\",skip_header=4)\n",
    "    kx = data[:,0]\n",
    "    ky = data[:,1]\n",
    "    kz = data[:,2]\n",
    "    energy = (data[:,3]*3.5)-0.31\n",
    "    spectral = data[:,4]\n",
    "    \n",
    "    cutoffMin = cutoff\n",
    "    cutoffMax = cutoff\n",
    "    \n",
    "    i = 0\n",
    "    j = 0\n",
    "    eigVec = []\n",
    "    spectralVec = []\n",
    "    \n",
    "    eig = np.zeros(nEnergy)\n",
    "    spect = np.zeros(nEnergy)\n",
    "    k = np.zeros(shape=(Nx*Ny,2))\n",
    "    \n",
    "    for el1, el2, el3, el4 in zip(kx,ky,energy,spectral):\n",
    "        i = i+1\n",
    "        if (i==1):\n",
    "            k[j,0] = el1\n",
    "            k[j,1] = el2\n",
    "        if (i<=nEnergy):\n",
    "            eig[i-1] = el3\n",
    "            spect[i-1] = el4\n",
    "        if (i==nEnergy):\n",
    "            eigVec.append(eig)\n",
    "            spectralVec.append(spect)\n",
    "            i = 0\n",
    "            j = j+1\n",
    "            eig = np.zeros(nEnergy)\n",
    "            spect = np.zeros(nEnergy)\n",
    "        \n",
    "    \n",
    "    eigTrueVec = []\n",
    "    for i, kki in enumerate(k):\n",
    "        eigTrue = np.zeros(nEnergy)\n",
    "        eig = eigVec[i]\n",
    "        spec = spectralVec[i]\n",
    "        eigTrue[np.logical_and(np.logical_and(eig>chosenEnergy-cutoffMin,eig<chosenEnergy+cutoffMax),spec>0.0)] = 1.0\n",
    "        specSum = np.sum(spec[np.logical_and(np.logical_and(eig>chosenEnergy-cutoffMin,eig<chosenEnergy+cutoffMax),spec>0.0)])\n",
    "        if (eigTrue.any()==1):\n",
    "            eigTrueVec.append(specSum)\n",
    "        else:\n",
    "            eigTrueVec.append(0)\n",
    "       \n",
    "\n",
    "    return k, eigTrueVec    \n",
    "    \n",
    "K1 = np.array([1.70276024584813, 0.0000])\n",
    "K1p = np.array([1.70243417150001, 3.332186294066186E-002])\n",
    "\n",
    "\n",
    "def getEnergyCutPlot2(dir1, chosenEnergy, Nx, Ny, cutoffMin, cutoffMax, nEnergy):\n",
    "\n",
    "    os.chdir(dir1)\n",
    "    \n",
    "\n",
    "    data = np.genfromtxt(\"generate.spectral\")\n",
    "    kx = data[:,0]\n",
    "    ky = data[:,1]\n",
    "    kz = data[:,2]\n",
    "    energy = (data[:,3]*3.5)-0.31\n",
    "    spectral = data[:,4]\n",
    "    \n",
    "    cutoffMin = cutoff\n",
    "    cutoffMax = cutoff\n",
    "    \n",
    "    i = 0\n",
    "    j = 0\n",
    "    eigVec = []\n",
    "    spectralVec = []\n",
    "    \n",
    "    eig = np.zeros(nEnergy)\n",
    "    spect = np.zeros(nEnergy)\n",
    "    k = np.zeros(shape=(Nx*Ny,2))\n",
    "    \n",
    "    for el1, el2, el3, el4 in zip(kx,ky,energy,spectral):\n",
    "        i = i+1\n",
    "        if (i==1):\n",
    "            k[j,0] = el1\n",
    "            k[j,1] = el2\n",
    "        if (i<=nEnergy):\n",
    "            eig[i-1] = el3\n",
    "            spect[i-1] = el4\n",
    "        if (i==nEnergy):\n",
    "            eigVec.append(eig)\n",
    "            spectralVec.append(spect)\n",
    "            i = 0\n",
    "            j = j+1\n",
    "            eig = np.zeros(nEnergy)\n",
    "            spect = np.zeros(nEnergy)\n",
    "        \n",
    "    \n",
    "    eigTrueVec = []\n",
    "    for i, kki in enumerate(k):\n",
    "        eigTrue = np.zeros(nEnergy)\n",
    "        eig = eigVec[i]\n",
    "        spec = spectralVec[i]\n",
    "        eigTrue[np.logical_and(np.logical_and(eig>chosenEnergy-cutoffMin,eig<chosenEnergy+cutoffMax),spec>0.0)] = 1.0\n",
    "        specSum = np.sum(spec[np.logical_and(np.logical_and(eig>chosenEnergy-cutoffMin,eig<chosenEnergy+cutoffMax),spec>0.0)])\n",
    "        if (eigTrue.any()==1):\n",
    "            eigTrueVec.append(specSum)\n",
    "        else:\n",
    "            eigTrueVec.append(0)\n",
    "       \n",
    "\n",
    "    return k, eigTrueVec    \n",
    "\n",
    "def getDOS(fileName):\n",
    "    data = np.genfromtxt(fileName)\n",
    "    xdata = data[:,0]\n",
    "    ydata = data[:,1]\n",
    "    return xdata, ydata\n",
    "    \n",
    "K1 = np.array([1.70276024584813, 0.0000])\n",
    "K1p = np.array([1.70243417150001, 3.332186294066186E-002])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot the bandstructure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nPlotBands = 40\n",
    "\n",
    "ncol = 1\n",
    "nrow = 1\n",
    "\n",
    "emin = -0.2\n",
    "emax = 0.2\n",
    "\n",
    "#emin = -0.002\n",
    "#emax = 0.007\n",
    "\n",
    "latexify(columns=1)#, fig_height=1.7)  \n",
    "\n",
    "g1 = gridspec.GridSpec(nrow, ncol)#, height_ratios=[6,1])\n",
    "g1.update(wspace=0.15, hspace=0.0) # set the spacing between axes.\n",
    "\n",
    "fig, ((ax0)) = plt.subplots(nrow, ncol, sharex='col', sharey='row')\n",
    "\n",
    "ax0 = subplot(g1[0]) \n",
    "\n",
    "\n",
    "kpoint = 0.0\n",
    "gpoint = 0.032230\n",
    "mpoint = 0.060143\n",
    "kppoint = 0.076259\n",
    "\n",
    "\n",
    "\n",
    "factor = kppoint\n",
    "max1_08 = 0.078838\n",
    "max1_12 = 0.078838\n",
    "\n",
    "EShift = 0.302\n",
    "dir1 = \"/Users/nihao/github/moire_course/grabnes2/outputFiles/tBG/\"#/useReadXYZ/\"\n",
    "kVec, EVec = getBandstructure(dir1, nPlotBands, EShift)\n",
    "i = 0\n",
    "for el1, el2 in zip(kVec,EVec):\n",
    "    i = i+1\n",
    "    if (i==1):\n",
    "        #ax0.plot(el1, el2,color=\"C1\",label=r\"KC$_{VV10}$\")\n",
    "        ax0.plot(el1, el2,color=\"C0\",label=r\"GAP$_{20}$\")\n",
    "    else:\n",
    "        ax0.plot(el1, el2,color=\"C0\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot the onsite energies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.chdir(\"/Users/nihao/github/moire_course/grabnes2/inputFiles/graphene/effectiveGBN/\")\n",
    "\n",
    "data = np.genfromtxt(\"generate.pos\")\n",
    "xdata = data[:,0][::2]\n",
    "ydata = data[:,1][::2]\n",
    "\n",
    "data = np.genfromtxt(\"generate.e\")\n",
    "zdata = data[:,0][::2]\n",
    "\n",
    "plt.figure(figsize=(20,10))\n",
    "plt.scatter(xdata, ydata, c=zdata,cmap=\"gist_ncar\")\n",
    "\n",
    "for DOS visualization:\n",
    "\n",
    "os.chdir(\"/Users/nihao/github/moire_course/grabnes2/outputFiles/graphene/\")\n",
    "\n",
    "xdata, ydata = getDOS(\"generate.DOS\")\n",
    "plt.plot(xdata,ydata)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}