"""
.. module:: Supercell
    :synopsis: Framework for constructing supercells.
.. moduleauthor:: D. Wang <dwang5@zoho.com>
"""
from math import atan
import numpy as np
[docs]class Supercell(object):
    """
    Args:
        n1,n2,n3 are the number of repeated unitcells along the three directions of
        the 'lattice'.
        lattice specifies the Bravais lattice of a **unit cell**.
 .. note::
    This version is not compitable with previous ones as
    * All the index starts from 0.
    * Arrays (e.g. lattice and reci lattice) are row majored.
    The reason is to use Pybind11 that binds to C++ instead of Fortran.
    """
    def __init__(self,n1,n2,n3,lattice):
        """
        Define most essential components of a supercell.
        """
        self.n1 = n1
        self.n2 = n2
        self.n3 = n3
        self.nsites = n1*n2*n3
        self.nxy = n1*n2
        #Lattice vector (Bravais vector)
        self.a = np.zeros((3,3))
        self.b = np.zeros((3,3))
        self.lattice = np.zeros((3,3))
        self.lattice = lattice
        self.set_supercell()
        self.celvol = self.volcell()
        #Calculate reciprocal lattice.
        self.rec_lat()
[docs]    def set_supercell(self):
        """
        Multiply the n1,n2,n3 to the unitcell vectors to form the superlattice.
        """
        self.a[0, :] = self.n1 * self.lattice[0, :]
        self.a[1, :] = self.n2 * self.lattice[1, :]
        self.a[2, :] = self.n3 * self.lattice[2, :]
        #'XYPOINTS'
        self.ix = np.zeros(self.nxy,dtype=np.int)
        self.iy = np.zeros(self.nxy,dtype=np.int)
        self.site = np.zeros((self.n1,self.n2),dtype=np.int)
        nxy = 0
        for i in range(self.n1):
            for j in range(self.n2):
                self.ix[nxy] = i
                self.iy[nxy] = j
                nxy += 1
                self.site[i,j] = nxy
        # Match ia with ix, iy, and iz.
        # This is the convention in all the following.
        self.iaa = np.zeros((self.n1,self.n2,self.n3),dtype=np.int)
        self.ixa = np.zeros(self.nsites,dtype=np.int)
        self.iya = np.zeros(self.nsites,dtype=np.int)
        self.iza = np.zeros(self.nsites,dtype=np.int)
        l = 0
        for k in range(self.n3):
            for j in range(self.n2):
                for i in range(self.n1):
                    self.iaa[i,j,k] = l
                    self.ixa[l] = i
                    self.iya[l] = j
                    self.iza[l] = k
                    l +=1 
        #for i in range(self.nsites):
        #    print(i,self.ixa[i],self.iya[i],self.iza[i])
[docs]    def volcell(self):
        """
        Volume of the supercell.
        """
        return np.inner(self.a[0,:], np.cross(self.a[1,:], self.a[2,:])) 
[docs]    def rec_lat(self):
        """
        Calculates reciprocal lattice vectors.
        The convention follows "Solid state physics" by
        Lu Dong, Jiang Ping ... Eq. (2.2.2)
        """
        pi = 4.0 * atan(1.0)
        self.b[0,:] = np.cross(self.a[1,:],self.a[2,:])*2*pi/self.celvol
        self.b[1,:] = np.cross(self.a[2,:],self.a[0,:])*2*pi/self.celvol
        self.b[2,:] = np.cross(self.a[0,:],self.a[1,:])*2*pi/self.celvol