r"""
Computes Kirkwood-Buff integrals (KBIs) from radial distribution functions (RDFs).
The KBI calculation typically suffers from finite-size effects because the RDF
converges slowly toward the bulk value and is constrained by the simulation box.
This class provides several correction schemes to recover the thermodynamic limit.
Default Behavior (``weight_type='geometric'``)
----------------------------------------------
By default, the rigorous framework outlined by `Simon (2022) <https://doi.org/10.1063/5.0106162>`_ is employed to address finite-size effects:
* **RDF Correction**: The RDF is adjusted for molecule excess/depletion (`Ganguly and van der Vegt (2013) <https://doi.org/10.1021/ct301017q>`_), ensuring the integral correctly reflects bulk density behavior.
* **Finite-Volume Correction**: Corrects finite-volume effects by applying an analytical hypersphere model, compensating for volume loss at large radial distances (`Krüger et al. (2013) <https://doi.org/10.1021/jz301992u>`_).
* **Extrapolation**: The integral is extrapolated to the thermodynamic limit, accounting for box-size dependence (`Dawass, Krüger, et al. (2020) <https://doi.org/10.3390/nano10040771>`_).
Correction Methods (``weight_type`` parameter)
----------------------------------------------
The ``weight_type`` parameter selects the specific weighting function :math:`w(r)` or :math:`u_k(r)` applied during integration.
The ``'geometric'`` option triggers the default behavior described above.
* ``'none'``: Uses the uncorrected RDF. Suitable for validation or systems where finite-size effects are negligible.
* ``'u0'``: Simple truncation of the integral at :math:`r=L`. Generally converges poorly as it lacks finite-volume corrections.
* ``'u1'``: First-order Taylor expansion in 1/L of the exact hypersphere weight function, :math:`w(r)`. Provides a more accurate estimate of :math:`G^\infty` than the uncorrected `u0` approach.
* ``'u2'``: Uses the exact hypersphere weight function with an analytical approximation of the surface term. This provides the most robust estimate of :math:`G^\infty` among the non-extrapolated methods.
* ``'geometric'``: The most rigorous method. Uses the geometric weight function :math:`w(r)` and performs explicit linear extrapolation to the thermodynamic limit.
.. note::
Linear extrapolation to the thermodynamic limit is performed **only** for ``weight_type='geometric'``.
The other weight functions (``u0``-``u2``) serve as direct estimates of :math:`G^\infty` as demonstrated by `Krüger and Vlugt (2018) <https://doi.org/10.1103/PhysRevE.97.051301>`_.
"""
from functools import cached_property
from pathlib import Path
from typing import TYPE_CHECKING, Literal
import matplotlib.pyplot as plt
import numpy as np
import scipy
from kbkit.config.mplstyle import load_mplstyle
from kbkit.io.rdf import RdfParser
from kbkit.utils.exceptions import KBIConvergenceError, LinearityError, handle_error
from kbkit.visualization.format import style_axes, style_legend
if TYPE_CHECKING:
from kbkit.systems.properties import SystemProperties
# load plotting style
load_mplstyle()
[docs]
class KBIntegrator:
"""
Class to compute the Kirkwood-Buff Integrals (KBI) from RDF data.
Parameters
----------
r: np.ndarray
Radial distance array (nm).
gr: np.ndarray
Radial distribution function values.
box_volume: float
Averaged simulation box volume (nm³).
delta: int
Kronecker delta for RDF molecules.
weight_type: str, optional
Type of weight function for finite-volume corrections. Options: ('none','u0','u1','u2','geometric')
errors : Literal["raise", "warn", "ignore"], optional
Error mode for handling KBIConvergenceErrors. Only for ``weight_type='geometric'``.
force: bool, optional
If KBIConvergenceError is raised, warns user and returns KBI for ``weight_type='u2'``. Only for ``weight_type='geometric'``.
"""
def __init__(
self,
r: np.ndarray,
gr: np.ndarray,
n_ref: int,
box_volume: float,
delta: int,
weight_type: Literal["none", "u0", "u1", "u2", "geometric"] = "geometric",
errors: Literal["raise", "warn", "ignore"] = "raise",
force: bool = False,
) -> None:
self.r = np.asarray(r)
self.gr = np.asarray(gr)
self.n_ref = int(n_ref)
self.box_volume = float(box_volume)
self.delta = int(delta)
self.weight_type = weight_type.lower()
self.rho_ref = self.n_ref / self.box_volume
self.errors = errors
self.force = force
[docs]
@classmethod
def from_rdf(
cls,
rdf: str | Path | RdfParser,
system_properties: "SystemProperties",
weight_type: Literal["none", "u0", "u1", "u2", "geometric"] = "geometric",
errors: Literal["raise", "warn", "ignore"] = "raise",
force: bool = False,
) -> "KBIntegrator":
"""
Create KBI object from RdfParser and SystemProperties.
Automatically extracts volume and molecule_count from system_properties.
Parameters
----------
rdf: str | Path | RdfParser
RdfParser object or file containing rdf data. Supported filetypes: ('.xvg','.txt','.csv','.xlsx')
system_properties: SystemProperties
System properties containing volume and topology.
weight_type: str, optional
Type of weight function for finite-volume corrections. Options: ('none','u0','u1','u2','geometric')
errors : Literal["raise", "warn", "ignore"], optional
Error mode for handling KBIConvergenceErrors. Only for ``weight_type='geometric'``.
force: bool, optional
If KBIConvergenceError is raised, warns user and returns KBI for ``weight_type='u2'``. Only for ``weight_type='geometric'``.
Returns
-------
KBI
Initialized integrator.
"""
if isinstance(rdf, (str, Path)):
rdf = RdfParser(rdf)
if not isinstance(rdf, RdfParser):
raise TypeError(
"'rdf' is not of type 'RdfParser' and could not create RdfParser object from 'rdf'. Check that 'rdf' is a valid RdfParser or filepath"
)
molecule_count = system_properties.get("molecule_count")
rdf_mols = rdf.extract_molecules(text=rdf.filepath.name, mol_list=molecule_count.keys())
return cls(
r=rdf.r,
gr=rdf.gr,
n_ref=molecule_count[rdf_mols[1]],
box_volume=system_properties.get("volume", units="nm^3"),
delta=int(rdf_mols[0] == rdf_mols[1]),
weight_type=weight_type,
errors=errors,
force=force,
)
@cached_property
def gr_vdv(self) -> np.ndarray:
r"""
Compute the corrected radial distribution function, accounting for excess/depletion of molecule :math:`i` around molecule :math:`j` due to a finite number of molecules, based on the approach by `Ganguly and Van der Vegt (2013) <https://doi.org/10.1021/ct301017q>`_.
Returns
-------
np.ndarray
Factor for correcting RDF so density follows the bulk density of the system.
Notes
-----
The correction is calculated as,
.. math::
g^{vdV}(r) = g(r) \frac{N_j f(r)}{N_j f(r) - \Delta N_j(r) - \delta_{ij}}
.. math::
f(r) = 1 - \frac{\frac{4}{3} \pi r^3}{\langle V \rangle}
.. math::
\rho_j = \frac{N_j}{\langle V \rangle}
.. math::
\Delta N_j(r) = \rho_j \int_0^{r_{max}} 4 \pi r^2 \bigl(g(r) - 1 \bigr) dr
where:
- :math:`r` is the distance
- :math:`\langle V \rangle` is the box volume
- :math:`N_j` is the number of particles of type \( j \)
- :math:`g(r)` is the radial distribution function directly from simulation
- :math:`\delta_{ij}` is a kronecker delta
.. note::
The cumulative integral :math:`\Delta N_j(r)` is approximated numerically using the trapezoidal rule.
"""
sphere_volume = (4 / 3) * np.pi * self.r**3
Vr = 1 - (sphere_volume / self.box_volume)
delta_n_ref = scipy.integrate.cumulative_trapezoid(
4 * np.pi * self.rho_ref * self.r**2 * (self.gr - 1), x=self.r, initial=0
)
vdv_correction = self.n_ref * Vr / (self.n_ref * Vr - delta_n_ref - self.delta)
return self.gr * vdv_correction
[docs]
def compute_hr(self, weight_type: str) -> np.ndarray:
r"""Total correlation function, :math:`h(r) = g(r) - 1`.
If ``weight_type='none'``,
.. math::
h(r) = g^{NpT}(r) - 1,
where :math:`g^{NpT}(r)` is the uncorrected radial distribution function.
For any other ``weight_type`` specified, the `Ganguly and Van der Vegt (2013) <https://doi.org/10.1021/ct301017q>`_ correction is applied,
.. math::
h(r) = g^{vdV}(r) - 1.
"""
return self.gr - 1 if weight_type == "none" else self.gr_vdv - 1
@cached_property
def hr(self) -> np.ndarray:
r"""np.ndarray: Total correlation function, :math:`h(r) = g(r) - 1` based on the ``weight_type`` intitalized in class."""
return self.compute_hr(self.weight_type)
[docs]
def geometric_weight(self, r: np.ndarray) -> np.ndarray:
r"""
Geometric weighting function for 3D hypersphere.
Correct KBI for finite volumes with an exact analytically correct form for hyperspheres, based on the method described by `Krüger et al. (2013) <https://doi.org/10.1021/jz301992u>`_.
Parameters
----------
r: np.ndarray
Radial distance for weight function.
Returns
-------
np.ndarray
Geometric weight function for hyperspheres.
Notes
-----
The geometric weight function, :math:`w(r)`, is defined as:
.. math::
w(r) = 4 \pi r^2 \left( 1 - \frac{3}{2}x + \frac{1}{2}x^3 \right),
where :math:`x` is the dimensionless distance, :math:`x = r/L`.
"""
return 4 * np.pi * r**2 * (1 - (3 / 2) * (r / r.max()) + (1 / 2) * (r / r.max()) ** 3)
[docs]
def u0_weight(self, r: np.ndarray) -> np.ndarray:
r"""
Simple truncation of :math:`r = L`.
This weighting function is well-known to converge badly, it lacks any finite-volume corrections [`Krüger and Vlugt (2018) <https://doi.org/10.1103/PhysRevE.97.051301>`_].
Parameters
----------
r: np.ndarray
Radial distance for weight function.
Returns
-------
np.ndarray
Weight function.
Notes
-----
The weight function, :math:`u_0(r)`, is defined as:
.. math::
u_0(r) = 4 \pi r^2
where:
- :math:`r` is the radial distance
"""
return 4 * np.pi * r**2
[docs]
def u1_weight(self, r: np.ndarray) -> np.ndarray:
r"""
First order Taylor series expansion in :math:`1/L` of the exact finite-volume integral of the sphere [`Krüger and Vlugt (2018) <https://doi.org/10.1103/PhysRevE.97.051301>`_].
Parameters
----------
r: np.ndarray
Radial distance for weight function.
Returns
-------
np.ndarray
Weight function.
Notes
-----
The weight function, :math:`u_1(r)`, is defined as:
.. math::
u_1(r) = 4 \pi r^2 \left( 1 - x^3 \right).
"""
return 4 * np.pi * r**2 * (1 - (r / r.max()) ** 3)
[docs]
def u2_weight(self, r: np.ndarray) -> np.ndarray:
r"""
Approximation of KBI in the thermodynamic limit using the exact geometric weighting function for hyperspheres and approximation of the surface term [`Krüger and Vlugt (2018) <https://doi.org/10.1103/PhysRevE.97.051301>`_].
Parameters
----------
r: np.ndarray
Radial distance for weight function.
Returns
-------
np.ndarray
Weight function.
Notes
-----
The weight function, :math:`u_2(r)`, is defined as:
.. math::
\begin{aligned}
u_2(r) &= w(r) \left( 1 + \frac{3}{2}x + \frac{9}{4}x^2 \right), \\
&= 4 \pi r^2 \left( 1 - \frac{23}{8}x^3 + \frac{3}{4}x^4 + \frac{9}{8}x^5 \right).
\end{aligned}
"""
return self.geometric_weight(r=r) * (1 + (3 / 2) * (r / r.max()) + (9 / 4) * (r / r.max()) ** 2)
@property
def _weight_mapped(self) -> dict:
return {
"none": self.u0_weight,
"u0": self.u0_weight,
"u1": self.u1_weight,
"u2": self.u2_weight,
"geometric": self.geometric_weight,
}
[docs]
def compute_running_kbi(self, weight_type: str):
r"""
Compute KBI as a function of radial distance between molecules, also referred to as the `running KBI`.
We follow the different correction schemes dicussed in `Krüger and Vlugt (2018) <https://doi.org/10.1103/PhysRevE.97.051301>`_, where various ``weight_type`` implement different corrections.
Parameters
----------
weight_type: str
Type of weight function for finite-volume corrections. Options: ('none', 'u0', 'u1', 'u2', 'geometric')
Returns
-------
np.ndarray
KBI values as a numpy array corresponding to distances :math:`r` from the RDF.
Notes
-----
For ``weight_type='geometric'``, the KBI is computed using the formula:
.. math::
\begin{aligned}
G^V(L) &= \int_0^L h(r) w(r) dr, \\
&= G^\infty + \frac{1}{L}F^\infty + \mathcal{O}(L^{-2}),
\end{aligned}
where:
- :math:`L \equiv 6V/A` is the linear dimension of the system.
- :math:`F^\infty` being the surface term, related to surface effects of the subvolume
- :math:`h(r)` is the total correlation function.
While the above equation provides an exact relation to :math:`G^\infty` it requires a linear region for extrapolation to the thermodynamic limit, which can be a disadvantage if a linear region is not identified.
Krüger and Vlugt proposed a method for direct estimation of :math:`G^\infty`, where the accuracy is depending upon weight functions of the form, :math:`u_k(r)`:
.. math::
G^\infty \approx G_k(L) = \int_0^L h(r) u_k(r) dr.
"""
weight_fn = self._weight_mapped[weight_type]
hr = self.compute_hr(weight_type)
rkbi = np.zeros(self.r.size)
for i in range(1, rkbi.size):
rkbi[i] = scipy.integrate.trapezoid(weight_fn(r=self.r[: i + 1]) * (hr[: i + 1]), x=self.r[: i + 1])
return rkbi
@cached_property
def running_kbi_map(self) -> dict[str, np.ndarray]:
"""dict[str, np.ndarray]: Compute the running KBI for each ``weight_type``."""
return {
"none": self.compute_running_kbi(weight_type="none"),
"u0": self.compute_running_kbi(weight_type="u0"),
"u1": self.compute_running_kbi(weight_type="u1"),
"u2": self.compute_running_kbi(weight_type="u2"),
"geometric": self.compute_running_kbi(weight_type="geometric"),
}
@property
def rkbi(self) -> np.ndarray:
"""np.ndarray: Running KBI for class initialized ``weight_type``."""
return self.running_kbi_map[self.weight_type]
@property
def geometric_extrapolation_result(self) -> dict:
"""dict: By default, returns the result of the linear extrapolation using ``maximize_r2=True``. This can be updated by running ``compute_geometric_extrapolation`` with different arguments and setting ``store=True``."""
if not hasattr(self, "_result"):
self._result = self.compute_geometric_extrapolation(maximize_r2=True, errors="raise")
return self._result
[docs]
def compute_kbi(self, weight_type: str) -> float:
r"""Get KBI value in the thermodynamic limit according to ``weight_type``.
Parameters
----------
weight_type: str
Type of weight function for finite-volume corrections. Options: ('none','u0','u1','u2','geometric')
Returns
-------
float
KBI in the thermodynamic limit
Notes
-----
For ``weight_type='geometric'``; KBI is compute via extrapolation to the thermodynamic limit. (see also: :meth:`compute_geometric_extrapolation`)
.. math::
L G^V(L) = L G^\infty + F^\infty
For all other type of ``weight_type``, KBI is approximated by averaging the value of :math:`G_k(L)` over the last 0.1 nm.
.. math::
G^\infty \approx G_k(L)
"""
# get mask for rkbi tail average
mask = (self.r > (self.r.max() - 0.1)) & (self.r < self.r.max())
if weight_type.lower() == "geometric":
try:
return self.geometric_extrapolation_result["G"]
except KBIConvergenceError as e:
if self.force:
handle_error(
error_mode="warn" if self.errors != "ignore" else "ignore",
message=f"KBI convergence failed: {e} Falling back to weight_type='u2'.",
error_type=type(e),
)
return float(np.nanmean(self.running_kbi_map["u2"][mask]))
handle_error(self.errors, str(e), type(e))
return np.nan
# for all other types; just extract last values in running kbi
return float(np.nanmean(self.running_kbi_map[weight_type.lower()][mask]))
@property
def kbi(self) -> float:
"""float: KBI value in the thermodynamic limit for class initialized ``weight_type``."""
return self.compute_kbi(weight_type=self.weight_type)
[docs]
def add_kbi(self, axhandle: plt.axes, weight_type: str, **kwargs) -> None:
"""Add KBI as a function of radial distance to axes.
Parameters
----------
axhandle: plt.axes
Axes to add plot.
weight_type: str
Type of weight function for finite-volume corrections. Options: ('none','u0','u1','u2','geometric')
"""
axhandle.plot(self.r, self.running_kbi_map[weight_type], **kwargs)
[docs]
def add_lkbi(self, axhandle: plt.axes, **kwargs) -> None:
"""Add scaled KBI as a function of radial distance to axes, for plotting geometric linear extrapolation.
Parameters
----------
axhandle: plt.axes
Axes to add plot.
"""
axhandle.plot(self.r, self.r * self.running_kbi_map["geometric"], **kwargs)
[docs]
def add_lkbi_fit(self, axhandle: plt.axes, result: dict | None = None, **kwargs) -> None:
"""Add fit results from geometric linear extrapolation to axes.
Parameters
----------
axhandle: plt.axes
Axes to add plot.
result: dict, optional
Geometric extrapolation result.
"""
result = result or self.geometric_extrapolation_result
rfit = result["r_fit"]
r_rkbi_pred = result["r_rkbi_pred"]
axhandle.plot(rfit, r_rkbi_pred, **kwargs)
[docs]
def add_kbi_value(self, axhandle: plt.axes, weight_type: str, **kwargs) -> None:
"""Add horizontal line for KBI in the thermodynamic limit for a given ``weight_type``.
Parameters
----------
axhandle: plt.axes
Axes to add plot.
weight_type: str
Type of weight function for finite-volume corrections. Options: ('none','u0','u1','u2','geometric')
"""
kbi = self.compute_kbi(weight_type)
if all(x not in kwargs for x in ("ls", "linestyle")):
kwargs.update({"ls": (0, (10, 5, 2, 5))})
axhandle.axhline(kbi, **kwargs)
@property
def _colors_by_weight(self) -> dict:
return {
"none": "gold",
"u0": "cyan",
"u1": "lawngreen",
"u2": "red",
"geometric": plt.colormaps["hsv"](np.linspace(0.12, 1, 5))[3],
}
def _get_colors(self, cmap: str | None = None) -> dict:
"""dict: Get colors mapped to ``weight_type``."""
if cmap is None:
return self._colors_by_weight
else:
colors_from_cmap = plt.colormaps[cmap](np.linspace(0, 1, 5))
return {weight: colors_from_cmap[i] for i, weight in enumerate(self._colors_by_weight)}
[docs]
def plot_kbi(
self,
weight_type: str | None = None,
cmap: str | None = None,
add_kbi_value: bool = False,
filepath: str | None = None,
**kwargs,
) -> None:
"""Create figure for plotting single KBI for a given weight type.
Parameters
----------
weight_type: str, optional
Weight function to plot. Options: ('none','u0','u1','u2','geometric'). Default will go to class initialized ``weight_type``.
cmap: str, optional
Matplotlib colormap name.
add_kbi_value: bool, optional
Add horizontal value for KBI with ``weight_type='geometric'``.
filepath: str, optional
Filepath to save figure as.
"""
weight_type = weight_type or self.weight_type
colors_by_weight = self._get_colors(cmap)
fig, ax = plt.subplots()
style_axes(ax, minorxticks=np.arange(0, self.r.max(), 0.5))
label = r"G"
if weight_type.startswith("u"):
label += rf"$_{weight_type[1]}^\infty$"
elif weight_type == "geometric":
label += r"$^\text{V}$"
self.add_kbi(ax, weight_type=weight_type, color=colors_by_weight[weight_type], label=label, **kwargs)
if add_kbi_value:
try:
self.add_kbi_value(ax, weight_type="geometric", color="k", lw=0.5, ls="-", label=r"G$^\infty$")
except KBIConvergenceError as e:
handle_error(error_mode="warn", message=str(e), error_type=type(e))
# Only add legend if labels exist
_, labels = plt.gca().get_legend_handles_labels()
if labels:
style_legend(ax)
ax.set_ylabel(rf"{label} (nm$^3$)", fontsize=16)
ax.set_xlabel(r"L (nm)", fontsize=16)
if filepath:
fig.savefig(filepath, dpi=100)
plt.close()
else:
plt.show()
[docs]
def plot_kbi_compare(
self,
weight_types: list[str],
cmap: str | None = None,
add_kbi_value: bool = False,
filepath: str | None = None,
) -> None:
"""Create figure for comparing KBI's from various correction methods.
Parameters
----------
weight_types: list[str]
Weight functions to include in plot. Options: ('none','u0','u1','u2','geometric')
cmap: str, optional
Matplotlib colormap name.
add_kbi_value: bool, optional
Add horizontal value for KBI with ``weight_type='geometric'``.
filepath: str, optional
Filepath to save figure as.
"""
colors_by_weight = self._get_colors(cmap)
if isinstance(weight_types, str):
weight_types = [weight_types]
fig, ax = plt.subplots()
style_axes(ax, minorxticks=np.arange(0, self.r.max(), 0.5))
for weight_type in weight_types:
label = r"G"
if weight_type.startswith("u"):
label += rf"$_{weight_type[1]}^\infty$"
elif weight_type == "geometric":
label += r"$^\text{V}$"
self.add_kbi(ax, weight_type=weight_type, lw=1, color=colors_by_weight[weight_type], label=label)
if add_kbi_value:
try:
self.add_kbi_value(ax, weight_type="geometric", color="k", lw=0.5, ls="-", label=r"G$^\infty$")
except KBIConvergenceError as e:
handle_error(error_mode="warn", message=str(e), error_type=type(e))
# Only add legend if labels exist
_, labels = plt.gca().get_legend_handles_labels()
if labels:
style_legend(ax)
ax.set_ylabel(r"G(L) (nm$^3$)", fontsize=16)
ax.set_xlabel(r"L (nm)", fontsize=16)
if filepath:
fig.savefig(filepath, dpi=100)
plt.close()
else:
plt.show()