import copy
import pathlib
from typing import Iterable
import numpy as np
import pandas
import scipy.stats
from ase.calculators.calculator import Calculator
from ase.calculators.singlepoint import SinglePointCalculator
from directupsampling.data import kB
from directupsampling.grid import Grid
from directupsampling.sampling import GridSampler
from directupsampling.snapshots import SnapshotContainer
[docs]
class Upsampling:
def __init__(
self,
snapshot_container: SnapshotContainer = None,
calc: str | Calculator = None,
vt_grid: Grid | Iterable[tuple[float]] | None = None,
energy_differences: dict[tuple, list] | None = None,
natoms: int | None = None,
per_atom: bool | None = None,
vibrational: bool = False,
restart: bool = False,
):
"""Class for upsampling with free energy perturbation.
Takes at least either a SnapshotContainer as first argument,
or energy_differences as keyword argument.
Parameters
----------
snapshot_container: SnapshotContainer
Snapshots to upsample. Corresponds to the reference system and
contains atoms with stored energy results. If the argument
`energy_differences` is given directly, the 'snapshot_container' is
optional. Can also contain already upsampled snapshots, see the
"restart" argument.
calc: str or ASE Calculator
Corresponds to the system/interactions which to upsample to with
free energy perturbation. If the argument `energy_differences` is
given directly, the calculator is optional.
vt_grid: Grid or list of tuples (optional)
Tuples should contain (volume, temperature) points. If not given
(None) it is extracted from `snapshot_container` or
`energy_differences`.
energy_differences: dict (optional)
Optional argument to directly provide energy differences for the
upsampling. Should be a dictionary with keys correspoding to vt_grid
and list/np.ndarray with energy-difference values. It is important
that these energies are in the unit eV (electron volt) for the
upsampling to be correct. `per_atom=False` can be specified if the
energies are not per atom, and in that case also the output unit is
not per atom.
natoms: int (optional)
Influences `self.normalizing_factor` and thereby the scaling of the
results, as well as `energy_differences` if given.
per_atom: bool (optional)
Default is `True` if `snapshot_container` is given. If
`energy_differences` is given, the default is `True if `natoms` is
given, else `False`.
vibrational: bool (optional)
Default is `False`. Decides if the normalization (if per_atom is
`True`) is done with natoms-1 (`True`) or natoms (`False`).
restart: bool (optional)
If true, the atoms energies in snapshot_container will be used as
upsampled energies, and it will search for references energies in
the container data.
Attributes
----------
snapshot_container: SnapshotContainer
Contains the snapshots sampled.
normalizing_factor: float
Scales the energy differences used in the free energy perturbation
as well as the results.
energy_differences: dict
Dictionary containing, for each vt_grid point, scaled energy
differences.
free_energy_differences: Pandas DataFrame
Contains the resulting free energy differences from the free energy
perturbation at each grid point.
Examples
--------
>>> upsampling = Upsampling(snapshots, dft, vt_grid=grid)
>>> upsampling.write_folders(
... root_path="my_calculations_folder",
... joblist_file="jobList",
... short=True, # write paths to job list only for each (V,T) point
... )
Then when the calculations are done, extract it by
>>> upsampling = Upsampling.read_folders("my_calculations_folder")
>>> feds_el = upsampling.free_energy_differences
The electronic free energies can be extracted by
>>> electronic_free_energies = { # per snapshot
... (v, t): upsampling.snapshot_container.filter(
... f"volume == {v} & temperature == {t}"
... ).electronic_free_energy
... / upsampling.natoms
... for v, t in upsampling.vt_grid
... }
"""
self._raw_energy_differences = {}
self.calc = calc
if (snapshot_container is None and energy_differences is None) or (
snapshot_container is not None and energy_differences is not None
):
raise TypeError(
"Use one and only one of 'snapshots_container' and 'energy_differences'."
)
elif energy_differences is not None:
if per_atom is True and natoms is None:
TypeError("'natoms' is required when 'per_atom' is True.")
elif per_atom is None:
per_atom = False if natoms is None else True
elif snapshot_container is not None:
if natoms is None:
natoms = len(snapshot_container.get("atoms")[0])
per_atom = True if per_atom is None else per_atom
self.natoms = natoms
if per_atom:
if vibrational:
self.normalizing_factor = 1 / (self.natoms - 1)
else:
self.normalizing_factor = 1 / self.natoms
else:
self.normalizing_factor = 1
if vt_grid is None and snapshot_container is not None:
vt_grid = snapshot_container.index.get_grid(["volume", "temperature"])
elif vt_grid is None and energy_differences is not None:
vt_grid = list(energy_differences.keys())
self.vt_grid = Grid.from_any(vt_grid)
self.snapshot_container = snapshot_container
self.energy_differences = energy_differences
if restart:
self._extract_reference_energies_from_container_data()
self._extract_energy_differences()
@property
def snapshot_container(self):
return self._snapshot_container
@snapshot_container.setter
def snapshot_container(self, snapshot_container):
if snapshot_container is not None:
for _ in snapshot_container.get("atoms"):
if len(_) != self.natoms:
raise ValueError("All snapshots must have the same number of atoms")
self._reference_snapshot_container = snapshot_container
self._snapshot_container = copy.deepcopy(snapshot_container)
else:
self._reference_snapshot_container = None
self._snapshot_container = None
@property
def reference_snapshot_container(self):
return self._reference_snapshot_container
@reference_snapshot_container.setter
def reference_snapshot_container(self):
raise TypeError(
"'reference_snapshot_container' cannot be set. "
"Set 'snapshots_container' instead."
)
@property
def energy_differences(self):
ene_diffs = {}
for vtp in self.vt_grid:
ene_diffs[vtp] = self._raw_energy_differences[vtp] * self.normalizing_factor
return ene_diffs
@energy_differences.setter
def energy_differences(
self, energy_differences: dict[tuple, list[float] | np.ndarray]
):
if isinstance(energy_differences, SnapshotContainer):
raise TypeError("'energy_differences' must be a dict-like.")
if energy_differences is not None:
raw_ene_diffs = {}
for vtp in self.vt_grid:
vtp_tuple = tuple(getattr(vtp, ax) for ax in self.vt_grid.axes)
key = vtp if vtp in energy_differences else vtp_tuple
val = np.array(energy_differences[key]) / self.normalizing_factor
raw_ene_diffs[vtp] = val
self._raw_energy_differences = raw_ene_diffs
else:
self._raw_energy_differences = None
@property
def free_energy_differences(self):
"""Property that gives the results of the free energy perturbation."""
free_ene_diffs = {}
for vt_point in self.vt_grid:
results = calculate_fep(
self._raw_energy_differences[vt_point],
vt_point.temperature,
)
key = tuple(float(getattr(vt_point, ax)) for ax in self.vt_grid.axes)
free_ene_diffs[key] = results * self.normalizing_factor
df = pandas.DataFrame(free_ene_diffs).transpose()
df.index.names = list(self.vt_grid.axes)
return df
[docs]
def calculate_snapshots(self):
"""Calculates `self.snapshot_container` with `self.calc`."""
sampler = GridSampler(
grid=self.vt_grid,
snapshot_container=self.snapshot_container,
)
sampler.calculate_with(self.calc, properties=["energy"], inplace=True)
self._extract_energy_differences()
[docs]
def write_folders(self, **kwargs):
"""Writes folders. See `SnapshotSampler.write_folders`.
Note: Only "energy" is saved from the results of
`self.reference_snapshots_container`.
"""
sampler = GridSampler(
grid=self.vt_grid,
snapshot_container=self.reference_snapshot_container,
)
if "calc" not in kwargs:
kwargs["calc"] = self.calc
sampler.write_folders(**kwargs)
def _extract_reference_energies_from_container_data(self):
for i in self.reference_snapshot_container.index:
atoms = self.reference_snapshot_container[i]["atoms"]
names = []
energies = []
for k, v in self.reference_snapshot_container[i].items():
if "potential_energy_" in k:
names.append(k.replace("potential_energy_", ""))
energies.append(v)
if len(energies) == 0:
raise TypeError("No reference energy found in container.")
elif len(energies) == 1:
ind = 0
elif len(energies) == 2 and atoms.get_potential_energy() in energies:
ind = energies.index(atoms.get_potential_energy()) - 1
else:
raise TypeError(
"More than one possible reference energies found in container."
)
atoms.calc = SinglePointCalculator(atoms, **{"energy": energies[ind]})
atoms.calc.name = names[ind]
[docs]
@classmethod
def read_folders(
cls,
path: str | pathlib.Path,
per_atom: bool = True,
vibrational: bool | None = None,
**kwargs,
):
"""Reads in calculation in a folder structure through `SnapshotSampler`.
Parameters
----------
path: str
The path to the root of the folder structure.
per_atom: bool (optional)
If `False` the results are given per cell, instead of per_atom.
Default is `True`.
vibrational: bool (optional)
Default is `True` if not "old_fermi_folders". Decides if the
normalization (if per_atom is `True`) is done with natoms-1 (`True`)
or natoms (`False`).
old_fermi_folders: bool (optional)
If reading old folders with calc_0K and calc_kBT subfolders, setting
this to `True` will read the Fermi smeared reasults in calc_kBT.
Returns
-------
A new `Upsampling` instance. For how to extract the electronic free
energies, see Examples for the `Upsampling` class.
"""
sampler = GridSampler.read_folders(path, **kwargs)
if len(sampler.snapshot_container) == 0:
raise RuntimeError(f"No finished calculations could be read from {path}")
if vibrational is None:
vibrational = not kwargs.get("old_fermi_folders", False)
new = cls(
sampler.snapshot_container,
calc=sampler.snapshot_container.get("atoms")[0].calc.name,
vt_grid=sampler.grid,
per_atom=per_atom,
vibrational=vibrational,
restart=True,
)
return new
def _extract_energy_differences(self):
"""Extracts energy differences between
`self.reference_snapshot_container` given at instantiation and
`self.snapshot_container` which is altered during `calculate_snapshots`
or `write_folders`/`read_folders`."""
energy_diffs = {}
for vtp in self.vt_grid:
energy_diff_list = []
ref_snapshots = self.reference_snapshot_container.filter(
f"volume == {vtp.volume} & temperature == {vtp.temperature}"
)
snapshots = self.snapshot_container.filter(
f"volume == {vtp.volume} & temperature == {vtp.temperature}"
)
for i in ref_snapshots.index:
ref_atoms = ref_snapshots[i]["atoms"]
atoms = snapshots[i]["atoms"]
energy_diff_list.append(
atoms.get_potential_energy(
force_consistent=(atoms.calc.parameters.get("ismear") == -1)
)
- ref_atoms.get_potential_energy()
)
energy_diffs[vtp] = np.array(energy_diff_list)
self._raw_energy_differences = energy_diffs
[docs]
def calculate_fep(energies, temperature):
"""Returns a pandas Series with results."""
energies = np.array(energies)
temperature = float(temperature)
nsamples = len(energies)
confidence = 0.95
statist_dof = nsamples - 1
t_value = scipy.stats.t.interval(confidence, statist_dof, loc=0, scale=1)[1]
fot_avg = np.mean(energies)
fot_std = np.std(energies, ddof=1)
sot_avg = -1 / (2 * kB * temperature) * np.var(energies)
energies -= fot_avg # Subtract avg to avoid numerical issues
pert_avg, pert_std, pert_bias = perturbation_average(energies, temperature)
pert_err_1 = t_value * pert_std / np.sqrt(nsamples)
pert_err_2 = t_value * (pert_bias[0] / nsamples + pert_bias[1] / nsamples**2)
pert_err = np.sqrt(pert_err_1**2 + pert_err_2**2)
results = {
"free energy": pert_avg + fot_avg, # Add avg again
"error": pert_err,
"1st order term": fot_avg,
"2nd order term": sot_avg,
"1st order stdev": fot_std,
}
return pandas.Series(results)
[docs]
def perturbation_average(energies, t):
exp_energies = np.exp(-energies / (kB * t))
avg_exp_energies = np.mean(exp_energies)
std_exp_energies = np.std(exp_energies, ddof=1)
thd_exp_energies = (
np.mean(exp_energies**3)
- 3 * avg_exp_energies * std_exp_energies**2
- avg_exp_energies**3
)
pert_avg = -kB * t * np.log(avg_exp_energies)
fot_std = np.std(energies, ddof=1)
pert_std = abs(kB * t * std_exp_energies / avg_exp_energies)
phi1 = std_exp_energies**2 / 2 / avg_exp_energies**2
phi2 = (
-(4 * avg_exp_energies * thd_exp_energies - 9 * std_exp_energies**4)
/ 12
/ avg_exp_energies**4
)
pert_bias = kB * t * np.array([phi1, phi2])
# take largest stdev estimate
std = max([pert_std, fot_std])
return pert_avg, std, pert_bias
[docs]
def rescale_to_different_supercell(value, natoms, natoms_new):
"""Returns a rescaled energy (vib) in ene/atom from one natoms to another
The reasoning for this is that the upsampled energy per cell is
approximated to be the same as for the cell with natoms_new atoms (usually
larger). In this case, the energy should not be normalized with natoms-1
but instead with natoms_new-1, to be consistent with the larger supercell
size, so that it can be added together with cancelling references.
Having canceling references resolves e.g. the problem of having to use MTP
0K references, which might not be a nice thing to do, while still keeping
the correct normalization of the vibrational energy from the MD with the
larger supercell size.
"""
if isinstance(value, (list, tuple, str)):
value = np.array(value).astype(float)
value *= (natoms - 1) / (natoms) * natoms_new / (natoms_new - 1)
return value