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

# Read the atomic configuration from an VNL file.
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
    )

# Define electrode method
electrode_bz = brillouinZoneIntegrationParameters( (1,1,100) )
electrode_params = ElectrodeParameters(
    brillouin_zone_integration_parameters = electrode_bz
    )

# Initialize the electron density
two_probe_params = twoProbeAlgorithmParameters(
    initial_density_type = InitialDensityType.NeutralAtom
    )

# Define the two-probe method
twoprobe_method=TwoProbeMethod(
    electrode_parameters = (electrode_params, electrode_params),
    basis_set_parameters = basis_set_params,
    algorithm_parameters = two_probe_params
    )

# Specify checkpoint file
runtime_params = runtimeParameters(
    checkpoint_filename = 'lih2li-scf-neutral-atom-density.nc' 
    )

# Perform the self-consistent two-probe calculation
scf = executeSelfConsistentCalculation(
    atomic_configuration = two_probe_conf,
    method = twoprobe_method,
    runtime_parameters = runtime_params
    )