from ATK.TwoProbe import *
import numpy 
import ATK

# Restore two-probe geometry
vnlfile = VNLFile('lih2li.vnl')
lih2li = vnlfile.readAtomicConfigurations()['lih2li']

# Use the same parameters for final bias as for zero bias
bz_int_param = brillouinZoneIntegrationParameters(
    monkhorst_pack_parameters=(1,1,100)
    )
electrode_params = ElectrodeParameters(
    brillouin_zone_integration_parameters = bz_int_param
    )
basis_set_params = basisSetParameters(
    type = SingleZeta,
    element = Lithium
    )
iteration_control_params = iterationControlParameters(
    tolerance = 1e-5
    )

ATK.setVerbosityLevel(0)

# Restore initial density from old calculation
scf = restoreSelfConsistentCalculation("lih2li-scf.nc")

print '# Bias (Volt)\tCurrent (Ampere)'

# Run bias from 0.0 and 1.0 in steps of 0.1
for voltage in numpy.arange(0.0, 1.01, 0.1):

    dft_method = TwoProbeMethod(
        electrode_parameters=(electrode_params,electrode_params),
        basis_set_parameters = basis_set_params,
        iteration_control_parameters = iteration_control_params,
        electrode_voltages = (voltage/2.0, -voltage/2.0)*Volt    
        )

    # Store each calculation in a separate NetCDF file    
    ATK.setCheckpointFilename ('lih2li-bias-%.1f.nc' % voltage)
    
    scf = executeSelfConsistentCalculation(
        atomic_configuration=lih2li,
        method = dft_method,
        initial_calculation = scf
        )

    current = calculateCurrent(scf)

    print "%.1f\t\t%.2e" %(voltage, current.inUnitsOf(Ampere))