Generate phonon trajectory from Phonopy's band.yaml

plot-phonon.py

#!/usr/bin/env python3
"""
A script to generate phonon trajectory from Phonopy's band.yaml.
Requirement: python3, numpy, PyYAML, ase
Author: @Ionizing
Acknowledgement: @QijingZheng
Date: 16:57, Jan 20th, 2025
"""


import sys
from argparse import ArgumentParser
import logging
import numpy as np
from typing import List
# from numpy.typing import NDArray
# import matplotlib.pyplot as plt
import yaml
from yaml import CLoader

# import ase
from ase.atoms import Atoms
from ase.atom import Atom
from ase.build import make_supercell


logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


HBAR   = 0.6582119514467407    # hbar, in (eV * fs)
CLIGHT = 2997.92458            # light speed, (Angstrom/fs)
KBOLTZ = 8.617330337217213e-05 # kB, Boltzmann factor, eV/K
MPROTON = 938272081.4796857    # proton mass, in eV/c^2
MPROTON = MPROTON / CLIGHT**2  # proton mass, in eV/(A/fs)^2


def msd_classical(w, m, T=300.0):
    """
    The classical mean square displacement (MSD) of a harmonic oscillator at
    temperature "T" with frequency "w" and mass "m".

       MSD_c:
           < x**2 > = k_B T / (M w**2)

    Inputs:
            w:  the frequency of the HO, in THz.
            m:  the mass of the HO in unified atomic mass unit. amu.
            T:  the temperature in Kelvin.
    freq_unit:  the unit of the frequency.

    Return:
        MSD in Angstrom**2

    Copied from https://github.com/QijingZheng/VaspVib2XSF/blob/41877166ff093469871631b0c37c904cec36842e/phonon_traj.py#L11
    """
    w = np.asarray(w, dtype=float)
    w = 2 * np.pi * w / 1000   # angular frequency, rad/fs
    T = np.asarray(T, dtype=float)
    return KBOLTZ * T / (m * MPROTON * w**2)


def msd_quantum(w, m, T=300.0, n=None):
    """
    The quantum mean square displacement (MSD) of a harmonic oscillator at
    temperature "T" with frequency "w" and mass "m".

       MSD_c:
                        hbar         hbar w
           < x**2 > = -------- coth(---------)
                       2 M w         2 k_B T

    Inputs:
            w:  the frequency of the HO, in THz.
            m:  the mass of the HO in unified atomic mass unit.
            T:  the temperature in Kelvin.
            n:  the number of phonons

    Return:
        MSD in Angstrom**2

    Copied from: https://github.com/QijingZheng/VaspVib2XSF/blob/41877166ff093469871631b0c37c904cec36842e/phonon_traj.py#L43
    """
    w = np.asarray(w, dtype=float)
    w = 2 * np.pi * w / 1000    # angular frequency, rad/fs
    if n is None:
        n = np.zeros_like(T, dtype=float)
        if T < 1E-12:
            n = 0.0
        else:
            n = 1.0 / (np.exp(HBAR * w / (KBOLTZ * T)) - 1.0)
    return HBAR / (2 * m * MPROTON * w) * (1.0 + 2*n)


class Phonon:
    def __init__(self, fname: str="band.yaml"):
        """
        Loading data from band.yaml.

        The band.yaml should be generated by phonopy-load --config band.conf
        where "EIGENVECTORS=.TRUE." should be present in band.conf
        """
        logger.info(f"Reading \"{fname}\" ...")

        dat = yaml.load(open(fname), Loader=CLoader)
        for (k, v) in dat.items():
            setattr(self, k, v)
            pass
        pass

        if 'eigenvector' not in dat["phonon"][0]['band'][0].keys():
            raise ValueError(f"\"eigenvector\" not present in {fname}, please regenerate it with \"EIGENVECTORS = .TRUE.\" ")

        self.lattice = np.array(self.lattice)
        self.reciprocal_lattice = np.array(self.reciprocal_lattice)
        self.qpositions = np.array([p['q-position'] for p in self.phonon])
        self.frequencies = np.array([[freq['frequency'] for freq in phonon['band']] for phonon in self.phonon])
        eigenvectors = np.array([[freq['eigenvector'] for freq in phonon['band']] for phonon in self.phonon])
        self.eigenvectors = eigenvectors[...,0] + 1j * eigenvectors[...,1]


    def get_amplitude(self, w: float, T: float, supcell: Atoms, *, msd_mode="quantum"):
        """
        Generate amplitude for each atom. See documentation of msd_classical and msd_quantum for detailed explanation.
        """
        m = supcell.get_masses()
        if msd_mode == "quantum":
            msd = msd_classical(w, m, T)
        elif msd_mode == "classical":
            msd = msd_quantum(w, m, T)
        else:
            raise ValueError("Invalid msd_mode, only quantum and classical available.")

        A = np.sqrt(2 * msd * m.size)
        return A


    def get_phonon(self, iq: int=0):
        """
        Return qposition (in fractional) and frequencies (in THz) at given q point index.
        """
        q = self.qpositions[iq]
        freqs = [b['frequency'] for b in self.phonon[iq]['band']]
        return (q, freqs)


    def get_primitivecell(self) -> Atoms:
        """
        Return an Atoms object of current structure
        """
        lattice = self.lattice
        atoms   = Atoms([Atom(symbol=atom['symbol'], position=atom['coordinates'] @ np.array(self.lattice), mass=atom['mass'])
                         for atom in self.points], cell=lattice, pbc=True)
        return atoms


    def get_supercell(self, M=[[1, 0, 0], [0, 1, 0], [0, 0, 1]], *, order="atom-major"):
        """
        Return an Atoms object of supercelled structure, where the supercell is defined by matrix M.

        order: whether sort the atoms by their symbol after make_supercell
        """
        primcell = self.get_primitivecell()
        return make_supercell(primcell, P=M, order=order)


    def get_eigenvecotr(self, iq=0, imode=0, M=[[1, 0, 0], [0, 1, 0], [0, 0, 1]], *, order="atom-major"):
        """
        Returns the eigenvector of the supercell.
        """
        eigvec = self.eigenvectors[iq, imode, :, :] # natom x 3, complex array

        primcell_re = self.get_primitivecell()
        primcell_re.set_velocities(eigvec.real)
        supcell_re = make_supercell(primcell_re, P=M, order=order)

        primcell_im = primcell_re.copy()
        primcell_im.set_velocities(eigvec.imag)
        supcell_im = make_supercell(primcell_im, P=M, order=order)

        supcell_eigvec = supcell_re.get_velocities() + 1j * supcell_im.get_velocities()
        return supcell_eigvec


    def get_phonon_trajectory(self, *,
                              iq: int = 0,                         # qpoint index
                              imode: int = 0,                      # mode index
                              M=[[1, 0, 0], [0, 1, 0], [0, 0, 1]], # supercell matrix
                              nframes=None,                        # None for one period
                              timestep: float = 1.0,               # fs
                              temperature: float = 300.0,          # Kelvin
                              ) -> List[Atoms]:
        """
        Generate the trajectory for specific phonon mode.
        """
        q, freqs = self.get_phonon(iq=iq)
        mode = freqs[imode] # in THz
        omega = mode * np.pi * 2 / 1000
        supercell = self.get_supercell(M)
        eigvecs   = self.get_eigenvecotr(iq=iq, imode=imode, M=M)
        
        q = q @ self.reciprocal_lattice     # in 1/Angstrom
        r = supercell.get_positions()       # in Angstrom
        A = self.get_amplitude(w=mode, T=temperature, supcell=supercell)

        if nframes is None:
            nframes = int(np.ceil(1000 / mode))

        logger.info("Generating trajectory ...")
        traj = []
        for iframe in range(nframes):
            t = iframe * timestep
            expiqr = np.exp(2j * np.pi * (r @ q))
            expiwt = np.exp(1j* omega * t)
            phase = expiwt * expiqr
            disps = A[:,None] * np.real(phase[:, None] * eigvecs)
            thisframe = supercell.copy()
            thisframe.set_positions(r + disps)
            traj.append(thisframe)

        return traj


    def save_trajectory(self, *,
                        out_prefix: str = "phonon",
                        out_format: str = "XDATCAR",
                        iq: int = 0,                          # qpoint index
                        imode: int = 0,                       # mode index
                        M=[[1, 0, 0], [0, 1, 0], [0, 0, 1]],  # supercell matrix
                        nframes=None,                         # None for one period
                        timestep: float = 1.0,                # fs
                        temperature: float = 300.0,           # Kelvin
                        ):
        from ase.io.vasp import write_vasp_xdatcar
        from ase.io.xyz import write_xyz

        traj = self.get_phonon_trajectory(iq=iq, imode=imode, M=M, nframes=nframes, timestep=timestep, temperature=temperature)
        
        if out_format.lower() == "xdatcar":
            out_prefix = f"XDATCAR_{out_prefix}"
            out_extension = "vasp"
            writer = write_vasp_xdatcar
        elif out_format.lower() == "xyz":
            out_extension = "xyz"
            writer = write_xyz
        else:
            raise ValueError(f"Unsupported outfile format: {out_format}. Only XDATCAR and xyz are supported for now.")

        _, freqs = self.get_phonon(iq=iq)
        mode = freqs[imode]
        fname = f"{out_prefix}_iq{iq:03d}_imode{imode:03d}_{mode:.2f}THz_T{temperature}.{out_extension}"
        comment = fname
        fd = open(fname, "w")

        logger.info(f"Writing to {fname} ...")
        writer(fd, traj, comment)
        pass


def parse_cml_argument():
    parser = ArgumentParser(add_help=True,
                            description="A script to generate phonon trajectory from Phonopy's band.yaml.",)
    parser.add_argument("input", type=str, action="store", default="band.yaml", nargs="?",
                        help="File name of band.yaml. Defult by band.yaml. Remember adding \"EIGENVECTORS=.TRUE.\" to \"band.conf\"",)
    parser.add_argument("--outprefix", type=str, action="store", default="phonon",
                        help="Prefix of output filename. Default by phonon",)
    parser.add_argument("--outformat", type=str, action="store", default="XDATCAR",
                        help="Format of output file. Default by XDATCAR", choices=["XDATCAR", "xyz"])
    parser.add_argument("-q", "--iqpoint", type=int, action="store", default=[0], nargs="+",
                        help="Index of q-point defined in band.conf, counting from 0. Default by 0.",)
    parser.add_argument("-m", "--imode", type=int, action="store", default=0, nargs="+",
                        help="Index of mode for q-point, counting from 0. Default by 0.",)
    parser.add_argument("-M", type=int, action="store", nargs="+", default=[1],
                        help="Supercell matrix. Can be 1, 3 and 9 elements. Row major. Default by 1",)
    parser.add_argument("-N", "--nframes", action="store", type=int, default=None,
                        help="How many frames to generate. Default by one period for selected mode.",)
    parser.add_argument("--timestep", action="store", type=float, default=1.0,
                        help="Time step between two frames, in fs. Default by 1.0 fs",)
    parser.add_argument("-T", "--temperature", action="store", type=float, default=[300.0], nargs="+",
                        help="Temperature for the trajectory, in Kelvin. Default by 300 K",)

    args = parser.parse_args(args=None if sys.argv[1:] else ['--help'])

    if len(args.M) == 1:
        M = np.eye(3, dtype=int) * args.M[0]
    elif len(args.M) == 3:
        M = np.diag(args.M)
    elif len(args.M) == 9:
        M = np.array(args.M).reshape(3, 3)
    else:
        raise ValueError("Length of supercell matrix could only be 1, 3 or 9.")
    args.M = M

    return args


if "__main__" == __name__:
    args = parse_cml_argument()
    phonon = Phonon(fname=args.input)

    for iq in args.iqpoint:
        for imode in args.imode:
            for T in args.temperature:
                phonon.save_trajectory(out_prefix=args.outprefix,
                                       out_format=args.outformat,
                                       iq=iq,
                                       imode=imode,
                                       M=args.M,
                                       nframes=args.nframes,
                                       timestep=args.timestep,
                                       temperature=T)

    # M = np.array([[4, 2, 0],
                  # [0, 3, 0],
                  # [0, 0, 1]])

    # phonon = Phonon()
    # phonon.save_trajectory(iq=242, imode=0, M=M)

Using Ovito to visualize the animation: