Ionizing · January 20, 2025 09:15 · Ionizing · Jan 20, 2025 · Ionizing · Jan 20, 2025
diff --git a/plot-phonon.py b/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)
	#!/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:
	< x2 > = k_B T / (M w2)

	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)