from gpaw import GPAW, PW
from ase import *
from phonopy import Phonopy
from phonopy.structure.atoms import Atoms as PhonopyAtoms
import numpy as np

# GPAW setting
a = 5.404
bulk = PhonopyAtoms(symbols=['Si']*8,
                    scaled_positions=[(0, 0, 0),
                                      (0, 0.5, 0.5),
                                      (0.5, 0, 0.5),
                                      (0.5, 0.5, 0),
                                      (0.25, 0.25, 0.25),
                                      (0.25, 0.75, 0.75),
                                      (0.75, 0.25, 0.75),
                                      (0.75, 0.75, 0.25)] )
bulk.set_cell(np.diag((a, a, a)))

calc = GPAW(mode=PW(300),
            kpts={'size': (4, 4, 4)},
            symmetry={'symmorphic': False})

phonon = Phonopy(bulk,
                 [[1,0,0],[0,1,0],[0,0,1]],
                 primitive_matrix=[[0, 0.5, 0.5],
                                   [0.5, 0, 0.5],
                                   [0.5, 0.5, 0]],
                 distance=0.01)
print "[Phonopy] Atomic displacements:"
disps = phonon.get_displacements()
for d in disps:
    print "[Phonopy]", d[0], d[1:]
supercells = phonon.get_supercells_with_displacements()

# Force calculations by calculator
set_of_forces = []
for scell in supercells:
    cell = Atoms(symbols=scell.get_chemical_symbols(),
                 scaled_positions=scell.get_scaled_positions(),
                 cell=scell.get_cell(),
                 pbc=True)
    cell.set_calculator(calc)
    forces = cell.get_forces()
    drift_force = forces.sum(axis=0)
    print "[Phonopy] Drift force:", "%11.5f"*3 % tuple(drift_force)
    # Simple translational invariance
    for force in forces:
        force -= drift_force / forces.shape[0]
    set_of_forces.append(forces)

# Phonopy post-process
phonon.produce_force_constants(forces=set_of_forces)
print
print "[Phonopy] Phonon frequencies at Gamma:"
for i, freq in enumerate(phonon.get_frequencies((0, 0, 0))):
    print "[Phonopy] %3d: %10.5f THz" %  (i + 1, freq) # THz


