In [1]:
import pythtb as tb
import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import minimize

class Orbital(object):
    """ Orbital object with a onsite energy """
    def __init__(self, onsite, label):
        """
        Initialize an instance of orbital
        
        :type  onsite: float
        :param onsite: The onsite energy of the orbital
        
        :type  label: string
        :param label: The name of the orbital (examples: s,p,d,f)
        """
        self.label = label
        self.onsite = onsite

    def __repr__(self):
        """ Representation of the orbital object """
        return "{} with onsite {}".format(self.label, self.onsite)
        
class Atom(object):
    """ Atom object hosting orbitals """
    def __init__(self, position, number):
        """ Initialize an instance of atom
        
        :type  position: ndarray, list, tuple
        :param position: The onsite energy of the orbital
        
        :type  number: int
        :param number: The atomic number of the atom (Z)
        """
        position = np.asarray(position)
        self.number = number
        self.position = position
        self.orbitals = []
    
    def add_orbital(self, orbital):
        """ Add an orbital to the hosting atom
        :type  orbital: orbital object
        :param orbital: an orbital in the atom
        """
        self.orbitals.append(orbital)
        
    def number_of_orbitals(self):
        """ Returns the number of orbitals in the atom 
        :returns: the number of orbitals in the atom
        """
        return len(self.orbitals)     

    
    def __repr__(self):
        """ Representation of the Atom object """

        return "Atom with Z={} at {}".format(self.number, self.position)
    
class Oxygen(Atom):
    """ A special case of Atom designed for Oxygen"""
    def __init__(self, position):
        Atom.__init__(self, position, 8)
        self.add_orbital(Orbital(-19.046, "S"))
        self.add_orbital(Orbital(  4.142, "Pz"))
        self.add_orbital(Orbital(  4.142, "Px"))
        self.add_orbital(Orbital(  4.142, "Py"))
    
    def __repr__(self):
        return "Oxygen at {}".format(self.position)
    
class Zinc(Atom):
    """ A special case of Atom designed for Zinc"""
    def __init__(self, position):
        
        Atom.__init__(self, position, 30)
        self.add_orbital(Orbital(  1.666, "S"))
        self.add_orbital(Orbital( 12.368, "Pz"))
        self.add_orbital(Orbital( 12.368, "Px"))
        self.add_orbital(Orbital( 12.368, "Py"))
        
    def __repr__(self):
        return "Zinc at {}".format(self.position)
    

class Gallium(Atom):
    """ A special case of Atom designed for Gallium"""
    def __init__(self, position):
        
        Atom.__init__(self, position, 30)
        self.add_orbital(Orbital(  1.438, "S"))
        self.add_orbital(Orbital( 10.896, "Pz"))
        self.add_orbital(Orbital( 10.896, "Px"))
        self.add_orbital(Orbital( 10.896, "Py"))
    
    def __repr__(self):
        return "Gallium at {}".format(self.position)
    
class Nitrogen(Atom):
    """ A special case of Atom designed for Nitrogen"""
    def __init__(self, position):
        
        Atom.__init__(self, position, 30)
        self.add_orbital(Orbital(-11.012, "S"))
        self.add_orbital(Orbital(  0.005, "Pz"))
        self.add_orbital(Orbital(  0.005, "Px"))
        self.add_orbital(Orbital(  0.005, "Py"))
        
    def __repr__(self):
        return "Nitrogen at {}".format(self.position)

In [2]:
class Crystal:
    
    def __init__(self, lattice):
        """ Initialize an instance of crystal with lattice hosting atom objects
        
        :type  lattice: ndarray
        :param lattice: a 3x3 array with lattice parameters [a,b,c]
        """
        self.lattice = lattice
        self.atoms = []
        
    def add_atom(self,atom):
        """ Add an atom object to the crystal
        :type  atom: atom object
        :param atom: the instance of atom to be placed in the crystal
        """
        existing = False
        for existing_atom in self.atoms:
            if np.all(atom.position == existing_atom.position):
                existing = True
        if not existing:
            self.atoms.append(atom)
            return "Placed atom at {}".format(atom.position)
        else:
            raise Warning("An atom is already at").format(atom.position)
            
    def __repr__(self):
        """ Representation of the crystal object """
        string = "Lattice:\n{}\nAtoms:\n".format(self.lattice)
        
        for i, atom in enumerate(self.atoms):
            string += "{}: {}\n".format(i,atom)
        return string

In [3]:
class TransitionMatrix:
    """ A matrix containing all hopping parameters between to atoms"""
    def __init__(self, initial_atom, final_atom, equivalent_final_cells=[[0, 0, 0]]):
        """ Create an instance of TransitionMatrix based on an initial and final atom
        
        :type  initial_atom: atom object
        :param initial_atom: the initial atom in the transition matrix
        
        :type  final_atom: atom object
        :param final_atom: the final atom in the transition matrix
        
        :type  equivalent_final_cells: list
        :param equivalent_final_cells: list of int-coordinates to neighbouring cells
        """
        
        self.set_transitions( 
            np.zeros((
                initial_atom.number_of_orbitals(),
                final_atom.number_of_orbitals())
            ) 
        )
        
        self.set_equivalent_final_cells(equivalent_final_cells)
        self.initial_atom = initial_atom
        self.final_atom = final_atom
        
        if initial_atom == final_atom:
            self.label = "Onsite element at {}\n".format(initial_atom)
        else:
            self.label = "From {} to {}\n".format(initial_atom,final_atom)
            
    def set_transitions(self, transition_matrix):
        """ Replace the full transition matrix
        
        :type  transition_matrix: ndarray
        :param transition_matrix: the new transition matrix
        """
        self.transitions = transition_matrix
        
    def set_equivalent_final_cells(self, equivalent_final_cells):
        """ set the list of equivalent final cells
        
        :type  equivalent_final_cells: list
        :param equivalent_final_cells: the new list of int-coordinates to neighbouring cells
        """
        self.equivalent_final_cells = equivalent_final_cells
            
    def __getitem__(self, index):
        """ Get the hopping parameter by the index of initial and final orbital [i][f]"""
        return self.transitions[index]
    
    def __repr__(self):
        """ Represent the transition matrix object"""
        string = self.label
    
        for final_orbital in self.final_atom.orbitals:
            string += "\t|{}>".format(final_orbital.label)
        string += "\n"
        for i, initial_orbital in enumerate(self.initial_atom.orbitals):
            string += "<{}|\t".format(initial_orbital.label)
            
            for f, final_orbital in enumerate(self.final_atom.orbitals):
                string += "{}\t".format(self.transitions.round(2)[i,f])
            string += "\n"
        string += "\n"
        string += "Equivalent cells: {}".format(self.equivalent_final_cells)
        return string

In [4]:
class AtomMatrix:
    """ A matrix containing all transition matrices between atoms """
    def __init__(self, crystal):
        """ Create instance of atom matrix
        :type  crystal: crystal object
        :param crystal: a crystal containing atoms
        """
        self.atom_matrix = []
        
        for i,initial_atom in enumerate(crystal.atoms):
            atom_matrix_row = []    

            for final_atom in crystal.atoms:    
                transition_matrix = TransitionMatrix(initial_atom, final_atom)

                if initial_atom == final_atom:
                    for i, orbital in enumerate(initial_atom.orbitals):
                        transition_matrix[i][i] = orbital.onsite
                        
                atom_matrix_row.append(transition_matrix)
                
            self.atom_matrix.append(atom_matrix_row)

    def __getitem__(self, index):
        """ Get a transition matrix by the index of initial and final atom [i][f]"""
        return self.atom_matrix[index]
                
    def __repr__(self):
        """ Representation of the atom matrix object"""
        string = "Atom matrix:\n"
        for i, atom_matrix_row in enumerate(self.atom_matrix):
            for f, transition_matrix in enumerate(atom_matrix_row):
                string += "[{},{}]".format(i,f)
            string += "\n"
        string += "\n"
        for i, atom_matrix_row in enumerate(self.atom_matrix):
            for f, transition_matrix in enumerate(atom_matrix_row[i:]):
                string += "[{},{}]\n{}\n\n".format(i,f,transition_matrix)
        return string

In [58]:
class TightBinding:
    """ Tight binding class constructed around the pythTB package """
    def __init__(self, crystal):
        """ Create instance of the Tight binding model 
        
        :type  crystal: crystal object
        :param crystal: a crystal object containing atoms
        """
        self.crystal = crystal
        
        self._orbital_positons = []
        for atom in crystal.atoms:
            for orbital in atom.orbitals:
                self._orbital_positons.append(atom.position)
        
        
        self.model = tb.tb_model(3,3,self.crystal.lattice, self._orbital_positons)
        
        
    """#############################################################333"""
        
    def get_hopping_parameters(self):
        """ Get a list of all hopping parameters in the model"""
        return [self.Vss, self.Vxx, self.Vxy, self.Vsapc, self.Vpasc ]
        
        
    def set_hopping_parameters(self, Vss, Vxx, Vxy, Vsapc, Vpasc):
        """ Set all hopping parameters in the model
        
        :type  Vss: float
        :param Vss: The hopping parameter from s to s
        
        :type  Vxx: float
        :param Vxx: The hopping parameter from px of the anion to px of the cation
        
        :type  Vxy: float
        :param Vxy: The hopping parameter from px of the anion to py of the cation
        
        :type  Vsapc: float
        :param Vsapc: The hopping parameter from s of the anion to p of the cation
        
        :type  Vpasc: float
        :param Vpasc: The hopping parameter from p of the anion to s of the cation
        """
        self.Vss = Vss
        self.Vxx = Vxx
        self.Vxy = Vxy
        self.Vsapc = Vsapc
        self.Vpasc = Vpasc
        
        """###############     Vertical bonding system    ##################"""      
        self.UVss = 0.25*Vss
        self.UVzz = 0.25*(Vxx+2*Vxy)
        self.UVxx = 0.25*(Vxx-Vxy)
        self.UVsz = -0.25*np.sqrt(3)*Vsapc
        self.UVzs =  0.25*np.sqrt(3)*Vpasc
        
        
        """###############    Horizontal bonding system   ##################"""
        self.UHss = self.UVss
        
        self.UHyy = self.UVxx
        self.UHzz = 1/9 * (8*self.UVxx +   self.UVzz)
        self.UHxx = 1/9 * (  self.UVxx + 8*self.UVzz)
        
        self.UHsz = -1/3 * self.UVsz
        self.UHzs = -1/3 * self.UVzs
        
        self.UHsx = -2*np.sqrt(2)/3 * self.UVsz
        self.UHxs = -2*np.sqrt(2)/3 * self.UVzs
        
        self.UHzx =  2*np.sqrt(2)/9 * (self.UVzz-self.UVxx)
        self.UHxz =  2*np.sqrt(2)/9 * (self.UVzz-self.UVxx)

        """##############################################################"""

        """  ONSITE  """

        for i, initial_atom in enumerate(self.crystal.atoms):
            for io, initial_orbital in enumerate(initial_atom.orbitals):
                self.model.set_onsite(initial_orbital.onsite, ind_i=(i*len(self.crystal.atoms)+io), mode='reset')
        


        """  HOPPING  """
                                                 
        self._M03(0,3)
        self._M12(1,2)
        
        self._M13(1,3)
        self._M02(0,2)
                
        
    def _M03(self, i, f):
        """ The M14 transition matrix designed by Kobayashi et Al."""
        self.model.set_hop(self.UVss, ind_i=i*4+0, ind_j=f*4+0, ind_R=[0, 0, -1], mode='reset')
        self.model.set_hop(self.UVsz, ind_i=i*4+0, ind_j=f*4+1, ind_R=[0, 0, -1], mode='reset')
        self.model.set_hop(self.UVzs, ind_i=i*4+1, ind_j=f*4+0, ind_R=[0, 0, -1], mode='reset')
        self.model.set_hop(self.UVzz, ind_i=i*4+1, ind_j=f*4+1, ind_R=[0, 0, -1], mode='reset')
        self.model.set_hop(self.UVxx, ind_i=i*4+2, ind_j=f*4+2, ind_R=[0, 0, -1], mode='reset')
        self.model.set_hop(self.UVxx, ind_i=i*4+3, ind_j=f*4+3, ind_R=[0, 0, -1], mode='reset')            
        
    def _M12(self, i, f):
        """ The M14 transition matrix designed by Kobayashi et Al."""
        self.model.set_hop(self.UVss, ind_i=i*4+0, ind_j=f*4+0, ind_R=[0, 0, 0], mode='reset')
        self.model.set_hop(self.UVsz, ind_i=i*4+0, ind_j=f*4+1, ind_R=[0, 0, 0], mode='reset')
        self.model.set_hop(self.UVzs, ind_i=i*4+1, ind_j=f*4+0, ind_R=[0, 0, 0], mode='reset')
        self.model.set_hop(self.UVzz, ind_i=i*4+1, ind_j=f*4+1, ind_R=[0, 0, 0], mode='reset')
        self.model.set_hop(self.UVxx, ind_i=i*4+2, ind_j=f*4+2, ind_R=[0, 0, 0], mode='reset')
        self.model.set_hop(self.UVxx, ind_i=i*4+3, ind_j=f*4+3, ind_R=[0, 0, 0], mode='reset')  
              
    def _M02(self, i, f):
        """ The M13 transition matrix designed by Kobayashi et Al."""
        #s-s
        self.model.set_hop(self.UHss, ind_i=i*4+0, ind_j=f*4+0, ind_R=[0, 0, 0], mode='reset')
        self.model.set_hop(self.UHss, ind_i=i*4+0, ind_j=f*4+0, ind_R=[-1, 0, 0], mode='reset')
        self.model.set_hop(self.UHss, ind_i=i*4+0, ind_j=f*4+0, ind_R=[-1, -1, 0], mode='reset')
        #s-z
        self.model.set_hop(self.UHsz, ind_i=i*4+0, ind_j=f*4+1, ind_R=[0, 0, 0], mode='reset')
        self.model.set_hop(self.UHsz, ind_i=i*4+0, ind_j=f*4+1, ind_R=[-1, 0, 0], mode='reset')
        self.model.set_hop(self.UHsz, ind_i=i*4+0, ind_j=f*4+1, ind_R=[-1, -1, 0], mode='reset')
        #s-x
        self.model.set_hop(     self.UHsx, ind_i=i*4+0, ind_j=f*4+2, ind_R=[0, 0, 0], mode='reset')
        self.model.set_hop(-0.5*self.UHsx, ind_i=i*4+0, ind_j=f*4+2, ind_R=[-1, 0, 0], mode='reset')
        self.model.set_hop(-0.5*self.UHsx, ind_i=i*4+0, ind_j=f*4+2, ind_R=[-1, -1, 0], mode='reset')
        #s-y
        self.model.set_hop( np.sqrt(3)/2*self.UHsx, ind_i=i*4+0, ind_j=f*4+3, ind_R=[-1, 0, 0], mode='reset')
        self.model.set_hop(-np.sqrt(3)/2*self.UHsx, ind_i=i*4+0, ind_j=f*4+3, ind_R=[-1, -1, 0], mode='reset')
        """##########################################"""
        #z-s
        self.model.set_hop(self.UHzs, ind_i=i*4+1, ind_j=f*4+0, ind_R=[0, 0, 0], mode='reset')
        self.model.set_hop(self.UHzs, ind_i=i*4+1, ind_j=f*4+0, ind_R=[-1, 0, 0], mode='reset')
        self.model.set_hop(self.UHzs, ind_i=i*4+1, ind_j=f*4+0, ind_R=[-1, -1, 0], mode='reset')
        #z-z
        self.model.set_hop(self.UHzz, ind_i=i*4+1, ind_j=f*4+1, ind_R=[0, 0, 0], mode='reset')
        self.model.set_hop(self.UHzz, ind_i=i*4+1, ind_j=f*4+1, ind_R=[-1, 0, 0], mode='reset')
        self.model.set_hop(self.UHzz, ind_i=i*4+1, ind_j=f*4+1, ind_R=[-1, -1, 0], mode='reset')
        #z-x
        self.model.set_hop(     self.UHzx, ind_i=i*4+1, ind_j=f*4+2, ind_R=[0, 0, 0], mode='reset')
        self.model.set_hop(-0.5*self.UHzx, ind_i=i*4+1, ind_j=f*4+2, ind_R=[-1, 0, 0], mode='reset')
        self.model.set_hop(-0.5*self.UHzx, ind_i=i*4+1, ind_j=f*4+2, ind_R=[-1, -1, 0], mode='reset')
        #z-y
        self.model.set_hop( np.sqrt(3)/2*self.UHzx, ind_i=i*4+1, ind_j=f*4+3, ind_R=[-1, 0, 0], mode='reset')
        self.model.set_hop(-np.sqrt(3)/2*self.UHzx, ind_i=i*4+1, ind_j=f*4+3, ind_R=[-1, -1, 0], mode='reset')
        """#########################################"""
        #x-s
        self.model.set_hop(     self.UHxs, ind_i=i*4+2, ind_j=f*4+0, ind_R=[0, 0, 0], mode='reset')
        self.model.set_hop(-0.5*self.UHxs, ind_i=i*4+2, ind_j=f*4+0, ind_R=[-1, 0, 0], mode='reset')
        self.model.set_hop(-0.5*self.UHxs, ind_i=i*4+2, ind_j=f*4+0, ind_R=[-1, -1, 0], mode='reset')
        #x-z
        self.model.set_hop(     self.UHxz, ind_i=i*4+2, ind_j=f*4+1, ind_R=[0, 0, 0], mode='reset')
        self.model.set_hop(-0.5*self.UHxz, ind_i=i*4+2, ind_j=f*4+1, ind_R=[-1, 0, 0], mode='reset')
        self.model.set_hop(-0.5*self.UHxz, ind_i=i*4+2, ind_j=f*4+1, ind_R=[-1, -1, 0], mode='reset')
        #x-x
        self.model.set_hop(     self.UHxx, ind_i=i*4+2, ind_j=f*4+2, ind_R=[0, 0, 0], mode='reset')
        self.model.set_hop(-0.5*self.UHxx, ind_i=i*4+2, ind_j=f*4+2, ind_R=[-1, 0, 0], mode='reset')
        self.model.set_hop(-0.5*self.UHxx, ind_i=i*4+2, ind_j=f*4+2, ind_R=[-1, -1, 0], mode='reset')
        self.model.set_hop(3/4*(self.UHxx+self.UHyy), ind_i=i*4+2, ind_j=f*4+2, ind_R=[-1, 0, 0], mode='add')
        self.model.set_hop(3/4*(self.UHxx+self.UHyy), ind_i=i*4+2, ind_j=f*4+2, ind_R=[-1, -1, 0], mode='add')        
        #x-y
        self.model.set_hop(-np.sqrt(3)/4*(self.UHxx-self.UHyy), ind_i=i*4+2, ind_j=f*4+3, ind_R=[-1, 0, 0], mode='reset')
        self.model.set_hop( np.sqrt(3)/4*(self.UHxx-self.UHyy), ind_i=i*4+2, ind_j=f*4+3, ind_R=[-1, -1, 0], mode='reset')
        """###########################################"""
        #y-s
        self.model.set_hop( np.sqrt(3)/2*self.UHxs, ind_i=i*4+3, ind_j=f*4+0, ind_R=[-1, 0, 0], mode='reset')
        self.model.set_hop(-np.sqrt(3)/2*self.UHxs, ind_i=i*4+3, ind_j=f*4+0, ind_R=[-1, -1, 0], mode='reset')
        #y-z
        self.model.set_hop( np.sqrt(3)/2*self.UHxz, ind_i=i*4+3, ind_j=f*4+1, ind_R=[-1, 0, 0], mode='reset')
        self.model.set_hop(-np.sqrt(3)/2*self.UHxz, ind_i=i*4+3, ind_j=f*4+1, ind_R=[-1, -1, 0], mode='reset')
        #y-x
        self.model.set_hop(-np.sqrt(3)/4*(self.UHxx-self.UHyy), ind_i=i*4+3, ind_j=f*4+2, ind_R=[-1, 0, 0], mode='reset')
        self.model.set_hop( np.sqrt(3)/4*(self.UHxx-self.UHyy), ind_i=i*4+3, ind_j=f*4+2, ind_R=[-1, -1, 0], mode='reset')  
        #y-y
        self.model.set_hop(     self.UHyy, ind_i=i*4+3, ind_j=f*4+3, ind_R=[0, 0, 0], mode='reset')
        self.model.set_hop(-0.5*self.UHyy, ind_i=i*4+3, ind_j=f*4+3, ind_R=[-1, 0, 0], mode='reset')
        self.model.set_hop(-0.5*self.UHyy, ind_i=i*4+3, ind_j=f*4+3, ind_R=[-1, -1, 0], mode='reset')
        self.model.set_hop(3/4*(self.UHxx+self.UHyy), ind_i=i*4+3, ind_j=f*4+3, ind_R=[-1, 0, 0], mode='add')
        self.model.set_hop(3/4*(self.UHxx+self.UHyy), ind_i=i*4+3, ind_j=f*4+3, ind_R=[-1, -1, 0], mode='add')    

    def _M13(self, i, f):
        """ The M24 transition matrix designed by Kobayashi et Al."""
        #s-s
        self.model.set_hop(self.UHss, ind_i=i*4+0, ind_j=f*4+0, ind_R=[0, 0, 0], mode='reset')
        self.model.set_hop(self.UHss, ind_i=i*4+0, ind_j=f*4+0, ind_R=[1, 0, 0], mode='reset')
        self.model.set_hop(self.UHss, ind_i=i*4+0, ind_j=f*4+0, ind_R=[1, 1, 0], mode='reset')
        #s-z
        self.model.set_hop(self.UHsz, ind_i=i*4+0, ind_j=f*4+1, ind_R=[0, 0, 0], mode='reset')
        self.model.set_hop(self.UHsz, ind_i=i*4+0, ind_j=f*4+1, ind_R=[1, 0, 0], mode='reset')
        self.model.set_hop(self.UHsz, ind_i=i*4+0, ind_j=f*4+1, ind_R=[1, 1, 0], mode='reset')
        #s-x
        self.model.set_hop(   -self.UHsx, ind_i=i*4+0, ind_j=f*4+2, ind_R=[0, 0, 0], mode='reset')
        self.model.set_hop(0.5*self.UHsx, ind_i=i*4+0, ind_j=f*4+2, ind_R=[1, 0, 0], mode='reset')
        self.model.set_hop(0.5*self.UHsx, ind_i=i*4+0, ind_j=f*4+2, ind_R=[1, 1, 0], mode='reset')
        #s-y
        self.model.set_hop(-np.sqrt(3)/2*self.UHsx, ind_i=i*4+0, ind_j=f*4+3, ind_R=[1, 0, 0], mode='reset')
        self.model.set_hop( np.sqrt(3)/2*self.UHsx, ind_i=i*4+0, ind_j=f*4+3, ind_R=[1, 1, 0], mode='reset')
        """##########################################"""
        #z-s
        self.model.set_hop(self.UHzs, ind_i=i*4+1, ind_j=f*4+0, ind_R=[0, 0, 0], mode='reset')
        self.model.set_hop(self.UHzs, ind_i=i*4+1, ind_j=f*4+0, ind_R=[1, 0, 0], mode='reset')
        self.model.set_hop(self.UHzs, ind_i=i*4+1, ind_j=f*4+0, ind_R=[1, 1, 0], mode='reset')
        #z-z
        self.model.set_hop(self.UHzz, ind_i=i*4+1, ind_j=f*4+1, ind_R=[0, 0, 0], mode='reset')
        self.model.set_hop(self.UHzz, ind_i=i*4+1, ind_j=f*4+1, ind_R=[1, 0, 0], mode='reset')
        self.model.set_hop(self.UHzz, ind_i=i*4+1, ind_j=f*4+1, ind_R=[1, 1, 0], mode='reset')
        #z-x
        self.model.set_hop(   -self.UHzx, ind_i=i*4+1, ind_j=f*4+2, ind_R=[0, 0, 0], mode='reset')
        self.model.set_hop(0.5*self.UHzx, ind_i=i*4+1, ind_j=f*4+2, ind_R=[1, 0, 0], mode='reset')
        self.model.set_hop(0.5*self.UHzx, ind_i=i*4+1, ind_j=f*4+2, ind_R=[1, 1, 0], mode='reset')
        #z-y
        self.model.set_hop(-np.sqrt(3)/2*self.UHzx, ind_i=i*4+1, ind_j=f*4+3, ind_R=[1, 0, 0], mode='reset')
        self.model.set_hop( np.sqrt(3)/2*self.UHzx, ind_i=i*4+1, ind_j=f*4+3, ind_R=[1, 1, 0], mode='reset')
        """#########################################"""
        #x-s
        self.model.set_hop(   -self.UHxs, ind_i=i*4+2, ind_j=f*4+0, ind_R=[0, 0, 0], mode='reset')
        self.model.set_hop(0.5*self.UHxs, ind_i=i*4+2, ind_j=f*4+0, ind_R=[1, 0, 0], mode='reset')
        self.model.set_hop(0.5*self.UHxs, ind_i=i*4+2, ind_j=f*4+0, ind_R=[1, 1, 0], mode='reset')
        #x-z
        self.model.set_hop(   -self.UHxz, ind_i=i*4+2, ind_j=f*4+1, ind_R=[0, 0, 0], mode='reset')
        self.model.set_hop(0.5*self.UHxz, ind_i=i*4+2, ind_j=f*4+1, ind_R=[1, 0, 0], mode='reset')
        self.model.set_hop(0.5*self.UHxz, ind_i=i*4+2, ind_j=f*4+1, ind_R=[1, 1, 0], mode='reset')
        #x-x
        self.model.set_hop(     self.UHxx, ind_i=i*4+2, ind_j=f*4+2, ind_R=[0, 0, 0], mode='reset')
        self.model.set_hop(-0.5*self.UHxx, ind_i=i*4+2, ind_j=f*4+2, ind_R=[1, 0, 0], mode='reset')
        self.model.set_hop(-0.5*self.UHxx, ind_i=i*4+2, ind_j=f*4+2, ind_R=[1, 1, 0], mode='reset')
        self.model.set_hop(3/4*(self.UHxx+self.UHyy), ind_i=i*4+2, ind_j=f*4+2, ind_R=[1, 0, 0], mode='add')
        self.model.set_hop(3/4*(self.UHxx+self.UHyy), ind_i=i*4+2, ind_j=f*4+2, ind_R=[1, 1, 0], mode='add')        
        #x-y
        self.model.set_hop(-np.sqrt(3)/4*(self.UHxx-self.UHyy), ind_i=i*4+2, ind_j=f*4+3, ind_R=[1, 0, 0], mode='reset')
        self.model.set_hop( np.sqrt(3)/4*(self.UHxx-self.UHyy), ind_i=i*4+2, ind_j=f*4+3, ind_R=[1, 1, 0], mode='reset')
        """###########################################"""
        #y-s
        self.model.set_hop(-np.sqrt(3)/2*self.UHxs, ind_i=i*4+3, ind_j=f*4+0, ind_R=[1, 0, 0], mode='reset')
        self.model.set_hop( np.sqrt(3)/2*self.UHxs, ind_i=i*4+3, ind_j=f*4+0, ind_R=[1, 1, 0], mode='reset')
        #y-z
        self.model.set_hop(-np.sqrt(3)/2*self.UHxz, ind_i=i*4+3, ind_j=f*4+1, ind_R=[1, 0, 0], mode='reset')
        self.model.set_hop( np.sqrt(3)/2*self.UHxz, ind_i=i*4+3, ind_j=f*4+1, ind_R=[1, 1, 0], mode='reset')
        #y-x
        self.model.set_hop(-np.sqrt(3)/4*(self.UHxx-self.UHyy), ind_i=i*4+3, ind_j=f*4+2, ind_R=[1, 0, 0], mode='reset')
        self.model.set_hop( np.sqrt(3)/4*(self.UHxx-self.UHyy), ind_i=i*4+3, ind_j=f*4+2, ind_R=[1, 1, 0], mode='reset')  
        #y-y
        self.model.set_hop(     self.UHyy, ind_i=i*4+3, ind_j=f*4+3, ind_R=[0, 0, 0], mode='reset')
        self.model.set_hop(-0.5*self.UHyy, ind_i=i*4+3, ind_j=f*4+3, ind_R=[1, 0, 0], mode='reset')
        self.model.set_hop(-0.5*self.UHyy, ind_i=i*4+3, ind_j=f*4+3, ind_R=[1, 1, 0], mode='reset')
        self.model.set_hop(3/4*(self.UHxx+self.UHyy), ind_i=i*4+3, ind_j=f*4+3, ind_R=[1, 0, 0], mode='add')
        self.model.set_hop(3/4*(self.UHxx+self.UHyy), ind_i=i*4+3, ind_j=f*4+3, ind_R=[1, 1, 0], mode='add')        
    
    def f0(self,conjugate=False):
    
        conjugate = -2*conjugate+1

        f = np.array([
            [ 0,  0, 0],
            [-1,  0, 0],
            [-1, -1, 0]
        ])*conjugate

        w = np.array([
            1, 
            1, 
            1
        ])
        return f,w
    
    def f1(self,conjugate=False):
    
        conjugate = -2*conjugate+1

        f = np.array([
            [ 0,  0, 0],
            [-1,  0, 0],
            [-1, -1, 0]
        ])*conjugate

        w = np.array([
            1, 
            -1/2, 
            -1/2
        ])
        return f,w
    
    def f2(self,conjugate=False):
    
        conjugate = -2*conjugate+1

        f = np.array([
            [ 0,  0, 0],
        ])*conjugate

        w = np.array([
            1
        ])
        return f,w
    
    def fplus(self,conjugate=False):
    
        conjugate = -2*conjugate+1

        f = np.array([
            [-1,  0, 0],
            [-1, -1, 0],
        ])*conjugate

        w = np.array([
            1, 
            1
        ])
        return f,w
    
    
    def fminus(self,conjugate=False):
    
        conjugate = -2*conjugate+1

        f = np.array([
            [-1,  0, 0],
            [-1, -1, 0],
        ])*conjugate

        w = np.array([
            1, 
            -1
        ])
        return f,w
    

    """#####################################################################"""

    """    
    def update(self):
        index_initial = 0
        for i, initial_atom in enumerate(self.crystal.atoms):
            index_final = index_initial
            for f, final_atom in enumerate(self.crystal.atoms[i:]):
                f = f+i

                for io, initial_orbital in enumerate(initial_atom.orbitals):
                    for fo, final_orbital in enumerate(initial_atom.orbitals):

                    # Set onsite elements
                        if initial_atom == final_atom:
                            if initial_orbital == final_orbital:
                                self.model.set_onsite(self.atom_matrix[i][f][io][fo], ind_i=(index_initial+io), mode='reset')
        
                    # Set hopping parameters
                        else:
                            for final_cell in self.atom_matrix[i][f].equivalent_final_cells:
                                self.model.set_hop(self.atom_matrix[i][f][io][fo], ind_i=(index_initial+io), ind_j=(index_final+fo), ind_R=final_cell, mode='reset')

                
                index_final += final_atom.number_of_orbitals()    
            index_initial += initial_atom.number_of_orbitals()
            """
            
    def display_pythtb(self):
        """ Displat the info from pythTB """
        self.model.display()
        
    def calculate(self, k_grid, eig_vectors = False):
        """ Calculate band energies for the given k-points 
        
        :type  k_grid: ndarray, list
        :param k_grid: an array/list of 3D k-points
        
        :type  eig_vectors: bool
        :param eig_vectors: if eigen vectors are returned
        """
        return self.model.solve_all(k_grid,eig_vectors=eig_vectors)

    
    def bandstructure(self, ylim=(None,None), color=None, ax=None):
        """ Plot a representation of the band structure
        
        :type  ylim: tuple, list
        :param ylim: lower and upper limit of y-values (ymin,ymax)
        """

        """ seekpath automatic lines"""
        #path = sp.get_explicit_k_path((lattice, positions, numbers), False, recipe="hpkot", threshold=1e-5,reference_distance=1)
        #expath = path['explicit_kpoints_abs'][:5]
        #labels = path['explicit_kpoints_labels'][:5]

        """ manual lines"""
        path=[[0.0,0.0,0.5],[0.5,0.0,0.5],[0.5,0,0.0],[0.0,0.0,0.0],[0,0,0.5],[2./3.,1./3.,0.5],[2./3.,1./3.,0.0],[0,0,0]]
        label=(r'$A $',      r'$L$',       r'$M$',   r'$\Gamma$', r'$A $', r'$H$',  r'$K$',r'$\Gamma $')

        
        # call function k_path to construct the actual path
        (k_vec,k_dist,k_node)=self.model.k_path(path,301,report=False)

        evals =self.model.solve_all(k_vec)
        
        fig = None
        if not ax:
            fig, ax = plt.subplots(figsize=(8,6))
            fig.tight_layout()

            ax.set_title("Bandstructure for Zno based on Kobayashi")
            ax.set_ylabel("Band energy")

            # specify horizontal axis details
            ax.set_xlim([0,k_node[-1]])
            # put tickmarks and labels at node positions
            ax.set_xticks(k_node)
            ax.set_xticklabels(label)
            # add vertical lines at node positions

            for n in range(len(k_node)):
                if label[n] == r'$\Gamma$':
                    ax.axvline(x=k_node[n],linewidth=1, color='k')
                else:
                    ax.axvline(x=k_node[n],linewidth=0.5, color='k')
    
        for band in evals:
            ax.plot(k_dist, band, color=color)

        if not fig:
            return ax
        else:
            ax.set_ylim(ylim)
            return ax, fig

#    def __repr__(self):
#        return "Tight binding model for: \n \n {} \n {}".format(self.crystal, self.atom_matrix)

In [59]:
class TbFitter():
    mu = 0
    T = 1
    def __init__(self, model, fitting_k, fitting_E, band_range=(0,-1), monitor=False, tolerance=None):
        self.model = model
        self.monitor = monitor
        self.fitting_k = fitting_k
        self.fitting_E = fitting_E
        self.tolerance = tolerance
        self.band_range = band_range
        
    def fit(self):
        initial_args = self.model.atom_matrix.get_hopping_parameters()
        
        return minimize(self.fit_function, initial_args, tol=self.tolerance)
        
    def fit_function(self, args):
        Vss, Vxx, Vxy, Vsapc, Vpasc = args
        self.model.set_hopping_parameters(Vss, Vxx, Vxy, Vsapc, Vpasc)
        E =self.model.calculate(self.fitting_k)[self.band_range[0]:self.band_range[1]]
    
        

        diff=self.fitting_E-E

        diff=abs(diff)**2
        
        val=sum((diff*self.weightfun(self.fitting_E,E)).ravel())
        
        if self.monitor:
            print(val)
            
        self.lastval=val
        self.weightsum=sum((self.weightfun(self.fitting_E,E)).ravel())
        return val
    
    def weightfun(self,E1,E2):
        w1=1./np.cosh((E1-self.mu)/self.T)**2/self.T/8
        w2=1./np.cosh((E2-self.mu)/self.T)**2/self.T/8
        return w1+w2

In [60]:
ZnO = Crystal(lattice = np.array([
        [ np.sqrt(3)/2, -0.5, 0.0],
        [ 0.0, 1.0,  0.0],
        [ 0.0, 0.0, 1.65]])*3.28)

ZnO.add_atom(Oxygen([0,0,0]))
ZnO.add_atom(Oxygen([2/3, 1/3, 1/2]))
ZnO.add_atom(Zinc([2/3, 1/3, 1/8]))
ZnO.add_atom(Zinc([0.0, 0.0, 5/8]))


ZnO_tb = TightBinding(ZnO)

Vss   = -6.043
Vxx   =  7.157
Vxy   = 10.578
Vsapc =  4.703
Vpasc =  8.634

ZnO_tb.set_hopping_parameters(Vss, Vxx, Vxy, Vsapc, Vpasc)

In [71]:
ax, fig = ZnO_tb.bandstructure(color="black",ylim=(None,None))
plt.show()



In [92]:
fitting_k=np.array([
    [  0.5,   0.0, 0.5],
    [  0.0,   0.0, 0.0],
    [  0.0,   0.0, 0.5],
    [2./3., 1./3., 0.5],
    [2./3., 1./3., 0.0]])

fitting_E = np.array([
    [-5.78, -5.85, -3.63, -5.82, -5.63], #2
    [-5.78, -1.52, -3.63, -5.82, -5.63], #3
    [-2.44, -1.52, -0.79, -3.13, -3.90], #4
    [-2.44,  0.00, -0.79, -3.13, -2.63], #5
    [-2.34,  0.00, -0.79, -2.30, -2.63], #6
    [-2.34,  0.00, -0.79, -2.30, -2.30], #7                 
    [ 9.05,  3.30,  6.11,  9.74, 10.55], #8                 
])



fitting_k = fitting_k[3:]


fitting_E = fitting_E[:,3:]

band_range = (2,10)

In [96]:
print(fitting_E)
print(ZnO_tb.calculate(fitting_k)[2:8])


[[ -5.82  -5.63]
 [ -5.82  -5.63]
 [ -3.13  -3.9 ]
 [ -3.13  -2.63]
 [ -2.3   -2.63]
 [ -2.3   -2.3 ]
 [  9.74  10.55]]
[[-5.14654924 -5.81770938]
 [-5.14654924 -4.27114878]
 [-2.30128942 -3.47985341]
 [-2.30128942 -2.7243382 ]
 [-1.9722264  -1.19709911]
 [-1.9722264  -1.16022449]]

In [ ]:
ZnO_fitter = TbFitter(model=ZnO_tb, fitting_k=fitting_k, fitting_E=fitting_E, band_range=band_range, monitor=False)

In [81]:
bands, waves = ZnO_tb.calculate([[0, 0, 0]],True)

In [85]:
bands[2:10],waves[2:10]


Out[85]:
(array([[ -5.85021301e+00],
        [ -1.51993352e+00],
        [ -1.51993352e+00],
        [  3.39619342e-04],
        [  3.39619342e-04],
        [  3.39619342e-04],
        [  3.30018703e+00],
        [  7.38985695e+00]]),
 array([[[  4.60156348e-02+0.j,  -4.89242107e-01+0.j,   1.66533454e-16+0.j,
            1.11022302e-16+0.j,  -4.60156348e-02+0.j,   4.89242107e-01+0.j,
           -2.40457479e-17+0.j,  -2.69760010e-17+0.j,  -4.68208634e-01+0.j,
           -1.98255889e-01+0.j,  -9.41114618e-17+0.j,  -1.15183043e-16+0.j,
            4.68208634e-01+0.j,   1.98255889e-01+0.j,   1.99421395e-17+0.j,
            2.13472683e-17+0.j]],
 
        [[  2.73566479e-18+0.j,  -1.29259798e-16+0.j,  -5.63327865e-01+0.j,
            1.94561682e-01+0.j,  -2.77555756e-17+0.j,  -5.55111512e-17+0.j,
            5.63327865e-01+0.j,  -1.94561682e-01+0.j,  -6.84966767e-17+0.j,
            1.14606573e-17+0.j,   3.59687050e-01+0.j,  -1.24228397e-01+0.j,
            1.29823440e-16+0.j,   4.40504939e-17+0.j,  -3.59687050e-01+0.j,
            1.24228397e-01+0.j]],
 
        [[  7.67992195e-18+0.j,  -2.44298317e-17+0.j,  -1.94561682e-01+0.j,
           -5.63327865e-01+0.j,  -7.63278329e-17+0.j,   5.20417043e-17+0.j,
            1.94561682e-01+0.j,   5.63327865e-01+0.j,  -8.30221386e-17+0.j,
            1.10912403e-16+0.j,   1.24228397e-01+0.j,   3.59687050e-01+0.j,
            1.59676089e-16+0.j,   1.11132202e-16+0.j,  -1.24228397e-01+0.j,
           -3.59687050e-01+0.j]],
 
        [[  2.74944061e-17+0.j,   4.58862131e-01+0.j,   1.29339482e-01+0.j,
           -3.83774465e-01+0.j,   2.08166817e-17+0.j,   4.58862131e-01+0.j,
            1.29339482e-01+0.j,  -3.83774465e-01+0.j,  -1.11022302e-16+0.j,
           -2.65537391e-01+0.j,  -7.48470322e-02+0.j,   2.22085161e-01+0.j,
            3.05311332e-16+0.j,  -2.65537391e-01+0.j,  -7.48470322e-02+0.j,
            2.22085161e-01+0.j]],
 
        [[  6.48081543e-18+0.j,   1.42720154e-01+0.j,   4.91114833e-01+0.j,
            3.36159449e-01+0.j,   5.55111512e-17+0.j,   1.42720154e-01+0.j,
            4.91114833e-01+0.j,   3.36159449e-01+0.j,   6.93889390e-17+0.j,
           -8.25902486e-02+0.j,  -2.84201600e-01+0.j,  -1.94530986e-01+0.j,
           -2.01227923e-16+0.j,  -8.25902486e-02+0.j,  -2.84201600e-01+0.j,
           -1.94530986e-01+0.j]],
 
        [[  2.02497722e-17-0.j,   3.79001978e-01+0.j,  -3.41531134e-01+0.j,
            3.38053224e-01+0.j,  -1.38777878e-17+0.j,   3.79001978e-01+0.j,
           -3.41531134e-01+0.j,   3.38053224e-01+0.j,   0.00000000e+00+0.j,
           -2.19323386e-01+0.j,   1.97639509e-01+0.j,  -1.95626889e-01+0.j,
            5.55111512e-17+0.j,  -2.19323386e-01+0.j,   1.97639509e-01+0.j,
           -1.95626889e-01+0.j]],
 
        [[ -1.84589888e-01+0.j,   1.11022302e-16+0.j,   1.11022302e-16+0.j,
            2.77555756e-17+0.j,  -1.84589888e-01+0.j,  -8.32667268e-17+0.j,
            1.74948464e-17+0.j,   1.11030744e-17+0.j,   6.82588143e-01+0.j,
           -5.55111512e-17+0.j,  -5.72717948e-17+0.j,  -5.24887412e-17+0.j,
            6.82588143e-01+0.j,   0.00000000e+00+0.j,  -7.95146072e-18+0.j,
           -6.76934794e-18+0.j]],
 
        [[  1.13245922e-01+0.j,   3.29808141e-01+0.j,   8.32667268e-17+0.j,
            1.11022302e-16+0.j,  -1.13245922e-01+0.j,  -3.29808141e-01+0.j,
            1.31427208e-17+0.j,   8.15571535e-18+0.j,  -4.90619332e-01+0.j,
            3.71072260e-01+0.j,  -1.03913628e-16+0.j,  -1.02443166e-16+0.j,
            4.90619332e-01+0.j,  -3.71072260e-01+0.j,  -5.76447980e-18+0.j,
           -7.62905814e-18+0.j]]]))

In [ ]: