from __future__ import annotations
import logging
from typing import TYPE_CHECKING, Any, NamedTuple
import numpy as np
import pandas as pd
import scipy.stats
from ase.calculators.mixing import MixedCalculator
from directupsampling.grid import Grid
from directupsampling.parallel import parlog
from directupsampling.sampling import GridSampler, SamplingSetting
from directupsampling.snapshots import SnapshotContainer
from directupsampling.tifit import ThermoIntFit
if TYPE_CHECKING:
from collections.abc import Iterable
from pathlib import Path
from ase.atoms import Atoms
from ase.calculators.calculator import Calculator
logger = logging.getLogger(__name__)
[docs]
class ThermoInt:
"""Class for thermodynamic integration from calc0 to calc1 on vt_grid.
Integrates over provided lambdas (from 0 to 1, corresponding to the
repsective calculators).
"""
def __init__(
self,
atoms: Atoms,
calc0: Calculator,
calc1: Calculator,
vt_grid: Grid | Iterable[tuple[float]],
lambdas: Grid | Iterable[float] = (0.0, 0.25, 0.5, 0.75, 1.0),
eos0: Any | None = None,
eos1: Any | None = None,
sampling_setting: SamplingSetting | dict[str, Any] | None = None,
) -> None:
"""Create ThermoInt.
Results are in eV/atom. If any of eos0 or eos1 are given, it assumes
vibrational energy and divides the results by (number of atoms) - 1,
consistent with Eq. 5 in the Supplementary Information to `npj2022`_.
Parameters
----------
atoms: Atoms
Initial structure for starting the TI/MD.
calc0, calc1: ASE calculator
Calculators to integrate between (from calc0 [lambda=0] to calc1
[lambda=1]).
vt_grid: Grid | Iterable[tuple[float, ...]]
Representing a (volume, temperature) grid
lambdas: Grid | Iterable[float]
Lambda "grid" over which to perform the TI. Should go from 0 to
1, and at least 5 points is recommended.
eos0, eos1: callable, optional
References for calc0, calc1, respectively, giving energy as functions
of volume. The energy will be subtracted from that of calc1 according
to Eq. 5 in the Supplementary Information to `npj2022`_.
sampler_setting: dict, optional
Keyword arguments to snapshot sampler. Typically used to adjust
`sampling_interval` (default 100 steps) and/or `sampling_offset`
(default 1000 steps).
For the arguments `thermostat_kwargs` and `thermostat_class`, see the
`Dynamics` class.
.. _npj2022: https://www.nature.com/articles/s41524-022-00956-8
Attributes
----------
snapshot_container: `SnapshotContainer`
Holds the snapshots sampled.
fits: dict[str, ThermoIntFit]
Fits of energy vs lambda for each grid point.
free_energy_differences: pd.DataFrame
Containg the free energy differences resulting from integration over
lambda at each grid point.
Notes
-----
`vt_grid` and `lambdas` elements can be repeated for parallel sampling
with different random seeds.
"""
self.atoms = atoms
self.calc0 = calc0
self.calc1 = calc1
self.lambda_grid = Grid.from_any(lambdas, "lambda_")
self.vt_grid = Grid.from_any(vt_grid)
self.eos0 = eos0
self.eos1 = eos1
self.snapshot_container = None
self.fits: dict[tuple, ThermoIntFit] = {} # energy-vs-lambda fits
self.sampling_setting = SamplingSetting.from_any(sampling_setting)
def __repr__(self) -> str:
return f"ThermoInt({self.__dict__!r})"
[docs]
def run_sampling(
self,
nsnapshots: int = 150,
sample_atoms: bool = False,
**kwargs: dict,
) -> None:
"""Sample `nsnapshots` (default 150) snapshots over vt_grid*lambdas.
Parameters
----------
nsnapshots: int
The number of snapshots to sample.
sample_atoms: bool
Whether or not to sample atoms objects.
Note that the default (False) differs from that of SnapshotSamplers
sample_snapshots and Dynamics.
kwargs: dict
Keyword arguments for Dynamics.
"""
sampler = GridSampler(
self.atoms,
grid=self.vt_grid * self.lambda_grid,
setting=self.sampling_setting,
snapshot_container=self.snapshot_container,
)
def set_mixed_calc(s=sampler, a=self.atoms) -> None:
lambda_ = s.grid.current.lambda_
a.calc = MixedCalculator(self.calc0, self.calc1, 1.0 - lambda_, lambda_)
sampler.attach(set_mixed_calc)
sampler.sample_snapshots(
nsnapshots,
sample_atoms=sample_atoms,
**kwargs,
)
self.snapshot_container = sampler.snapshot_container # needed if first is None
[docs]
def write_folders(self, **kwargs: dict) -> None:
sampler = GridSampler(
self.atoms,
grid=self.vt_grid * self.lambda_grid,
setting=self.sampling_setting,
)
sampler.write_folders(
kind="sampling",
calc=(self.calc0, self.calc1),
**kwargs,
)
@property
def free_energy_differences(self) -> pd.DataFrame:
feds = {}
for vt_point in self.vt_grid:
if not self.fits.get(vt_point):
self.extract_free_energy_difference(vt_point)
key = tuple(float(getattr(vt_point, ax)) for ax in self.vt_grid.axes)
if self.fits[vt_point] is not None:
best_f = self.fits[vt_point].best_fit
feds[key] = best_f[["free energy", "fit error"]]
else:
feds[key] = pd.Series([np.nan, np.nan])
df = pd.DataFrame(feds).transpose()
df.index.names = list(self.vt_grid.axes)
return df
def _average_energy_difference_per_atom(
self,
snapshot_container: SnapshotContainer,
volume: float,
) -> tuple[float, float]:
"""Return <ediff> in eV/atom."""
energy_pairs = [
(_["potential_energy_0"], _["potential_energy_1"])
for _ in snapshot_container.data
]
if len(energy_pairs) > 0:
energies0, energies1 = [np.array(x) for x in zip(*energy_pairs)]
natoms = len(self.atoms)
if self.eos0 is not None or self.eos1 is not None:
ref_ene0, ref_ene1 = self._get_reference_energies(volume)
delta_e = (
(energies1 - ref_ene1 * natoms) - (energies0 - ref_ene0 * natoms)
) / (natoms - 1)
else:
delta_e = (energies1 - energies0) / natoms
delta_u, err = self.get_avg_and_err(delta_e)
return delta_u, err
else:
return np.nan, np.nan
def _get_reference_energies(self, volume: float) -> tuple[float, float]:
if self.eos0 is not None:
ref_ene0 = self.eos0.get_energy(volume)
else:
ref_ene0 = 0
if self.eos1 is not None:
ref_ene1 = self.eos1.get_energy(volume)
else:
ref_ene1 = 0
return ref_ene0, ref_ene1
def _check_temperature_convergence(self, vt_point):
raise NotImplementedError
[docs]
@classmethod
def read(
cls,
filename: str | Path,
atoms: Atoms = None,
) -> ThermoInt:
"""Return a ThermoInt instance corresponding to .db file.
Parameters
----------
filename : str or Path
Path to file from which to read the snapshots for the thermodynamic
integration. The .db file must contain indices over volume,
temperature and lambda.
atoms : ASE Atoms (optional)
Will be assigend to returned ThermoInt instance. This is not needed
if the .db file contains sampled atoms.
Returns
-------
A new ThermoInt instance.
"""
snapshots = SnapshotContainer.read(filename)
if atoms is None:
atoms = snapshots.get("atoms")[0]
vt_grid = snapshots.index.get_grid(["volume", "temperature"])
lambdas = np.array(snapshots.index.get_grid(["lambda_"])).ravel()
new = cls(atoms, None, None, vt_grid, lambdas=lambdas)
new.snapshot_container = snapshots
return new
[docs]
@staticmethod
def get_avg_and_err(data) -> tuple[float, float]:
nsamples = data.shape[0]
if nsamples == 0:
avg, err = np.nan, np.nan
else:
t_value = scipy.stats.t.interval(0.95, df=nsamples - 1, loc=0, scale=1)[1]
avg = np.mean(data, axis=0)
err = t_value * np.std(data, axis=0, ddof=1) / np.sqrt(nsamples)
return avg, err
[docs]
def check_if_empty(snapshot_container: SnapshotContainer, vt_point: NamedTuple) -> bool:
q = f"volume == {vt_point.volume} & temperature == {vt_point.temperature}"
n = len(snapshot_container.filter(q))
return n == 0