from __future__ import annotations
import logging
import pathlib
import warnings
from typing import TYPE_CHECKING, Self
import numpy as np
import numpy.typing as npt
from numpy.polynomial.polynomial import polyfit as np_polyfit
from directupsampling.effqh.io import read_fc_coeffs as _read_fc_coeffs
from directupsampling.effqh.utils import get_ref_atoms
from directupsampling.io.io import read as _io_read
from directupsampling.io.io import write as _io_write
from directupsampling.parallel import parcall
if TYPE_CHECKING:
import os
from ase.atoms import Atoms
from hiphive import ForceConstantPotential
logger = logging.getLogger(__name__)
[docs]
class ForceConstantFit:
"""Class for fitting force constants."""
directupsampling_objtype = "forceconstantfit"
def __init__(self, fit_order: int | None = None) -> None:
"""Initialize a ForceConstantFit.
Parameters
----------
fit_order : int, optional
The fitting order of polynomials. If None, the order is 3 or less.
Attributes
----------
atoms : Atoms
Reference supercell structure the force constants are sized to.
fit_order : int
The fitting order of polynomials.
coefs : npt.NDArray[np.float64]
Fitting coefficients of the elements of force constants.
The shape is (fit_order + 1, natoms, natoms, 3, 3).
raw_data : tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]
Tuple of raw volumes and force constants.
fit_results : dict
Contains the sum of squared residuals per force constant "residuals",
"max_residual" and the RMSE "rmse".
"""
self.atoms = None
self.fit_order = fit_order
self.volume_bounds = (None, None)
self.raw_data = None
self.fit_results = None
self._fc_coeffs = None
@property
def coefficients(self) -> npt.NDArray[np.float64]:
"""Fitted force constant coefficients."""
return self._fc_coeffs
def __call__(self, volume: npt.ArrayLike) -> npt.NDArray[np.float64]:
"""Get force constants at the given volume(s) by interpolation."""
return self.get_force_constants(volume)
[docs]
def get_force_constants(self, volume: npt.ArrayLike) -> npt.NDArray[np.float64]:
"""Get force constants at the given volume(s) by interpolation."""
v = np.array(volume, dtype=float)
if self.coefficients.shape[0] == 1:
return self.coefficients[0, ...]
fc_coefs = self.coefficients
coefs = fc_coefs.reshape(fc_coefs.shape + v.ndim * (1,))
fcs = np.zeros(coefs.shape[1:5] + v.shape)
for i, c in enumerate(coefs):
fcs += c * v**i
return fcs
[docs]
def fit_force_constants(self, volumes, force_constants):
"""Fit each element of force constants to polynomials of volumes"""
volumes = np.array(list(volumes), dtype=float)
self.raw_data = {"volume": volumes, "force_constant": force_constants}
vmin = np.min(self.raw_data["volume"])
vmax = np.max(self.raw_data["volume"])
self.volume_bounds = (vmin - (vmax - vmin) * 0.1, vmax + (vmax - vmin) * 0.1)
if self.fit_order is None:
self.fit_order = len(volumes) - 1 if len(volumes) < 4 else 3
fc_coefs, residuals = _fit_force_constants(
volumes, force_constants, self.fit_order
)
if len(residuals) == 0:
msg = (
"Non-overdetermined fitting of FCs over volume: "
f"number of volumes: {len(volumes)}, order: {self.fit_order}"
)
parcall(logger.warning, msg)
residuals = np.array([np.nan])
self._fc_coeffs = fc_coefs
self.fit_results = {
"residuals": residuals,
"max_residual": np.max(residuals),
"rmse": (np.sqrt(np.mean(residuals)) / residuals.size),
}
[docs]
def todict(self) -> dict:
return {
"atoms": self.atoms,
"fit_order": self.fit_order,
"volume_bounds": self.volume_bounds,
"raw_data": self.raw_data,
"fit_results": self.fit_results,
"coefficients": self.coefficients,
}
[docs]
@classmethod
def fromdict(cls, dct: dict) -> Self:
"""Reconstruct a :class:`ForceConstantFit` from its ``todict()`` form.
Returns
-------
ForceConstantFit
"""
dct = _migrate_legacy_fcfit_dict(dict(dct))
new = cls(_as_int(dct.get("fit_order")))
new.atoms = dct.get("atoms")
vb = dct.get("volume_bounds")
new.volume_bounds = tuple(vb) if vb is not None else (None, None)
new.raw_data = dct.get("raw_data")
new.fit_results = dct.get("fit_results")
new._fc_coeffs = dct.get("coefficients")
return new
[docs]
@classmethod
def fit_from_hiphive(
cls,
volumes: list[float],
force_constant_potentials: list[ForceConstantPotential],
atoms: Atoms | list[Atoms],
fit_order: int | None = None,
) -> Self:
"""Fit a ForceConstantFit from a list of hiphive ForceConstantPotentials.
Parameters
----------
volumes: list of float
List of volumes corresponding to the ForceConstantPotentials.
force_constant_potentials: list of ForceConstantPotential
List of hiphive ForceConstantPotentials to fit to.
atoms: Atoms or list of Atoms
Reference Atoms or list of Atoms to get the force constants with.
fit_order: int, optional
The fitting order of polynomials. If None, the order is 3 or less.
Returns
-------
ForceConstantFit
"""
parcall(logger.debug, "Fitting to force constants from fcps ...")
force_constants = []
for vol, fcp in zip(volumes, force_constant_potentials, strict=True):
ref_atoms = get_ref_atoms(atoms, vol)
fc = fcp.get_force_constants(ref_atoms)
force_constants.append(fc.get_fc_array(order=2))
new = cls(fit_order)
new.atoms = atoms if hasattr(atoms, "get_masses") else atoms[0]
new.fit_force_constants(volumes, force_constants)
return new
[docs]
def write(self, file: os.PathLike | str) -> None:
file = pathlib.Path(file)
if file.suffix not in {".fcfit", ".hdf5", ".h5"}:
file = file.with_suffix(".fcfit")
_io_write(self, file)
[docs]
@classmethod
def read(cls, filename: os.PathLike | str) -> Self:
"""Read force constants from a .fcfit (HDF5) file or phonopy txt file(s).
Parameters
----------
filename : str or list of str
The filename(s) to read.
"""
try:
return _io_read(filename, cls=cls)
except OSError:
# Not an HDF5 file -- fall back to the phonopy txt format.
poly_coeffs = _read_fc_coeffs(filename)
fit_order = poly_coeffs.shape[0] - 1
new = cls(fit_order)
new._fc_coeffs = poly_coeffs
return new
def _as_int(value: int | None) -> int | None:
return None if value is None else int(value)
def _migrate_legacy_fcfit_dict(dct: dict) -> dict:
"""Rename legacy keys in an fcfit dict, emitting DeprecationWarnings.
Returns
-------
dict
"""
for old in ("coeffs", "polynomial_coefficients"):
if old in dct:
if "coefficients" not in dct:
warnings.warn(
f"Legacy fcfit key '{old}' found; "
"regenerate the file to use 'coefficients'.",
DeprecationWarning,
stacklevel=3,
)
dct["coefficients"] = dct.pop(old)
else:
dct.pop(old)
if "volume_bounds" not in dct and "vmin" in dct and "vmax" in dct:
warnings.warn(
"Legacy fcfit keys 'vmin'/'vmax' found; "
"regenerate the file to use 'volume_bounds'.",
DeprecationWarning,
stacklevel=3,
)
dct["volume_bounds"] = np.array([dct.pop("vmin"), dct.pop("vmax")])
else:
dct.pop("vmin", None)
dct.pop("vmax", None)
raw_data = dct.get("raw_data")
if isinstance(raw_data, dict):
for old, new in (("volumes", "volume"), ("force_constants", "force_constant")):
if old in raw_data:
if new not in raw_data:
warnings.warn(
f"Legacy fcfit key 'raw_data/{old}' found; "
f"regenerate the file to use '{new}'.",
DeprecationWarning,
stacklevel=3,
)
raw_data[new] = raw_data.pop(old)
else:
raw_data.pop(old)
return dct
def _fit_force_constants(volumes, force_constants, order=3):
volumes = np.array(list(volumes), dtype=float)
force_constants = np.array(force_constants)
assert volumes.ndim == 1 and volumes.shape[0] == force_constants.shape[0]
natoms = force_constants.shape[1]
# Fit each FC value over volumes (mask where values over all volumes are zero)
fc_flat = force_constants.reshape(volumes.shape + (-1,))
where_nonzero = np.invert(np.isclose(fc_flat, 0, atol=1e-12).all(axis=0))
fc_reduced = fc_flat[:, where_nonzero]
poly_coef_reduced, fit_results = np_polyfit(volumes, fc_reduced, order, full=True)
poly_coef_flat = np.zeros((order + 1,) + (natoms * natoms * 3 * 3,))
poly_coef_flat[:, where_nonzero] = poly_coef_reduced
poly_coef = poly_coef_flat.reshape((order + 1,) + (natoms, natoms, 3, 3))
residuals = fit_results[0]
# Symmetrize the polynomial coefficients
poly_coef_sym = np.zeros([order + 1, natoms, natoms, 3, 3])
for i in range(order + 1):
poly_coef_sym[i] = _symmetrize_force_constants(poly_coef[i])
return poly_coef_sym, residuals
def _symmetrize_force_constants(force_constants):
n0 = force_constants.shape[0]
n1 = force_constants.shape[1]
fc_reshaped = force_constants.transpose([0, 2, 1, 3]).reshape(n0 * 3, n1 * 3)
fc_sym = (fc_reshaped + np.transpose(fc_reshaped)) / 2
fc_sym = fc_sym.reshape(n0, 3, n1, 3).transpose([0, 2, 1, 3])
return fc_sym
def _zerorize_force_constants(force_constants):
"""For the moment not used"""
n0 = force_constants.shape[0]
n1 = force_constants.shape[1]
force_constants = force_constants.transpose([0, 2, 1, 3]).reshape(n0 * 3, n1 * 3)
eigvals, eigvecs = np.linalg.eigh(force_constants)
eigvals_mod = np.where(eigvals < -1e-6, 1e-6, eigvals)
# eigvals_mod = np.where(eigvals < -1e-6, -1*eigvals, eigvals)
# Putting eigvals to ~0 seems sensitive to Fourier interpolation
# (probably?), inverting is therefore done for now
logger.debug("Zerorized modes:", np.where(eigvals != eigvals_mod)[0])
fc_wout_neg_freq = (eigvecs * eigvals_mod) @ eigvecs.T
# The transpose is used here because eigvecs is orthonormal (normalized
# Hermitian?) so the transpose equals the inverse.
fc_wout_neg_freq = fc_wout_neg_freq.reshape(n0, 3, n0, 3).transpose([0, 2, 1, 3])
return fc_wout_neg_freq