from ATK.TwoProbe import *

# Restore old two-probe configuration
vnl_file = VNLFile("lih2li.vnl")
configurations = vnl_file.readAtomicConfigurations()
two_probe_conf = configurations["lih2li"]

# Reduce basis set size for Li
basis_set_params = basisSetParameters(
    type = SingleZeta,
    element = Lithium
    )

# Set k-points for electrodes
bz_int_param = brillouinZoneIntegrationParameters( (1,1,100) )

# Create parameters for electrodes
electrode_params = ElectrodeParameters(
    brillouin_zone_integration_parameters = bz_int_param
    )

# Tolerance for convergence (default)
iteration_control_params = iterationControlParameters(
    tolerance = 1e-5
    )

# Collect parameters into a two-probe calculation method
method = TwoProbeMethod(
    (electrode_params,electrode_params),
    basis_set_parameters = basis_set_params,
    iteration_control_parameters = iteration_control_params
    )

# Specify verbosity and checkpoint file
runtime_params = runtimeParameters(
    verbosity_level = 10,
    checkpoint_filename = 'lih2li-scf.nc' 
    )

# Perform SCF calculation with chosen parameters
scf = executeSelfConsistentCalculation(
    atomic_configuration = two_probe_conf,
    method = method,
    runtime_parameters = runtime_params
    )