template-1.py

# This file as a whole is distributed under the GPL v3.
# https://www.gnu.org/licenses/gpl-3.0.en.html

from gpaw import GPAW, PW
from ase import Atoms
from phonopy import Phonopy
from phonopy.structure.atoms import PhonopyAtoms
from phonopy.phonon.band_structure import get_band_qpoints_and_path_connections
from ase.io.vasp import read_vasp
import phonopy

import sys
import shutil
import os.path
import numpy as np

def main():
    cell = read_vasp("mos2.vasp")
    phonon = phonopy_pre_process(
        cell,
        # supercell_matrix = np.array([[-1, 1, 1], [1, -1, 1], [1, 1, -1]]),  # arbitrary supercells
        supercell_matrix = np.diag([2, 2, 1]),  # simple supercells
        displacement_distance = 1e-3,
    )

    calc = GPAW(
        mode=PW(300),
        kpts={'size': (2, 2, 1)},
    )

    # The second argument here is an ASE path string.
    # You can have commas to represent discontinuities, e.g. ``'GXMGRX,MR'`` for cubic.
    path = get_band_path(cell, 'GKMG', 51)
    
    # 'None' will use a default path
    #path = get_band_path(cell, None, 51)
    # you can also override the path points if necessary
    #path = get_band_path(cell, 'GKMG', 51, path_frac=[[[0, 0, 0], [1/3, 1/3, 0], [0.5, 0, 0], [1, 0 ,0]]])

    load_or_compute_force_constants('force-constants.npy', calc, phonon, sum_rule=True)
    phonopy_post_process(phonon)

    plot_bands(phonon, path)

def plot_bands(phonon, path):
    qpoints, labels, connections = path
    phonon.run_band_structure(qpoints, path_connections=connections, labels=labels)

    # without DOS
    # fig = phonon.plot_band_structure()
    
    # with DOS
    phonon.run_mesh([20, 20, 20])
    phonon.run_total_dos()
    fig = phonon.plot_band_structure_and_dos()

    # with PDOS
    # phonon.run_mesh([20, 20, 20], with_eigenvectors=True, is_mesh_symmetry=False)
    # fig = phonon.plot_band_structure_and_dos(pdoc_indices=[[0], [1]])

    fig.savefig('band.pdf')
    fig.show()

# Stuff below this line should not need to change as frequently.
#--------------------------------------------------------------------
# The remaining function bodies in this file are MIT-licensed. (C) 2020 Michael Lamparski

def get_band_path(atoms, path_str, npoints, path_frac=None, labels=None):
    from ase.dft.kpoints import bandpath

    atoms = convert_atoms_to_ase(atoms)
    if path_str is None:
        path_str = atoms.get_cell().bandpath().path

    # Commas are part of ase's supported syntax, but we'll take care of them
    # ourselves to make it easier to get things the way phonopy wants them
    if path_frac is None:
        path_frac = []
        for substr in path_str.split(','):
            path = bandpath(substr, atoms.get_cell()[...], npoints=1)
            path_frac.append(path.kpts)

    if labels is None:
        labels = []
        for substr in path_str.split(','):
            path = bandpath(substr, atoms.get_cell()[...], npoints=1)

            _, _, substr_labels = path.get_linear_kpoint_axis()
            labels.extend(['$\\Gamma$' if s == 'G' else s for s in substr_labels])

    qpoints, connections = get_band_qpoints_and_path_connections(path_frac, npoints=npoints)
    return qpoints, labels, connections

def phonopy_pre_process(cell, supercell_matrix, displacement_distance):
    cell = convert_atoms_to_phonopy(cell)
    phonon = Phonopy(cell, supercell_matrix, log_level=1)
    phonon.generate_displacements(distance=displacement_distance)
    print("[Phonopy] Atomic displacements:")
    disps = phonon.get_displacements()
    for d in disps:
        print("[Phonopy] %d %s" % (d[0], d[1:]))
    return phonon

def run_gpaw_all(calc, phonon):
    return [ run_gpaw(calc, supercell) for supercell in phonon.get_supercells_with_displacements() ]

def run_gpaw(calc, cell):
    cell = convert_atoms_to_ase(cell)
    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]
    return forces

#--------------------------------------------------------------------

def load_or_compute_force_constants(path, calc, phonon, sum_rule):
    if os.path.exists(path):
        print('Reading FCs from {!r}'.format(path))
        phonon.force_constants = np.load(path)

    else:
        print('Computing FCs')
        os.makedirs('force-sets', exist_ok=True)
        supercells = list(phonon.get_supercells_with_displacements())
        fnames = ['force-sets/sc-{:04}.npy'.format(i) for i in range(len(supercells))]
        set_of_forces = [
            load_or_compute_force(fname, calc, supercell)
            for (fname, supercell) in zip(fnames, supercells)
        ]
        print('Building FC matrix')
        phonon.produce_force_constants(forces=set_of_forces, calculate_full_force_constants=False)
        if sum_rule:
            phonon.symmetrize_force_constants()

        print('Writing FCs to {!r}'.format(path))
        np.save(path, phonon.get_force_constants())
        shutil.rmtree('force-sets')

def load_or_compute_force(path, calc, atoms):
    if os.path.exists(path):
        print('Reading {!r}'.format(path))
        return np.load(path)

    else:
        print('Computing {!r}'.format(path))
        force_set = run_gpaw(calc, atoms)
        np.save(path, force_set)
        return force_set

#--------------------------------------------------------------------

def convert_atoms_to_ase(atoms):
    return Atoms(
        symbols=atoms.get_chemical_symbols(),
        scaled_positions=atoms.get_scaled_positions(),
        cell=atoms.get_cell(),
        pbc=True
    )

def convert_atoms_to_phonopy(atoms):
    return PhonopyAtoms(
        symbols=atoms.get_chemical_symbols(),
        scaled_positions=atoms.get_scaled_positions(),
        cell=atoms.get_cell(),
        pbc=True
    )

def phonopy_post_process(phonon):
    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

    # DOS
    phonon.set_mesh([21, 21, 21])
    phonon.set_total_DOS(tetrahedron_method=True)
    print('')
    print("[Phonopy] Phonon DOS:")
    for omega, dos in np.array(phonon.get_total_DOS()).T:
        print("%15.7f%15.7f" % (omega, dos))

#--------------------------------------------------------------------

def die(*args):
    print('FATAL:', *args)
    sys.exit(1)

if __name__ == "__main__":
    main()