import Numeric
from ASE.Transport import TransmissionTools
from ASE.Transport.ScatteringHamiltonianMatrix import ScatteringHamiltonianMatrix
from ASE.Transport.LeadHamiltonianMatrix import LeftLeadHamiltonianMatrix,RightLeadHamiltonianMatrix

# Construct a lead Hamiltonian matrix describing 2 principal layers of 
# a 1d TB chain with up to second-nearest neighbor hopping                

H_lead=Numeric.zeros([4,4],Numeric.Complex)

# On-site energies are zero
for i in range(4):
    H_lead[i,i]=0.0
# Nearest neighbor hopping is -1.0
for i in range(3):
    H_lead[i,i+1]=-1.0
    H_lead[i+1,i]=-1.0
# Next-nearest neighbor hopping is 0.2
for i in range(2):
    H_lead[i,i+2]=0.2
    H_lead[i+2,i]=0.2

# Construct the Hamiltonian of the scattering region including 1 principal 
# layer of the lead. The scattering region consists of 2 sites at zero        
# energy and with external coupling of 0.2 and internal coupling og 0.8.      

H_scat=Numeric.zeros([6,6],Numeric.Complex)

# Principal layers
H_scat[:2,:2]=H_lead[:2,:2]
H_scat[-2:,-2:]=H_lead[:2,:2]
# Scattering region
H_scat[2,2]=0.0
H_scat[3,3]=0.0
H_scat[2,3]=0.8
H_scat[3,2]=0.8
# External coupling
H_scat[1,2]=0.2
H_scat[2,1]=0.2
H_scat[3,4]=0.2
H_scat[4,3]=0.2


# Initializing classes to contain scattering and lead Hamiltonians
# Principal layer is two sites and is the same for left and right leads 
prin=2
hmat=ScatteringHamiltonianMatrix(leftprincipallayer=prin,rightprincipallayer=prin,hamiltonianmatrix=H_scat)
left=LeftLeadHamiltonianMatrix(principallayer=prin,hamiltonianmatrix=H_lead)
right=RightLeadHamiltonianMatrix(principallayer=prin,hamiltonianmatrix=H_lead)

# Calculate transmission function at energies
en=Numeric.arange(-3,3,0.02)

#####################################################################################
# The following illustrates how to manipulate the Hamiltonian for analysis          #
# purposes. In particular it is shown how to (i) sub-diagonalize a subspace to      #
# obtain renormalized energy levels in a certain region. (ii) How to cut the        #
# coupling to selected orbitals e.g. some of the renormalized energy levels.        # 
#####################################################################################

# Sub-diagonalize the subspace of the two sites to obtain two (renormalized) energy levels
from ASE.Utilities.Wannier import HamiltonianTools
H=hmat.GetHamiltonianMatrix()
# H has dimension 6x6 corresponding to the scattering region (two levels).
# The two levels thus have indices 2 and 3
states=[2,3]
# Sub-diagonalize the subspace of the two sites to obtain two (renormalized) energy levels
H2,U,eig=HamiltonianTools.SubDiagonalize(H,states)
print "The renormalized molecular energy levels",eig
# Cut all couplings to one of the two levels. This level is now effectively removed 
# from the calculation, and only the other level will show up as a peak in the transmission.
H3=HamiltonianTools.CutCoupling(H2,[3])
hmat.SetHamiltonianMatrix(H3)

# Initialize Transmission object
trans=TransmissionTools.TransmissionTool(hmat,left,right,en,prin)
trans.InitGreenFunctions()

# We want to calculate: 
# the transmission function and the spectral functions of all sites in scattering region
trans.UpdateTransmission()
trans.UpdateSpectralFunctions()

# Run calculation
t=trans.UpdateToNetCDFFile('TBchain_trans.nc')