from __future__ import annotations
import logging
import pathlib
import warnings
from tarfile import ReadError
from typing import TYPE_CHECKING
import numpy as np
from ase.parallel import parallel_function
from hiphive import (
ClusterSpace,
ForceConstantPotential,
ForceConstants,
StructureContainer,
)
from hiphive.utilities import extract_parameters, get_displacements
from trainstation import Optimizer
from directupsampling.calculators.effqh import EffQHCalculator
from directupsampling.fecontrib import FreeEnergyContrib
from directupsampling.fesurface import FreeEnergySurface
from directupsampling.parallel import DummyMPI, parcall, world
from directupsampling.snapshots import SnapshotContainer
from directupsampling.utils import chunks, gather_chunks
from .fcfit import ForceConstantFit
from .io import write_force_constants
from .phonons import ExactPhonons, MeshPhonons
from .utils import get_ref_atoms
if TYPE_CHECKING:
import os
from ase.atoms import Atoms
logger = logging.getLogger(__name__)
[docs]
class EffQH:
"""Class representing an effective quasi-harmonic model."""
def __init__(
self,
atoms: Atoms | list[Atoms],
fit_order: int | None = None,
cutoffs: list[float] | None = None,
comm: DummyMPI = world,
) -> None:
"""Initialize an EffQH instance.
Parameters
----------
atoms: Atoms or list of Atoms
Ideal structure Atoms instance used as reference. May be a list with
an Atoms instance corresponding to each volume. In that case, the
EffQH instance is only valid at the exact volumes, and no
interpolation can be done.
fit_order: int or None
Order of the force constant fit over volume.
cutoffs : list of float
list of cutoffs for the effective harmonic model, in Ångström. The
cutoffs will correspond to the first volume given as input in
`fit_snapshots`, and is otherwise scaled by (volume /
volumes[0])^(1/3). Each element in the list corresponds to 2nd
order force constants, 3rd order... and so on. The default ([5])
corresponds only to 2nd order.
comm: MPI communicator
MPI communicator for parallelization.
Attributes
----------
atoms: Atoms
Reference Atoms instance.
fc_potentials: dict[float, ForceConstantPotential]
Dictionary with volume keys and ForceConstantPotential values.
fc_fit: ForceConstantFit
The ForceConstantFit instance containing coefficients for the force constants.
cutoffs: list[float]
Cutoffs in Ångström for the effective quasi-harmonic model obtained
through the ForceConstantPotentials. See Parameters.
fit_order: int
Order of the force constant fit over volume. Passed on to ForceConstantFit.
exact_phonons: ExactPhonons
Property returning an ExactPhonons instance.
mesh_phonons: MeshPhonons
Property returning a MeshPhonons instance.
phonon_mesh: list[int]
Mesh in x, y, z passed on to mesh_phonons.
comm: MPI communicator
"""
self.atoms = atoms
self.fc_potentials = {}
self.fc_fit = None
self.cutoffs = cutoffs if cutoffs is not None else [5]
self.fit_order = fit_order
self._exact_phonons = None
self._mesh_phonons = None
self.phonon_mesh = [30, 30, 30]
self.comm = comm
def __repr__(self) -> str:
return f"EffQH({self.__dict__!r})"
def __str__(self) -> str:
return "EffQH(ForceConstantCalculator(volume))"
@property
def calc(self) -> EffQHCalculator:
return EffQHCalculator(self.atoms, self.fc_fit)
[docs]
def fit_snapshots(
self,
snapshots: SnapshotContainer | list[list[Atoms]],
method: str | None = None,
) -> None:
"""Fits the EffQH to snapshots over volumes.
Internally fits a `hiphive ForceConstantPotential` to the snapshots
for each volume, and then fits each resulting force-constant component
with a polynomial over volume.
Parameters
----------
snapshots : SnapshotContainer or list of list of Atoms
Either a SnapshotContainer with volumes in the index or a list
grouped by volume with lists of Atoms.
method : str (optional)
Method for fitting the effective force constants. Defaults to
'least-squares' in general and lasso for underdetermined systems.
Notes
-----
For large systems, the 'least-squares' method is preferred for speed
reasons.
"""
parcall(logger.debug, "Fitting FCPs and EffQH to snapshots ...")
if isinstance(snapshots, SnapshotContainer):
volumes = snapshots.to_pandas().index.to_frame()["volume"].unique()
lst = [snapshots.filter(f"volume == {_}").get("atoms") for _ in volumes]
else: # should be list
lst = snapshots
volumes = [_[0].get_volume() / len(_[0]) for _ in lst]
fc_potentials = []
scattered = self.comm.scatter(
chunks(zip(volumes, lst, strict=False), self.comm.size)
)
for volume, images in scattered:
ref_atoms = get_ref_atoms(self.atoms, volume)
scaled_cutoffs = np.array(self.cutoffs) * (volume / volumes[0]) ** (1 / 3)
fcp = _fit_hiphive_force_constant_potential(
ref_atoms,
images,
scaled_cutoffs,
method,
)
fc_potentials.append(fcp)
fc_potentials = gather_chunks(self.comm.allgather(fc_potentials))
for volume, fcp in zip(volumes, fc_potentials, strict=False):
self.fc_potentials[volume] = fcp
self.fc_fit = ForceConstantFit.fit_from_hiphive(
volumes,
self.fc_potentials.values(),
self.atoms,
fit_order=self.fit_order,
)
[docs]
def change_size(self, new_atoms: Atoms) -> EffQH:
"""Return a new EffQH instance for ``new_atoms``.
Parameters
----------
new_atoms : Atoms
Must have the same symmetry as ``effqh.atoms`` and be large enough
for the cutoff to fit inside, but can differ in number of atoms and
cell shape.
Returns
-------
EffQH
"""
warnings.warn(
"change_size is deprecated and will be removed in a future version. "
"Use change_reference instead.",
DeprecationWarning,
stacklevel=2,
)
return self.change_reference(new_atoms)
[docs]
def change_reference(self, new_atoms: Atoms) -> EffQH:
"""Return a new EffQH instance for ``new_atoms``.
Parameters
----------
new_atoms : Atoms
Must have the same symmetry as ``effqh.atoms`` and be large enough
for the cutoff to fit inside, but can differ in number or order of
atoms, and cell shape.
Returns
-------
EffQH
"""
if len(self.fc_potentials) == 0:
self.extract_fc_potentials()
new = type(self)(new_atoms)
new.fc_fit = ForceConstantFit.fit_from_hiphive(
self.fc_potentials.keys(),
self.fc_potentials.values(),
new_atoms,
fit_order=self.fit_order,
)
return new
[docs]
def write_hessian(
self,
volume: float,
file: os.PathLike | str,
filetype: str | None = None,
fmt: str = "sphinx",
unit: str = "eVperAng2",
) -> None:
"""Write the force-constant matrix for a given volume to file.
Parameters
----------
volume : float
Volume for which to write the force-constant matrix.
file : os.PathLike or str
Path to the output file.
filetype : str
Filetype to write. Must be 'txt', 'hdf5' or None.
If None, the filetype is guessed from the file suffix.
fmt : str
Format of the output file. Must be 'sphinx' or 'phonopy'.
Default is 'sphinx'.
unit : str
Unit of the force constants, only relevant for ``fmt="sphinx"``.
Must be 'eVperAng2' or 'HaperBohr2'. Default is 'eVperAng2'.
"""
force_constants = self.fc_fit(volume)
write_force_constants(
pathlib.Path(file),
force_constants,
filetype=filetype,
fmt=fmt,
unit=unit,
)
[docs]
def write_fc_fit(self, filename: str) -> None:
self.fc_fit.write(filename)
[docs]
@classmethod
def read_fc_fit(cls, filename: str, atoms: Atoms | list[Atoms] | None = None):
fc_fit = ForceConstantFit.read(filename)
if fc_fit.atoms is not None:
new = cls(fc_fit.atoms)
elif atoms is not None:
new = cls(atoms)
else:
msg = f"{filename} does not contain an atoms object. Pass atoms explicitly."
raise ValueError(msg)
new.fc_fit = fc_fit
return new
[docs]
@classmethod
def read_fcp(
cls,
pathname: str,
atoms: Atoms | list[Atoms],
fit_order: int | None = None,
):
new = cls(atoms)
new.fc_potentials = cls._load_fcp(pathname)
volumes = new.fc_potentials.keys()
fcps = new.fc_potentials.values()
new.fc_fit = ForceConstantFit.fit_from_hiphive(
volumes,
fcps,
atoms,
fit_order=fit_order,
)
return new
@staticmethod
@parallel_function
def _load_fcp(pathname: str) -> dict:
path = pathlib.Path(pathname).expanduser()
if path.is_dir():
files = sorted(path.glob("*"))
else:
files = sorted(path.parent.glob(path.name))
fc_potentials = {}
for file in files:
try:
fcp = ForceConstantPotential.read(str(file))
except ReadError:
continue
atoms: Atoms = fcp.primitive_structure
fc_potentials[atoms.get_volume() / len(atoms)] = fcp
if not fc_potentials:
msg = f"No fcps found in {pathname}"
raise RuntimeError(msg)
return fc_potentials
@property
def exact_phonons(self):
if not self._exact_phonons:
atoms = self.atoms if hasattr(self.atoms, "get_masses") else self.atoms[0]
self._exact_phonons = ExactPhonons.from_fcfit(atoms, self.fc_fit)
return self._exact_phonons
@property
def mesh_phonons(self):
if not self._mesh_phonons:
atoms = self.atoms if hasattr(self.atoms, "get_masses") else self.atoms[0]
self._mesh_phonons = MeshPhonons.from_fcfit(
atoms,
self.fc_fit,
mesh=self.phonon_mesh,
)
return self._mesh_phonons
@property
def exact_free_energy_surface(self):
ph = self.exact_phonons
fec = FreeEnergyContrib(
"qh",
bounds=[ph.volume_bounds, [1, 1500]],
)
fec.parametrization = ph
return FreeEnergySurface([fec])
@property
def mesh_free_energy_surface(self):
ph = self.mesh_phonons
fec = FreeEnergyContrib(
"qh",
bounds=[ph.volume_bounds, [1, 1500]],
)
fec.parametrization = ph
return FreeEnergySurface([fec])
def _fit_hiphive_force_constant_potential(
ref_atoms: Atoms,
images: list[Atoms],
cutoffs: list[float],
method: str | None = None,
) -> ForceConstantPotential:
"""Fit a hiphive ForceConstantPotential.
Parameters
----------
ref_atoms : Atoms
Atoms object representing reference positions.
images : list of Atoms
List of Atoms objects containing the forces and positions for the fit.
cutoffs : list of float
Cutoffs for the Hiphive fitting. Contains only one element for 2nd order
force constants.
method : str (optional)
Method for fitting the effective force constants. Defaults to lstsq in
general and lasso for underdetermined systems.
Returns
-------
ForceConstantPotential
Notes
-----
For large systems, the lstsq method is preferred for speed.
"""
parcall(logger.debug, "Preparing input ...")
hph_atoms = _prepare_atoms(images[0], ref_atoms)
cluster_space = ClusterSpace(hph_atoms, list(cutoffs))
structure_container = StructureContainer(cluster_space)
for atoms in images:
hph_atoms = _prepare_atoms(atoms, ref_atoms)
structure_container.add_structure(hph_atoms)
train_size = len(atoms) * 3 * len(images)
underdetermined = train_size < cluster_space.n_dofs
warning_msg = (
"\n"
"Fitting ForceConstantPotential is underdetermined. If `method` is not "
"given, it will default to 'lasso', which might be slow for large number "
"of parameters. The method can be choosen through the `method` arg.\n"
"See https://trainstation.materialsmodeling.org/moduleref.html#trainstation.Optimizer."
"\n"
)
if method is None:
if underdetermined:
method = "lasso"
warnings.warn(warning_msg)
else:
method = "least-squares"
opt = Optimizer(
structure_container.get_fit_data(),
train_size=1, # fraction used for training
fit_method=method,
seed=np.random.randint(1000, 1000000),
)
logger.debug(
f"Fitting ForceConstantPotential with {cluster_space.n_dofs}"
f" degrees of freedom and {train_size} data points.",
)
with warnings.catch_warnings():
warnings.filterwarnings(
"ignore",
message=(
"Objective did not converge. You might want to "
"increase the number of iterations."
),
)
opt.train()
parcall(logger.info, opt)
fcp = ForceConstantPotential(cluster_space, opt.parameters)
fcp.metadata["rmse_train"] = opt.rmse_train
return fcp
def _prepare_atoms(atoms: Atoms, atoms_ideal: Atoms) -> Atoms:
try:
forces = atoms.calc.get_property("forces", allow_calculation=False)
except:
if hasattr(atoms, "forces"):
forces = atoms.forces
elif "forces" in atoms.arrays:
forces = atoms.get_array("forces")
else:
forces = atoms.get_forces()
displacements = get_displacements(atoms, atoms_ideal)
new_atoms = atoms.copy()
new_atoms.positions = atoms_ideal.get_positions()
new_atoms.new_array("displacements", displacements)
new_atoms.new_array("forces", forces)
return new_atoms