# Define a periodic potential along z
def potentialShift(f, shift):
    """
    :param  f:      Position in fractional coordinates
    :type  f:       array of floats

    :param shift:   Amplitude of the shift with unit energy
    :type shift:    float

    :returns:   The external potential at fractional coordinate position.
    :rtype:     float
    """
    return shift*numpy.cos(2.*numpy.pi*f[2])


# Set up a silicon lattice
configuration = BulkConfiguration(
    SimpleCubic(5.4306*Angstrom),
    [Silicon,]*8,
    fractional_coordinates = [[0.0, 0.0, 0.0], [0.25, 0.25, 0.25],
                              [0.5, 0.5, 0.0], [0.75, 0.75, 0.25],
                              [0.5, 0.0, 0.5], [0.75, 0.25, 0.75],
                              [0.0, 0.5, 0.5], [0.25, 0.75, 0.75]])

# Repeat the structure along z
configuration = configuration.repeat(1,1,3)

# Define a selfconsistent Huckel calculator
calculator = HuckelCalculator(
    basis_set=CerdaHuckelParameters.Silicon_GW_diamond_Basis,
    numerical_accuracy_parameters = NumericalAccuracyParameters(
    k_point_sampling=(4, 4, 2)),
    iteration_control_parameters = IterationControlParameters()
    )

# Set a calculator on the configuration
configuration.setCalculator(calculator())

# Perform a loop with 4 different shifts
for shift in [0.0, 1.0, 5.0, 10.0]*eV:

    fractional = configuration.fractionalCoordinates()
    atom_potentials = [(i,potentialShift(f, shift)) for i,f in enumerate(fractional)]

    configuration.setExternalPotential(AtomicShift(atom_potentials))


    # Calculate the bandstructure
    bandstructure = Bandstructure(
        configuration=configuration,
        route=['G', 'Z'],
        )

    nlsave('atomic_shift.nc', bandstructure)