from Dacapo import Dacapo
from ASE import Atom, ListOfAtoms
from ASE.Visualization.VMD import VMD
from ASE.Visualization.gnuplot import gnuplot

bulk = ListOfAtoms([Atom('Al', (0, 0, 0))] )
a = 4.05
b = a / 2
bulk.SetUnitCell([(0, b, b),
                  (b, 0, b),
                  (b, b, 0)])

bulk_plot = VMD(bulk, repeat=(5,5,5))

calc = Dacapo(kpts=(1, 1, 1),          # set the k-points (Monkhorst-Pack)
              planewavecutoff=200,     # planewavecutoff in eV
              nbands=6,                # set the number of electronic bands
              usesymm=True,            # use symmetry to reduce the k-point set
              out='Al-fcc.nc',         # define the out netcdf file
              txtout='Al-fcc.txt')     # define the ascii out file

bulk.SetCalculator(calc)
calc.SetDensityCutoff(400)

# make a plot of the convergence with respect to k-points
kpt_energies = []
for n in [1, 2, 4, 6]: 
    calc.SetBZKPoints((n, n, n))
    energy = bulk.GetPotentialEnergy() 
    kpt_energies.append((n, energy))

kpt_plot = gnuplot(kpt_energies) 


# make a plot of the convergence with respect to planewavecutoff
n = 2
calc.SetBZKPoints((n, n, n))
pw_energies = []
for planewavecutoff in [75, 100, 150, 200, 250, 300]:
    calc.SetPlaneWaveCutoff(planewavecutoff)
    energy = bulk.GetPotentialEnergy()
    pw_energies.append((planewavecutoff, energy))

pw_plot = gnuplot(pw_energies)