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

# 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
    )

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

# Define methods for constrained algorithm
constrained_algorithm = twoProbeAlgorithmParameters(
    electrode_constraint = ElectrodeConstraints.RealSpaceDensity
    )
constrained_method = TwoProbeMethod(
    electrode_parameters = (electrode_params,electrode_params),
    basis_set_parameters = basis_set_params,
    algorithm_parameters = constrained_algorithm
    )  

# Execute self-consistent constrained calculation
constrained_scf = constrained_method.apply(two_probe_conf)

# Define methods for unconstrained algorithm (default algorithm)
unconstrained_algorithm = twoProbeAlgorithmParameters(
    electrode_constraint = ElectrodeConstraints.Off
    )
unconstrained_method = TwoProbeMethod(
    electrode_parameters = (electrode_params,electrode_params),
    basis_set_parameters = basis_set_params,
    algorithm_parameters = unconstrained_algorithm
    )

# Specify checkpoint file
runtime_params = runtimeParameters(
    checkpoint_filename = 'lih2li-scf-constrained-unconstrained.nc' 
    )

# Execute self-consistent unconstrained calculation
unconstrained_scf = executeSelfConsistentCalculation(
    atomic_configuration = two_probe_conf,
    method = unconstrained_method,
    runtime_parameters = runtime_params,
    initial_calculation = constrained_scf
    )