from QuantumATK import *
# -------------------------------------------------------------
# Bulk configuration
# -------------------------------------------------------------

# Set up lattice
lattice = FaceCenteredCubic(5.4306*Angstrom)

# Define elements
elements = [Silicon, Silicon]

# Define coordinates
fractional_coordinates = [[ 0.  ,  0.  ,  0.  ],
                          [ 0.25,  0.25,  0.25]]

# Set up configuration
bulk_configuration = BulkConfiguration(
    bravais_lattice=lattice,
    elements=elements,
    fractional_coordinates=fractional_coordinates
    )

# -------------------------------------------------------------
# Calculator
# -------------------------------------------------------------
numerical_accuracy_parameters = NumericalAccuracyParameters(
    k_point_sampling=(9, 9, 9),
    density_mesh_cutoff=150.0*Hartree,
    )

calculator = LCAOCalculator(
    numerical_accuracy_parameters=numerical_accuracy_parameters,
    )

bulk_configuration.setCalculator(calculator)
nlprint(bulk_configuration)
bulk_configuration.update()

# -------------------------------------------------------------
# Optimize Geometry
# -------------------------------------------------------------
bulk_configuration = OptimizeGeometry(
        bulk_configuration,
        max_forces=0.0005*eV/Ang,
        max_stress=0.0005*eV/Ang**3,
        max_steps=200,
        max_step_length=0.2*Ang,
        trajectory_filename=None,
        optimizer_method=LBFGS(),
        )
nlsave('Silicon_uniaxial.nc', bulk_configuration)
nlprint(bulk_configuration)

# -------------------------------------------------------------
# Bandstructure
# -------------------------------------------------------------
bandstructure = Bandstructure(
    configuration=bulk_configuration,
    route=['L', 'G', 'X'],
    points_per_segment=201,
    bands_above_fermi_level=All
    )
nlsave('Silicon_uniaxial.nc', bandstructure)

# -------------------------------------------------------------
# Optimize Geometry
# -------------------------------------------------------------
bulk_configuration = OptimizeGeometry(
        bulk_configuration,
        max_forces=0.0005*eV/Ang,
        max_stress=0.0005*eV/Ang**3,
        max_steps=200,
        max_step_length=0.2*Ang,
        trajectory_filename=None,
        optimizer_method=LBFGS(),
        target_stress=[1, 0, 0, 0, 0, 0]*GPa,
        )
nlsave('Silicon_uniaxial.nc', bulk_configuration)
nlprint(bulk_configuration)

# -------------------------------------------------------------
# Bandstructure
# -------------------------------------------------------------
bandstructure = Bandstructure(
    configuration=bulk_configuration,
    route=['L', 'G', 'B'],
    points_per_segment=201,
    bands_above_fermi_level=All
    )
nlsave('Silicon_uniaxial.nc', bandstructure)