from ATK.TwoProbe import *
import ATK
ATK.setVerbosityLevel(2)

# Restore initial and optimized two-probe configurations
ini_vnl_file = VNLFile("lih2li-ini.vnl")
ini_configurations = ini_vnl_file.readAtomicConfigurations()
ini_two_probe_conf = ini_configurations["lih2li-ini"]
opt_vnl_file = VNLFile("lih2li-opt.vnl")
opt_configurations = opt_vnl_file.readAtomicConfigurations()
opt_two_probe_conf = opt_configurations["lih2li-opt"]

# Create parameters for electrodes
electrode_bz = brillouinZoneIntegrationParameters( (1,1,80) )
electrode_occ = eigenstateOccupationParameters(
    temperature = 1300*Kelvin
    )
electrode_params = ElectrodeParameters(
    eigenstate_occupation_parameters = electrode_occ,
    brillouin_zone_integration_parameters = electrode_bz
    )

# Collect parameters into a two-probe calculation method
two_probe_method=TwoProbeMethod(
    electrode_parameters = (electrode_params, electrode_params),
    exchange_correlation_type = LDA.PZ,
    )
    
# Execute self-consistent calculations
ini_scf = executeSelfConsistentCalculation(
    ini_two_probe_conf,
    two_probe_method,
    )
opt_scf = executeSelfConsistentCalculation(
    opt_two_probe_conf,
    two_probe_method,
    )

# Calculate the transmission eigenvalues at the Fermi level
k_point_list = [(0.0,0.0)]
ini_transmission_eigenvalues = calculateTransmissionEigenvalues(
    ini_scf,
    energy = 0*electronVolt,
    quantum_numbers = (k_point_list)
    )
opt_transmission_eigenvalues = calculateTransmissionEigenvalues(
    opt_scf,
    energy = 0*electronVolt,
    quantum_numbers = (k_point_list)
    )

# Write the transmission eigenvalues
print "Initial transmission eigenvalue: %s" % (ini_transmission_eigenvalues[0])
print "Optimized transmission eigenvalue: %s" % (opt_transmission_eigenvalues[0])

# Calculate the corresponding transmission eigenstates
ini_transmission_eigenstates = calculateTransmissionEigenstates(
    ini_scf,
    energy = 0*electronVolt,
    quantum_numbers = (0,k_point_list)
    )
opt_transmission_eigenstates = calculateTransmissionEigenstates(
    opt_scf,
    energy = 0*electronVolt,
    quantum_numbers = (0,k_point_list)
    )

# Save the results
ini_vnl_file.addToSample(ini_transmission_eigenstates[0],"lih2li-ini")
opt_vnl_file.addToSample(opt_transmission_eigenstates[0],"lih2li-opt")