def create_base(mat='graphene', layers=1, size=1):
    """Return a relaxed structure of the base material with n layers."""
    name = 'vasp/type=base/mat={0}/layers={1}'.format(mat, layers)
    atoms = Vasp(name).get_atoms()

    atoms = atoms.repeat([size, size, 1])
    return atoms

def test0():
    atoms = Atoms([Atom('O', [4, 5, 5], magmom=1),
                   Atom('C', [5, 5, 5], magmom=2),
                   Atom('O', [6, 5, 5], magmom=3)],
                   cell=(10, 10, 10))

    calc = Vasp('vasp',
                sigma=0.01,
                atoms=atoms)

    calc.write_db(fname='DB.db', append=False)

    with connect('DB.db') as con:
        data = con.get(1).data
        assert data['parameters']['sigma'] == 0.01

    print 'done'

def db_update(db_path, dft_path, delete=False, silent=False):
    """Update the database to include Vasp() calculations nested in dft_path."""
    VASPRC['mode'] = None
    db = connect(db_path)
    old_size = sum(1 for _ in db.select())

    db_paths = []
    for d in db.select():
        try:
            db_paths.append(d.data.path)
        except:
            pass

    for path in utils.calc_paths(dft_path):
        if os.path.abspath(path) in db_paths:
            continue
        calc = Vasp(path)

        if not calc.in_queue() and calc.potential_energy is None:
            for output_file in utils.calc_output_files(path):
                dead_file = os.path.join(path, output_file)
                if delete:
                    os.remove(dead_file)
                if not silent:
                    print("Dead output file: {}. Deleted: {}".format(dead_file, delete))
        else:
            ctime = calc.get_elapsed_time()

            # The write_db method throws an AttributeError for new calcs. Remove this to debug.
            old_stdout = sys.stdout
            sys.stdout = open(os.devnull, "w")
            calc.write_db(db_path, parser='=',
                        overwrite=False,
                        data={'ctime': ctime})
            sys.stdout.close()
            sys.stdout = old_stdout

            if not silent:
                print("Added calc to DB: {}".format(path))

    new_size = sum(1 for _ in db.select())
    added = new_size - old_size
    if not silent:
        print("{} total entries. {} new entries added.".format(new_size, added))

from vasp import Vasp
from ase.lattice.cubic import FaceCenteredCubic
from ase import Atoms, Atom
# bulk system
atoms = FaceCenteredCubic(directions=[[0, 1, 1],
                                      [1, 0, 1],
                                      [1, 1, 0]],
                                      size=(1, 1, 1),
                                      symbol='Rh')
calc = Vasp('bulk/bulk-rh',
            xc='PBE',
            encut=350,
            kpts=[4, 4, 4],
            isif=3,
            ibrion=2,
            nsw=10,
            atoms=atoms)
bulk_energy = atoms.get_potential_energy()
# atomic system
atoms = Atoms([Atom('Rh',[5, 5, 5])],
              cell=(7, 8, 9))
calc = Vasp('bulk/atomic-rh',
            xc='PBE',
            encut=350,
            kpts=[1, 1, 1],
            atoms=atoms)
atomic_energy = atoms.get_potential_energy()
calc.stop_if(None in (bulk_energy, atomic_energy))
cohesive_energy = atomic_energy - bulk_energy
print 'The cohesive energy is {0:1.3f} eV'.format(cohesive_energy)

from vasp import Vasp
calc = Vasp('molecules/O2-sp-singlet')
calc.clone('molecules/O2-sp-singlet-magmoms')
calc.set(lorbit=11)
atoms = calc.get_atoms()
magmoms = atoms.get_magnetic_moments()
print('singlet ground state')
for i, atom in enumerate(atoms):
    print('atom {0}: magmom = {1}'.format(i, magmoms[i]))
print(atoms.get_magnetic_moment())
calc = Vasp('molecules/O2-sp-triplet')
calc.clone('molecules/O2-sp-triplet-magmoms')
calc.set(lorbit=11)
atoms = calc.get_atoms()
magmoms = atoms.get_magnetic_moments()
print()
print('triplet ground state')
for i, atom in enumerate(atoms):
    print('atom {0}: magmom = {1}'.format(i, magmoms[i]))
print(atoms.get_magnetic_moment())

from ase.thermochemistry import IdealGasThermo
from vasp import Vasp
import numpy as np
import matplotlib.pyplot as plt
# first we get the electronic energies
c1 = Vasp('molecules/wgs/CO')
E_CO = c1.potential_energy
CO = c1.get_atoms()
c2 = Vasp('molecules/wgs/CO2')
E_CO2 = c2.potential_energy
CO2 = c2.get_atoms()
c3 = Vasp('molecules/wgs/H2')
E_H2 = c3.potential_energy
H2 = c3.get_atoms()
c4 = Vasp('molecules/wgs/H2O')
E_H2O = c4.potential_energy
H2O = c4.get_atoms()
# now we get the vibrational energies
h = 4.1356675e-15  # eV * s
c = 3.0e10  # cm / s
calc = Vasp('molecules/wgs/CO-vib')
vib_freq = calc.get_vibrational_frequencies()
CO_vib_energies = [h * c * nu for nu in vib_freq]
calc = Vasp('molecules/wgs/CO2-vib')
vib_freq = calc.get_vibrational_frequencies()
CO2_vib_energies = [h * c * nu for nu in vib_freq]
calc = Vasp('molecules/wgs/H2-vib')
vib_freq = calc.get_vibrational_frequencies()
H2_vib_energies = [h * c * nu for nu in vib_freq]
calc = Vasp('molecules/wgs/H2O-vib')
vib_freq = calc.get_vibrational_frequencies()

from vasp import Vasp
atoms = molecule('N2')
atoms.set_cell((10,10,10), scale_atoms=False)
# first we relax a molecule
calc = Vasp('molecules/n2-relax',
            xc='PBE',
            encut=300,
            ibrion=2,
            nsw=5,
            atoms=atoms)
electronicenergy = atoms.get_potential_energy()
# next, we get vibrational modes
calc2 = Vasp('molecules/n2-vib',
             xc='PBE',
             encut=300,
             ibrion=6,
             nfree=2,
             potim=0.15,
             nsw=1,
             atoms=atoms)
calc2.wait()
vib_freq = calc2.get_vibrational_frequencies() # in cm^1
#convert wavenumbers to energy
h = 4.1356675e-15 # eV*s
c = 3.0e10 #cm/s
vib_energies = [h*c*nu for nu in vib_freq]
print('vibrational energies\n====================')
for i,e in enumerate(vib_energies):
    print('{0:02d}: {1} eV'.format(i,e))
# # now we can get some properties. Note we only need one vibrational
# energy since there is only one mode. This example does not work if
# you give all the energies because one energy is zero.

from vasp import Vasp
from ase.lattice.cubic import BodyCenteredCubic
atoms = BodyCenteredCubic(directions=[[1, 0, 0],
                                      [0, 1, 0],
                                      [0, 0, 1]],
                                      size=(1, 1, 1),
                                      symbol='Fe')
NUPDOWNS = [0.0, 2.0, 4.0, 5.0, 6.0, 8.0]
energies = []
for B in NUPDOWNS:
    calc = Vasp('bulk/Fe-bcc-fixedmagmom-{0:1.2f}'.format(B),
                xc='PBE',
                encut=300,
                kpts=[4, 4, 4],
                ispin=2,
                nupdown=B,
                atoms=atoms)
    energies.append(atoms.get_potential_energy())
if None in energies:
    calc.abort()
import matplotlib.pyplot as plt
plt.plot(NUPDOWNS, energies)
plt.xlabel('Total Magnetic Moment')
plt.ylabel('Energy (eV)')
plt.savefig('images/Fe-fixedmagmom.png')

from vasp import Vasp
from ase.lattice.surface import fcc111
import matplotlib.pyplot as plt
Nlayers = [3, 4, 5, 6, 7, 8, 9, 10, 11]
energies = []
sigmas = []
for n in Nlayers:
    slab = fcc111('Cu', size=(1, 1, n), vacuum=10.0)
    slab.center()
    calc = Vasp('bulk/Cu-layers/{0}'.format(n),
                xc='PBE',
                encut=350,
                kpts=[8, 8, 1],
                atoms=slab)
    calc.set_nbands(f=2)  # the default nbands in VASP is too low for Cu
    energies.append(slab.get_potential_energy())
calc.stop_if(None in energies)
for i in range(len(Nlayers) - 1):
    N = Nlayers[i]
    DeltaE_N = energies[i + 1] - energies[i]
    sigma = 0.5 * (-N * energies[i + 1] + (N + 1) * energies[i])
    sigmas.append(sigma)
    print 'nlayers = {1:2d} sigma = {0:1.3f} eV/atom'.format(sigma, N)
plt.plot(Nlayers[0:-1], sigmas, 'bo-')
plt.xlabel('Number of layers')
plt.ylabel('Surface energy (eV/atom)')
plt.savefig('images/Cu-unrelaxed-surface-energy.png')

from vasp import Vasp
# get relaxed geometry
calc = Vasp('molecules/wgs/CO2')
CO2 = calc.get_atoms()
# now do the vibrations
calc = Vasp('molecules/wgs/CO2-vib',
            xc='PBE',
            encut=350,
            ismear=0,
            ibrion=6,
            nfree=2,
            potim=0.02,
            nsw=1,
            atoms=CO2)
calc.wait()
vib_freq = calc.get_vibrational_frequencies()
for i, f in enumerate(vib_freq):
    print('{0:02d}: {1} cm^(-1)'.format(i, f))

from ase import Atoms, Atom
from vasp import Vasp
Vasp.vasprc(mode=None)
#Vasp.log.setLevel(10)
import matplotlib.pyplot as plt
import numpy as np
from ase.dft import DOS
import pylab as plt
a = 3.9  # approximate lattice constant
b = a / 2.
bulk = Atoms([Atom('Pd', (0.0, 0.0, 0.0))],
             cell=[(0, b, b),
                   (b, 0, b),
                   (b, b, 0)])
kpts = [8, 10, 12, 14, 16, 18, 20]
calcs = [Vasp('bulk/pd-dos-k{0}-ismear-5'.format(k),
              encut=300,
              xc='PBE',
              kpts=[k, k, k],
              atoms=bulk) for k in kpts]
Vasp.wait(abort=True)
for calc in calcs:
    # this runs the calculation
    if calc.potential_energy is not None:
        dos = DOS(calc, width=0.2)
        d = dos.get_dos() + k / 4.0
        e = dos.get_energies()
        plt.plot(e, d, label='k={0}'.format(k))
    else:
        pass
plt.xlabel('energy (eV)')

from vasp import Vasp
atoms = Vasp('bulk/Al-lda-vasp').get_atoms()
atoms2 = Vasp('bulk/Al-lda-ase').get_atoms()
import numpy as np
cellA = atoms.get_cell()
cellB = atoms2.get_cell()
print((np.abs(cellA - cellB) < 0.01).all())

from vasp import Vasp
from ase.lattice.cubic import FaceCenteredCubic
from ase.dft import DOS
atoms = FaceCenteredCubic(directions=[[0, 1, 1],
                                      [1, 0, 1],
                                      [1, 1, 0]],
                                      size=(1, 1, 1),
                                      symbol='Ni')
atoms[0].magmom = 1
calc = Vasp('bulk/Ni-PBE',
            ismear=-5,
            kpts=[5, 5, 5],
            xc='PBE',
            ispin=2,
            lorbit=11,
            lwave=True, lcharg=True,  # store for reuse
            atoms=atoms)
e = atoms.get_potential_energy()
print('PBE energy:   ',e)
calc.stop_if(e is None)
dos = DOS(calc, width=0.2)
e_pbe = dos.get_energies()
d_pbe = dos.get_dos()
calc.clone('bulk/Ni-PBE0')
calc.set(xc='pbe0')
atoms = calc.get_atoms()
pbe0_e = atoms.get_potential_energy()
if atoms.get_potential_energy() is not None:
    dos = DOS(calc, width=0.2)
    e_pbe0 = dos.get_energies()
    d_pbe0 = dos.get_dos()

from vasp import Vasp
from ase import Atom, Atoms
encuts = [250, 300, 350, 400, 450, 500, 550]
D = []
for encut in encuts:
    atoms = Atoms([Atom('O', [5, 5, 5], magmom=2)],
                  cell=(10, 10, 10))
    calc = Vasp('molecules/O-sp-triplet-{0}'.format(encut),
                xc='PBE',
                encut=encut,
                ismear=0,
                ispin=2,
                atoms=atoms)
    E_O = atoms.get_potential_energy()
    # now relaxed O2 dimer
    atoms = Atoms([Atom('O', [5,    5, 5], magmom=1),
                   Atom('O', [6.22, 5, 5], magmom=1)],
                  cell=(10, 10, 10))
    calc = Vasp('molecules/O2-sp-triplet-{0}'.format(encut),
                xc='PBE',
                encut=encut,
                ismear=0,
                ispin=2,   # turn spin-polarization on
                ibrion=2,  # this turns relaxation on
                nsw=10,
                atoms=atoms)
    E_O2 = atoms.get_potential_energy()
    if None not in (E_O, E_O2):
        d = 2*E_O - E_O2
        D.append(d)
        print('O2 -> 2O encut = {0}  D = {1:1.3f} eV'.format(encut, d))

from vasp import Vasp
calc = Vasp('bulk/tio2/step1-0.90')
calc.clone('bulk/tio2/step2-0.90')
#calc.set(isif=4)
print calc.set(isif=4)
print calc.calculation_required()

from vasp import Vasp
from ase.structure import molecule
atoms = molecule('CO')
atoms.center(vacuum=5)
calc = Vasp('molecules/CO-vacuum',
            encut=600,
            prec='Accurate',
            ismear=0,
            sigma=0.05,
            ibrion=2,
            nsw=0,
            ediff=1e-6,
            atoms=atoms)
print(atoms.get_potential_energy())
print(atoms.get_forces())
print('Calculation time: {} seconds'.format(calc.get_elapsed_time()))

from vasp import Vasp
# don't forget to normalize your total energy to a formula unit. Cu2O
# has 3 atoms, so the number of formula units in an atoms is
# len(atoms)/3.
calc = Vasp('bulk/Cu2O')
atoms1 = calc.get_atoms()
cu2o_energy = atoms1.get_potential_energy()
calc = Vasp('bulk/CuO')
atoms2 = calc.get_atoms()
cuo_energy = atoms2.get_potential_energy()
# make sure to use the same cutoff energy for the O2 molecule!
calc = Vasp('molecules/O2-sp-triplet-400')
atoms3 = calc.get_atoms()
o2_energy = atoms3.get_potential_energy()
calc.stop_if(None in [cu2o_energy, cuo_energy, o2_energy])
cu2o_energy /= (len(atoms1) / 3)  # note integer math
cuo_energy /= (len(atoms2) / 2)
rxn_energy = 4.0 * cuo_energy - o2_energy - 2.0 * cu2o_energy
print 'Reaction energy = {0} eV'.format(rxn_energy)

from vasp import Vasp
import matplotlib.pyplot as plt
calc = Vasp('surfaces/Al-Na-nodip')
atoms = calc.get_atoms()
x, y, z, lp = calc.get_local_potential()
nx, ny, nz = lp.shape
axy_1 = [np.average(lp[:, :, z]) for z in range(nz)]
# setup the x-axis in realspace
uc = atoms.get_cell()
xaxis_1 = np.linspace(0, uc[2][2], nz)
e1 = atoms.get_potential_energy()
calc = Vasp('surfaces/Al-Na-dip-step2')
atoms = calc.get_atoms()
x, y, z, lp = calc.get_local_potential()
nx, ny, nz = lp.shape
axy_2 = [np.average(lp[:, :, z]) for z in range(nz)]
# setup the x-axis in realspace
uc = atoms.get_cell()
xaxis_2 = np.linspace(0, uc[2][2], nz)
ef2 = calc.get_fermi_level()
e2 = atoms.get_potential_energy()
print 'The difference in energy is {0} eV.'.format(e2-e1)
plt.plot(xaxis_1, axy_1, label='no dipole correction')
plt.plot(xaxis_2, axy_2, label='dipole correction')
plt.plot([min(xaxis_2), max(xaxis_2)], [ef2, ef2], 'k:', label='Fermi level')
plt.xlabel('z ($\AA$)')
plt.ylabel('xy-averaged electrostatic potential')
plt.legend(loc='best')
plt.savefig('images/dip-vs-nodip-esp.png')

from vasp import Vasp
from ase.units import Debye
calc = Vasp('molecules/co-centered')
dipole_moment = calc.get_dipole_moment()
print('The dipole moment is {0:1.2f} Debye'.format(dipole_moment))

from ase import Atoms, Atom
from vasp import Vasp
import matplotlib.pyplot as plt
import numpy as np
a = 3.9  # approximate lattice constant
b = a / 2.
bulk = Atoms([Atom('Pd', (0.0, 0.0, 0.0))],
             cell=[(0, b, b),
                   (b, 0, b),
                   (b, b, 0)])
RWIGS = [1.0, 1.1, 1.25, 1.5, 2.0, 2.5, 3.0, 4.0, 5.0 ]
ED, WD, N = [], [], []
for rwigs in RWIGS:
    calc = Vasp('bulk/pd-ados')
    calc.clone('bulk/pd-ados-rwigs-{0}'.format(rwigs))
    calc.set(rwigs={'Pd': rwigs})
    if calc.potential_energy is None:
        continue
        # now get results
        ados = VaspDos(efermi=calc.get_fermi_level())
        energies = ados.energy
        dos = ados.site_dos(0, 'd')
        #we will select energies in the range of -10, 5
        ind = (energies < 5) & (energies > -10)
        energies = energies[ind]
        dos = dos[ind]
        Nstates = np.trapz(dos, energies)
        occupied = energies <= 0.0
        N_occupied_states = np.trapz(dos[occupied], energies[occupied])
        ed = np.trapz(energies * dos, energies) / np.trapz(dos, energies)
        wd2 = np.trapz(energies**2 * dos, energies) / np.trapz(dos, energies)