import logging
import warnings
from abc import abstractmethod
from collections.abc import Callable
from typing import Any
import numpy as np
from ase.atoms import Atoms
from numpy.polynomial.polynomial import polyfit as np_polyfit
from directupsampling.data import Ang, THz_to_eV, amu, eV, hbar, kB
from directupsampling.effqh.fcfit import ForceConstantFit
logger = logging.getLogger(__name__)
def _ensure_positive(array: np.typing.ArrayLike) -> None:
"""Raise ValueError if any element of array is negative."""
if np.any(np.asarray(array) < 0):
raise ValueError(array)
def _prepare_input(
in1: np.typing.ArrayLike,
in2: np.typing.ArrayLike,
) -> tuple[np.ndarray, np.ndarray]:
"""Validate and broadcast two non-negative inputs to float arrays."""
_ensure_positive(in1)
_ensure_positive(in2)
return np.broadcast_arrays(
np.asarray(in1, dtype=float),
np.asarray(in2, dtype=float),
)
_ZERO_TOLERANCE_KELVIN = 1e-15
_OVERFLOW_TOLERANCE = 200
[docs]
def evib_classical(
temperature: np.typing.ArrayLike,
frequency: np.typing.ArrayLike = None,
) -> np.ndarray:
"""Classical harmonic internal energy: k_B * T.
Parameters
----------
temperature :
Temperature in Kelvin.
frequency :
Ignored; present for interface compatibility with the QM form.
Returns
-------
np.ndarray
k_B * T.
"""
_ensure_positive(temperature)
return kB * np.asarray(temperature, dtype=float)
[docs]
def fvib_classical(
temperature: np.typing.ArrayLike,
frequency: np.typing.ArrayLike,
) -> np.ndarray:
r"""Classical harmonic vibrational Helmholtz energy.
Parameters
----------
temperature :
Temperature in Kelvin.
frequency :
Frequency in eV.
Returns
-------
np.ndarray
k_B * T * ln(omega / (k_B*T)).
"""
t, f = _prepare_input(temperature, frequency)
p = t > _ZERO_TOLERANCE_KELVIN
res = np.zeros_like(t, dtype=float)
res[p] += kB * t[p] * np.log(f[p] / (kB * t[p]))
return res
[docs]
def evib_qm(
temperature: np.typing.ArrayLike,
frequency: np.typing.ArrayLike,
) -> np.ndarray:
r"""Quantum harmonic internal energy.
Parameters
----------
temperature :
Temperature in Kelvin.
frequency :
Frequency in eV.
Returns
-------
np.ndarray
omega/2 + omega / (exp(omega / (k_B*T)) - 1).
At T=0 returns the zero-point energy omega/2.
"""
t, f = _prepare_input(temperature, frequency)
safe_t = t.copy()
safe_t[t == 0] += _ZERO_TOLERANCE_KELVIN
x = f / (kB * safe_t)
p = x < _OVERFLOW_TOLERANCE
res = np.asarray(f / 2, dtype=float)
res[p] += f[p] / np.expm1(x[p])
return res
[docs]
def fvib_qm(
temperature: np.typing.ArrayLike,
frequency: np.typing.ArrayLike,
) -> np.ndarray:
r"""Quantum harmonic vibrational Helmholtz energy.
Parameters
----------
temperature :
Temperature in Kelvin.
frequency :
Frequency in eV.
Returns
-------
np.ndarray
omega/2 + k_B*T * ln(1 - exp(-omega / (k_B*T))).
At T=0 returns the zero-point energy omega/2.
"""
t, f = _prepare_input(temperature, frequency)
p = t > _ZERO_TOLERANCE_KELVIN
res = np.asarray(f / 2, dtype=float)
res[p] += kB * t[p] * np.log(-np.expm1(-f[p] / (kB * t[p])))
return res
[docs]
class BasePhonons:
"""Base class for phonon models fitted as polynomial functions of volume.
Subclasses implement ``force_constants_to_frequencies`` for a specific
diagonalisation scheme (Gamma-point exact diagonalisation or BZ mesh
via phonopy).
"""
FREQ_WARN_TOL = 1e-5
def __init__(
self,
atoms: Atoms,
volume_bounds: list | None = None,
freq_coeffs: np.ndarray | None = None,
weights: np.ndarray | None = None,
*,
invert_neg: bool = False,
freq_cutoff: float = 1e-5,
) -> None:
"""Initiate the phonon model.
Parameters
----------
atoms :
Reference structure.
volume_bounds :
[V_min, V_max] in ų defining the valid range of the polynomial fit.
freq_coeffs :
Polynomial coefficients with shape (order, n_freq) mapping volume
to frequencies in eV.
weights :
k-point weights aligned with the frequency array.
invert_neg :
If True, reflect negative frequencies to |omega| before evaluation.
The negative-frequency warning is suppressed for inverted modes.
freq_cutoff :
Frequencies at or below this value (eV) are excluded from sums.
"""
self.atoms = atoms
self.volume_bounds = volume_bounds
self.freq_coeffs = freq_coeffs
self.weights = weights
self.invert_neg = invert_neg
self.freq_cutoff = freq_cutoff
[docs]
def todict(self) -> dict:
dct = self.__dict__.copy()
dct.pop("_phonopy_handler", None)
return dct
[docs]
@classmethod
def from_fcfit(
cls,
atoms: Atoms,
force_constant_fit: ForceConstantFit,
mesh: list | None = None,
*,
invert_neg: bool = False,
freq_cutoff: float = 1e-5,
) -> "BasePhonons":
"""Construct a phonon model by fitting frequencies from a ForceConstantFit.
Parameters
----------
atoms :
Reference structure.
force_constant_fit :
Fitted force constants as a function of volume.
mesh :
BZ mesh passed to MeshPhonons; omit for ExactPhonons.
invert_neg :
Forwarded to ``__init__``.
freq_cutoff :
Forwarded to ``__init__``.
Returns
-------
BasePhonons
Instance with ``freq_coeffs`` and ``weights`` set.
"""
init_kwargs = {"invert_neg": invert_neg, "freq_cutoff": freq_cutoff}
if mesh is None:
new = cls(atoms, force_constant_fit.volume_bounds, **init_kwargs)
else:
new = cls(atoms, force_constant_fit.volume_bounds, mesh=mesh, **init_kwargs)
volumes = force_constant_fit.raw_data["volume"]
fq_list = []
for v in volumes:
fc = force_constant_fit(v)
fq = new.force_constants_to_frequencies(fc)
fq_list.append(fq)
coefficients = np_polyfit(volumes, fq_list, force_constant_fit.fit_order)
new.freq_coeffs = coefficients
return new
[docs]
@abstractmethod
def force_constants_to_frequencies(
self,
force_constants: np.ndarray,
) -> np.ndarray:
"""Return frequencies in eV and set ``self.weights``.
Parameters
----------
force_constants :
Force-constant matrix for the structure.
Returns
-------
np.ndarray
Frequencies in eV, sorted in ascending order.
"""
...
[docs]
def get_frequencies(self, volume: np.typing.ArrayLike) -> np.ndarray:
"""Evaluate the polynomial frequency model at the given volume(s).
Parameters
----------
volume :
Volume(s) in ų.
Returns
-------
np.ndarray
Frequencies in eV with shape (n_freq, \*volume.shape).
"""
v = np.asarray(volume, dtype=float)
freqs = np.zeros(self.freq_coeffs.shape[1:2] + v.shape)
coeffs = self.freq_coeffs.reshape(self.freq_coeffs.shape + v.ndim * (1,))
for i, c in enumerate(coeffs):
freqs += c * v**i
return freqs
[docs]
def get_energy(
self,
volume: np.typing.ArrayLike,
temperature: np.typing.ArrayLike,
*,
classical: bool = False,
) -> np.ndarray:
"""Harmonic internal energy per atom.
Parameters
----------
volume :
Volume(s) in ų.
temperature :
Temperature(s) in Kelvin.
classical :
Use the classical (k_B*T) expression instead of quantum.
Returns
-------
np.ndarray
Internal energy in eV/atom.
"""
evib = evib_classical if classical else evib_qm
return self.sum_over_frequencies(evib, volume, temperature)
[docs]
def get_free_energy(
self,
volume: np.typing.ArrayLike,
temperature: np.typing.ArrayLike,
*,
classical: bool = False,
) -> np.ndarray:
"""Harmonic Helmholtz energy per atom.
Parameters
----------
volume :
Volume(s) in ų.
temperature :
Temperature(s) in Kelvin.
classical :
Use the classical expression instead of quantum.
Returns
-------
np.ndarray
Helmholtz energy in eV/atom.
"""
fvib = fvib_classical if classical else fvib_qm
return self.sum_over_frequencies(fvib, volume, temperature)
[docs]
def sum_over_frequencies(
self,
fun: Callable[[np.ndarray, np.ndarray], np.ndarray],
volume: np.typing.ArrayLike,
temperature: np.typing.ArrayLike,
) -> np.ndarray:
"""Weighted sum of ``fun`` over modes above ``freq_cutoff``.
Normalised per vibrational atom (see denominator note below).
Modes at or below ``freq_cutoff`` are excluded from the sum via a
binary mask; they contribute zero to the numerator regardless of
``fun``.
The denominator is ``weights.sum() / 3 - 1``, equal to N_atoms - 1
(ExactPhonons) or N_atoms * N_kpoints - 1 (MeshPhonons). This
normalises per *vibrational* atom: after removing the 3 translational
modes the system has 3N-3 physical degrees of freedom, giving N-1
effective 3-DOF "atoms". This matches the MD convention where COM
kinetic energy is subtracted before normalising. For large supercells
the difference from 1/N is negligible; for small unit cells (e.g.
N=2) the factor of 2 is intentional and ensures consistency with MD.
Parameters
----------
fun :
Function ``(temperature, frequency) -> energy``.
volume :
Volume(s) in ų.
temperature :
Temperature(s) in Kelvin.
Returns
-------
np.ndarray
Weighted sum of ``fun`` over active frequencies, per atom.
Warns
-----
RuntimeWarning
If any frequency remains below ``-FREQ_WARN_TOL`` after optional
inversion.
"""
v, t = _prepare_input(volume, temperature)
f = self.get_frequencies(v)
if self.invert_neg:
f = np.abs(f)
tol = self.FREQ_WARN_TOL
n_neg = int(np.sum(f < -tol))
if n_neg > 0:
msg = f"Hessian gave {n_neg} negative frequencies."
logger.warning(msg)
warnings.warn(msg, RuntimeWarning, 2)
w = self.weights.reshape(self.weights.shape + (1,) * v.ndim)
m = (f > self.freq_cutoff).astype(float)
f[~m.astype(bool)] = self.freq_cutoff
return np.sum(fun(t, f) * w * m, axis=0) / (self.weights.sum() / 3 - 1)
[docs]
class ExactPhonons(BasePhonons):
"""Gamma-point phonons via exact diagonalisation of the dynamical matrix."""
directupsampling_objtype = "exactphonons"
[docs]
def force_constants_to_frequencies(
self,
force_constants: np.ndarray,
) -> np.ndarray:
"""Diagonalise the dynamical matrix and return frequencies in eV.
Uses ``eigvalsh`` (real-symmetric solver) for numerical stability.
Negative eigenvalues are mapped to negative frequencies via
``sqrt(|lambda|) * sign(lambda)`` rather than producing imaginary values.
Sets ``self.weights = ones(3N)``.
Parameters
----------
force_constants :
Force-constant matrix for the structure.
Returns
-------
np.ndarray
Frequencies in eV sorted in ascending order, shape (3N,).
"""
masses = self.atoms.get_masses()
dynmat_to_freq = np.sqrt(eV / Ang**2 / amu) * hbar # 0.064654148 eV
mass_squared = np.outer(masses, masses)
dyn_mat = force_constants / np.sqrt(mass_squared)[..., None, None]
n0, n1 = dyn_mat.shape[0], dyn_mat.shape[1]
dyn_mat = dyn_mat.transpose([0, 2, 1, 3]).reshape(n0 * 3, n1 * 3)
eigenvalues = np.linalg.eigvalsh(dyn_mat)
frequencies = np.sort(
np.sqrt(np.abs(eigenvalues)) * np.sign(eigenvalues) * dynmat_to_freq,
)
self.weights = np.ones_like(frequencies)
return frequencies
[docs]
class MeshPhonons(BasePhonons):
"""BZ-mesh phonons computed via phonopy."""
directupsampling_objtype = "meshphonons"
def __init__(
self,
*args: Any,
mesh: list,
**kwargs: Any,
) -> None:
"""Initiate the phonon model with a specified BZ sampling mesh.
Parameters
----------
*args :
Positional arguments forwarded to ``BasePhonons.__init__``.
mesh :
BZ sampling mesh, e.g. ``[8, 8, 8]``.
**kwargs :
Keyword arguments forwarded to ``BasePhonons.__init__``.
"""
super().__init__(*args, **kwargs)
self.mesh = mesh
[docs]
def force_constants_to_frequencies(
self,
force_constants: np.ndarray,
) -> np.ndarray:
"""Compute BZ-mesh frequencies via phonopy and return them in eV.
Sets ``self.weights`` to the phonopy k-point weights repeated over
branches (shape: n_qpoints * 3 * N_primitive).
Parameters
----------
force_constants :
Force-constant matrix for the structure.
Returns
-------
np.ndarray
Frequencies in eV, shape (n_qpoints * 3 * N_primitive,).
"""
ph = PhonopyHandler(self.atoms, force_constants, self.mesh)
frequencies = ph.get_frequencies()
self.weights = ph.get_weights()
self._phonopy_handler = ph
return frequencies
[docs]
class PhonopyHandler:
"""Thin wrapper around a phonopy Phonopy object for a single structure."""
def __init__(self, atoms: Atoms, force_constants: np.ndarray, mesh: list) -> None:
"""Initiate the phonopy handler.
Parameters
----------
atoms :
Structure to compute phonons for.
force_constants :
Force-constant matrix.
mesh :
BZ sampling mesh, e.g. ``[8, 8, 8]``.
"""
self.atoms = atoms
self.force_constants = force_constants
self.mesh = mesh
self._phonons = None
[docs]
def get_frequencies(self) -> np.ndarray:
"""Return flattened mesh frequencies in eV."""
return np.ravel(self.phonons.get_mesh_dict()["frequencies"]) * THz_to_eV
@property
def phonons(self):
if self._phonons is None:
self._phonons = self.make_phonons()
return self._phonons
[docs]
def get_weights(self) -> np.ndarray:
"""Return k-point weights repeated over branches (3 * N_primitive per q-point)."""
p = self.phonons
return np.repeat(p.get_mesh_dict()["weights"], 3 * len(p.primitive))
[docs]
def make_phonons(self):
"""Build and run the phonopy Phonopy object."""
from phonopy import Phonopy
from phonopy.structure.atoms import PhonopyAtoms
atoms = self.atoms
phonopy_atoms = PhonopyAtoms(
cell=atoms.cell,
positions=atoms.positions,
symbols=atoms.get_chemical_symbols(),
)
phonons = Phonopy(phonopy_atoms, np.eye(3), primitive_matrix="auto")
phonons.force_constants = self.force_constants
phonons.run_mesh(self.mesh)
return phonons
[docs]
def plot_phonon_dispersion(self):
self.phonons.auto_band_structure()
self.phonons.run_total_dos()
plt = self.phonons.plot_band_structure_and_dos()
plt.title(f"{self.atoms.get_volume() / len(self.atoms)} Ang3/atom")