import pylab

# Try lattice constants in steps of 0.05 Angstrom
a_values = numpy.arange(3.9,4.4,0.05)
# Create a list to store the computed total energies
energies = []
# For the first step there is no initial state
initial_state = None
for a in a_values:
    # For each lattice constant, create the configuration
    lattice = FaceCenteredCubic(a*Angstrom)
    elements = [Gold]
    cartesian_coordinates = [[ 0.,  0.,  0.]]*Angstrom
    bulk_configuration = BulkConfiguration(
        bravais_lattice=lattice,
        elements=elements,
        cartesian_coordinates=cartesian_coordinates
    )
    # 13x13x13 is a sufficient k-point sampling for gold
    numerical_accuracy_parameters = NumericalAccuracyParameters(
        k_point_sampling=(13, 13, 13)
    ) 
    iteration_control_parameters = IterationControlParameters(
        tolerance=4e-07
    )
    calculator = LCAOCalculator(
        numerical_accuracy_parameters=numerical_accuracy_parameters,
        iteration_control_parameters=iteration_control_parameters,
    )
    # Initial state is set here
    # Note that None is a valid value for the initial_state
    bulk_configuration.setCalculator(
        calculator,
        initial_state=initial_state,
    )
    # Compute and store the total energy
    energies.append(TotalEnergy(bulk_configuration).evaluate())
    # Save the converged result in a separate variable, to use as initial state for the next value of a
    initial_state = bulk_configuration
# Print results (on the master node only)
if processIsMaster():
    for a,energy in zip(a_values,energies):
        print(a, energy)
    # Simplistic approach to finding the minimum: fit a parabola
    coefs = numpy.polyfit(numpy.array(a_values),numpy.array(energies),2)
    print("Fitted minimum: a=%s Angstrom" % (-coefs[1]/coefs[0]/2))

pylab.plot(a_values,energies,'bo-')
pylab.xlabel('Lattice Constant ( Angstrom )',fontsize=16)
pylab.ylabel('Energy (eV)',fontsize=16 )
pylab.title('Lattice Constant Scanning')

# Show the plot.
pylab.show()