"""
Bulk structural properties
==========================
.. moduleauthor:: Benjamin Ye <GitHub: @bbye98>
This module contains classes to analyze the structure of bulk fluid and
electrolyte systems.
"""
from itertools import combinations_with_replacement
from typing import Union
import warnings
import MDAnalysis as mda
from MDAnalysis.lib.distances import capped_distance
import numba
import numpy as np
from scipy.integrate import simpson
from scipy.signal import argrelextrema
from scipy.special import jv
from .base import Hash, DynamicAnalysisBase, NumbaAnalysisBase
from .. import FOUND_OPENMM, Q_, ureg
from ..algorithm import accelerated
from ..algorithm.molecule import center_of_mass
from ..algorithm.topology import convert_cell_representation
from ..algorithm.unit import is_unitless, strip_unit
from ..algorithm.utility import get_closest_factors
if FOUND_OPENMM:
from openmm import unit
[docs]
def radial_histogram(
positions_a: np.ndarray[float],
positions_b: np.ndarray[float],
/,
n_bins: int,
range_: np.ndarray[float],
dimensions: np.ndarray[float],
bin_edges: np.ndarray[float] = None,
*,
exclusion: tuple[int] = None,
) -> np.ndarray[float]:
r"""
Computes the radial histogram of distances between particles of
the same species :math:`\alpha` or two different species
:math:`\alpha` and :math:`\beta`.
Parameters
----------
positions_a : `numpy.ndarray`, positional-only
Positions or centers of mass of entities belonging to species
:math:`\alpha`.
**Shape**: :math:`(N_\alpha,\,3)`.
**Reference unit**: :math:`\mathrm{Å}`.
positions_b : `numpy.ndarray`, positional-only
Positions or centers of mass of entities belonging to species
:math:`\beta`.
**Shape**: :math:`(N_\beta,\,3)`.
**Reference unit**: :math:`\mathrm{Å}`.
n_bins : `int`
Number of histogram bins :math:`N_\mathrm{bins}`.
range_ : array-like
Range of radii values.
**Shape**: :math:`(2,)`.
**Reference unit**: :math:`\mathrm{Å}`.
dimensions : array-like
System dimensions and orthogonality.
**Shape**: :math:`(6,)`.
**Reference unit**: :math:`\mathrm{Å}` (dimensions),
:math:`^\circ` (orthogonality).
bin_edges : `numpy.ndarray`, optional
Bin edges.
**Shape**: :math:`(N_\mathrm{bins}+1,)`.
**Reference unit**: :math:`\mathrm{Å}`.
exclusion : array-like, keyword-only, optional
Tiles to exclude from the interparticle distances.
**Shape**: :math:`(2,)`.
**Example**: :code:`(1, 1)` to exclude self-interactions.
Returns
-------
histogram : `numpy.ndarray`
Radial histogram.
**Shape**: :math:`(N_\mathrm{bins},)`.
"""
# Get pair separation distances of atom pairs within range
pairs, distances = capped_distance(
positions_a,
positions_b,
range_[1],
range_[0] - np.finfo(np.float64).eps,
box=dimensions,
)
# Exclude atom pairs with the same atoms or atoms from the
# same residue
if exclusion is not None:
distances = distances[
np.where(pairs[:, 0] // exclusion[0] != pairs[:, 1] // exclusion[1])[0]
]
return accelerated.numba_histogram(
distances,
n_bins,
bin_edges
or accelerated.numba_histogram_bin_edges(
np.asarray(range_, dtype=float), n_bins
),
)
[docs]
def calculate_coordination_numbers(
bins: np.ndarray[float],
rdf: np.ndarray[float],
rho: float,
*,
n_coord_nums: int = 2,
n_dims: int = 3,
threshold: float = 0.1,
) -> np.ndarray[float]:
r"""
Calculates coordination numbers :math:`n_k` from a radial
distribution function :math:`g_{\alpha\beta}(r)`.
It is defined as
.. math::
n_k=4\pi\rho_j\int_{r_{k-1}}^{r_k}r^2g_{\alpha\beta}(r)\,dr
for three-dimensional systems and
.. math::
n_k=2\pi\rho_j\int_{r_{k-1}}^{r_k}rg_{\alpha\beta}(r)\,dr
for two-dimensional systems, where :math:`k` is the index,
:math:`\rho_\beta` is the bulk number density of species
:math:`\beta` and :math:`r_k` is the :math:`(k + 1)`-th local
minimum of :math:`g_{\alpha\beta}(r)`.
If the radial distribution function :math:`g_{\alpha\beta}(r)` does
not contain as many local minima as `n_coord_nums`, this method will
return `numpy.nan` for the coordination numbers that could not be
calculated.
Parameters
----------
bins : `numpy.ndarray`
Centers of the histogram bins.
**Shape**: :math:`(N_\mathrm{bins},)`.
**Reference unit**: :math:`\mathrm{Å}`.
rdf : `numpy.ndarray`
Radial distribution function :math:`g_{\alpha\beta}(r)`.
**Shape**: :math:`(N_\mathrm{bins},)`.
rho : `float`
Number density :math:`\rho_\beta` of species :math:`\beta`.
**Reference unit**: :math:`\mathrm{Å}^{n_\mathrm{dims}}`.
n_coord_nums : `int`, keyword-only, default: :code:`2`
Number of coordination numbers :math:`n_\mathrm{coord}` to
calculate.
n_dims : `int`, keyword-only, default: :code:`3`
Number of dimensions :math:`n_\mathrm{dims}`.
threshold : `float`, keyword-only, default: :code:`0.1`
Minimum :math:`g_{\alpha\beta}(r)` value that must be reached
before local minima are found.
Returns
-------
coord_nums : `numpy.ndarray`
Coordination numbers :math:`n_k`.
"""
if n_dims not in {2, 3}:
raise ValueError("Invalid number of dimensions.")
def f(
r: np.ndarray[float], rdf: np.ndarray[float], rho: float, start: int, stop: int
) -> np.ndarray[float]:
if n_dims == 3:
return 4 * np.pi * rho * simpson(r**2 * rdf[start:stop], r)
else:
return 2 * np.pi * rho * simpson(r * rdf[start:stop], r)
coord_nums = np.empty(n_coord_nums)
coord_nums[:] = np.nan
# Find indices of minima in the radial distribution function
(i_min,) = argrelextrema(rdf, np.less)
i_min = i_min[rdf[i_min] >= threshold]
n_min = len(i_min)
# Integrate the radial distribution function to get the coordination
# number(s)
if n_min:
r = bins[: i_min[0] + 1]
coord_nums[0] = f(r, rdf, rho, None, i_min[0] + 1)
for i in range(min(n_coord_nums, n_min) - 1):
r = bins[i_min[i] : i_min[i + 1] + 1]
coord_nums[i + 1] = f(r, rdf, rho, i_min[i], i_min[i + 1] + 1)
else:
warnings.warn("No local minima found.")
return coord_nums
[docs]
def calculate_structure_factor(
r: np.ndarray[float],
g: np.ndarray[float],
equal: bool,
rho: float,
x_a: float = 1,
x_b: float = None,
q: np.ndarray[float] = None,
*,
q_lower: float = None,
q_upper: float = None,
n_q: int = 1_000,
n_dims: int = 3,
formalism: str = "FZ",
) -> tuple[np.ndarray[float], np.ndarray[float]]:
r"""
Calculates the static or partial structure factor
:math:`S_{\alpha\beta}(q)` using the radial histogram bins :math:`r`
and the radial distribution function :math:`g_{\alpha\beta}(r)` for
an isotropic fluid.
Parameters
----------
r : `numpy.ndarray`
Radii :math:`r`.
**Reference unit**: :math:`\mathrm{Å}`.
g : `numpy.ndarray`
Radial distribution function :math:`g_{\alpha\beta}(r)`.
**Shape**: Same as `r`.
equal : `bool`
Specifies whether `g` is between the same species
(:math:`\alpha=\beta`). If :code:`False`, the number
concentrations of species :math:`\alpha` and :math:`\beta` must
be specified in `x_a` and `x_b`.
rho : `float`
Bulk number density :math:`\rho` or surface density
:math:`\sigma`.
**Reference unit**: :math:`\mathrm{Å}^{-3}` or
:math:`\mathrm{Å}^{-2}`.
x_a : `float`, default: :code:`1`
Number concentration of species :math:`\alpha`. Required if
:code:`equal=False`.
x_b : `float`, optional
Number concentration of species :math:`\beta`. Required if
:code:`equal=False`.
q : `numpy.ndarray`, optional
Wavenumbers :math:`q`. If not specified, wavenumbers are
generated using `q_lower`, `q_upper`, and `n_q`.
**Reference unit**: :math:`\mathrm{Å}^{-1}`.
q_lower : `float`, keyword-only, optional
Lower bound for the wavenumbers :math:`q`. Has no effect if `q`
is specified.
**Reference unit**: :math:`\mathrm{Å}^{-1}`.
q_upper : `float`, keyword-only, optional
Upper bound for the wavenumbers :math:`q`. Has no effect if `q`
is specified.
**Reference unit**: :math:`\mathrm{Å}^{-1}`.
n_q : `int`, keyword-only, default: :code:`1_000`
Number of wavenumbers :math:`q` to generate. Has no effect if `q`
is specified.
n_dims : `int`, keyword-only, default: :code:`3`
Number of dimensions :math:`n_\mathrm{dims}`.
formalism : `str`, keyword-only, default: :code:`"FZ"`
Formalism to use for the partial structure factor. Has no effect
if :code:`equal=True`.
.. container::
**Valid values**:
* :code:`"general"`: A general formalism given by
.. math::
S_{\alpha\beta}(q)=1+x_\alpha x_\beta\frac{4\pi\rho}{q}
\int_0^\infty(g_{\alpha\beta}(r)-1)\sin{(qr)}r\,dr
* :code:`"FZ"`: Faber–Ziman formalism [1]_
.. math::
S_{\alpha\beta}(q)=1+\frac{4\pi\rho}{q}\int_0^\infty
(g_{\alpha\beta}(r)-1)\sin{(qr)}r\,dr
* :code:`"AL"`: Ashcroft–Langreth formalism [2]_
.. math::
S_{\alpha\beta}(q)=\delta_{ij}+(x_\alpha x_\beta)^{1/2}
\frac{4\pi\rho}{q}
\int_0^\infty(g_{\alpha\beta}(r)-1)\sin{(qr)}r\,dr
In two-dimensional systems, the second term is
.. math::
2\pi\rho\int_0^\infty (g_{\alpha\beta}(r)-1)J_0(qr)r\,dr
instead, where :math:`J_0` is the zeroth-order Bessel
function.
Returns
-------
q : `numpy.ndarray`
Wavenumbers :math:`q`.
**Shape**: :math:`(N_q,)`.
S : `numpy.ndarray`
Static or partial structure factor :math:`S_{\alpha\beta}(q)`.
**Shape**: :math:`(N_q,)`.
References
----------
.. [1] T. E. Faber and J. M. Ziman, A Theory of the Electrical
Properties of Liquid Metals: III. the Resistivity of Binary
Alloys, *Philosophical Magazine* **11**, **153** (1965).
https://doi.org/10.1080/14786436508211931
.. [2] N. W. Ashcroft and D. C. Langreth, Structure of Binary Liquid
Mixtures. I, *Phys. Rev.* **156**, **685** (1967).
https://doi.org/10.1103/PhysRev.156.685
"""
if q is None:
if q_lower is None:
q_lower = 2 * np.pi / r[-1]
if q_upper is None:
q_upper = 2 * np.pi / r[0]
q = np.linspace(
q_lower, q_upper, int((q_upper - q_lower) / q_lower) if n_q is None else n_q
)
if n_dims == 3:
_transform = radial_fourier_transform
elif n_dims == 2:
_transform = zeroth_order_hankel_transform
else:
raise ValueError("Invalid number of dimensions.")
rho_sft = rho * _transform(r, g - 1, q)
if equal or formalism == "FZ":
return q, 1 + rho_sft
elif not equal:
if formalism == "AL":
return q, (x_a == x_b) + np.sqrt(x_a * x_b) * rho_sft
elif formalism == "general":
return q, 1 + x_a * x_b * rho_sft
raise ValueError("Invalid formalism.")
[docs]
class RadialDistributionFunction(DynamicAnalysisBase):
"""
Serial and parallel implementations to calculate the radial
distribution function (RDF) :math:`g_{\\alpha\\beta}(r)` between
types :math:`\\alpha` and :math:`\\beta` and its related properties
for two- and three-dimensional systems.
The RDF is given by
.. math::
g_{\\alpha\\beta}^\\mathrm{3D}(r)
=\\frac{V}{4\\pi r^2N_\\alpha N_\\beta}\\sum_{i=1}^{N_\\alpha}
\\sum_{j=1}^{N_\\beta}\\left\\langle\\delta
\\left(|\\mathbf{r}_i-\\mathbf{r}_j|-r\\right)\\right\\rangle\\\\
g_{\\alpha\\beta}^\\mathrm{2D}(r)
=\\frac{A}{2\\pi rN_\\alpha N_\\beta}\\sum_{i=1}^{N_\\alpha}
\\sum_{j=1}^{N_\\beta}\\left\\langle\\delta
\\left(|\\mathbf{r}_i-\\mathbf{r}_j|-r\\right)\\right\\rangle
where :math:`V` and :math:`A` are the system volume and area,
:math:`N_\\alpha` and :math:`N_\\beta` are the number of particles
of species :math:`\\alpha` and :math:`\\beta`, and
:math:`\\mathbf{r}_i` and :math:`\\mathbf{r}_j` are the positions of
particles :math:`i` and :math:`j` belonging to species
:math:`\\alpha` and :math:`\\beta`, respectively. The RDF is
normalized such that
:math:`\\lim_{r\\rightarrow\\infty}g_{\\alpha\\beta}(r)=1` in a
homogeneous system.
A closely related quantity is the single particle density
:math:`n_{\\alpha\\beta}(r)=\\rho_\\beta g_{\\alpha\\beta}(r)`,
where :math:`\\rho_\\beta` is the number density of species
:math:`\\beta`.
The cumulative RDF is
.. math::
G_{\\alpha\\beta}^\\mathrm{3D}(r)
=4\\pi\\int_0^rR^2g_{\\alpha\\beta}(R)\\,dR\\\\
G_{\\alpha\\beta}^\\mathrm{2D}(r)
=2\\pi\\int_0^rRg_{\\alpha\\beta}(R)\\,dR
and the average number of :math:`\\beta` particles found within
radius :math:`r` is
.. math::
N_{\\alpha\\beta}(r)=\\rho_\\beta G_{\\alpha\\beta}(r)
The expression above can be used to compute the coordination numbers
(number of neighbors in each solvation shell) by setting :math:`r`
to values where :math:`g_{\\alpha\\beta}(r)` is locally minimized,
which signify the solvation shell boundaries.
.. container::
The RDF can also be used to obtain other relevant structural
properties, such as
* the potential of mean force
.. math::
w_{\\alpha\\beta}(r)
=-k_\\mathrm{B}T\\ln{g_{\\alpha\\beta}(r)}
where :math:`k_\\mathrm{B}` is the Boltzmann constant and
:math:`T` is the system temperature, and
* the static or partial structure factor (see
:func:`calculate_structure_factor` for the possible
definitions).
Parameters
----------
group_a : `MDAnalysis.AtomGroup`
Atom group containing species :math:`\\alpha`.
group_b : `MDAnalysis.AtomGroup`
Atom group containing species :math:`\\beta`. If not specified,
`group_a` is used for both species.
n_bins : `int`, default: :code:`201`
Number of histogram bins :math:`N_\\mathrm{bins}`.
range_ : array-like, default: :code:`(0.0, 15.0)`
Range of radii values. The upper bound should be less than half
the smallest system dimension.
**Shape**: :math:`(2,)`.
**Reference unit**: :math:`\\mathrm{Å}`.
exclusion : array-like, keyword-only, optional
Tiles to exclude from the interparticle distances. The
`groupings` parameter dictates what a tile represents.
**Shape**: :math:`(2,)`.
**Example**: :code:`(1, 1)` to exclude self-interactions.
groupings : `str` or array-like, keyword-only, \
default: :code:`"atoms"`
Determines whether the centers of mass are used in lieu of
individual atom positions. If `groupings` is a `str`, the same
value is used for all `groups`.
.. note::
If the desired grouping is not :code:`"atoms"`,
* the trajectory file should have segments (or chains)
containing residues (or molecules) and residues containing
atoms, and
* residues and segments should be locally unwrapped at the
simulation box edges, if not already, using
:class:`MDAnalysis.transformations.wrap.unwrap`,
:meth:`MDAnalysis.core.groups.AtomGroup.unwrap`, or
:func:`MDAnalysis.lib.mdamath.make_whole`.
.. container::
**Valid values**:
* :code:`"atoms"`: Atom positions (generally or for
coarse-grained simulations).
* :code:`"residues"`: Residues' centers of mass (for
atomistic simulations).
* :code:`"segments"`: Segments' centers of mass (for
atomistic polymer simulations).
drop_axis : `int` or `str`, keyword-only, optional
Axis in three-dimensional space to ignore in the two-dimensional
analysis.
**Valid values**: :code:`0` or :code:`x` for the :math:`x`-axis,
:code:`1` or :code:`y` for the :math:`y`-axis, and :code:`2` or
:code:`z` for the :math:`z`-axis.
norm : `str`, keyword-only, default: :code:`"rdf"`
Determines how the radial histograms are normalized.
.. container::
**Valid values**:
* :code:`norm="rdf"`: The radial distribution function
:math:`g_{\\alpha\\beta}(r)` is computed.
* :code:`norm="density"`: The single particle density
:math:`n_{\\alpha\\beta}(r)` is computed.
* :code:`norm=None`: The raw particle pair count in the
radial histogram bins is returned.
n_batches : `int`, keyword-only, optional
Number of batches to divide the histogram calculation into.
This is useful for large systems that cannot be processed in a
single pass.
.. note::
If you use too few bins and too many batches, the histogram
counts may be off by a few due to the floating-point nature
of the cutoffs. However, when the RDF is averaged over a
long trajectory with many particles, the difference should
be negligible.
reduced : `bool`, keyword-only, default: :code:`False`
Specifies whether the data is in reduced units.
parallel : `bool`, keyword-only, default: :code:`False`
Determines whether the analysis is performed in parallel.
.. note::
The Dask backend generally provides the best performance.
verbose : `bool`, keyword-only, default: :code:`True`
Determines whether detailed progress is shown.
**kwargs
Additional keyword arguments to pass to
:class:`MDAnalysis.analysis.base.AnalysisBase`.
Attributes
----------
universe : `MDAnalysis.Universe`
:class:`MDAnalysis.core.universe.Universe` object containing all
information describing the system.
results.units : `dict`
Reference units for the results. For example, to get the
reference units for :code:`results.bins`, call
:code:`results.units["bins"]`.
results.bin_edges : `numpy.ndarray`
Bin edges.
**Shape**: :math:`(N_\\mathrm{bins}+1,)`.
**Reference unit**: :math:`\\mathrm{Å}`.
results.bins : `numpy.ndarray`
Bin centers :math:`r`.
**Shape**: :math:`(N_\\mathrm{bins},)`.
**Reference unit**: :math:`\\mathrm{Å}`.
results.counts : `numpy.ndarray`
Raw particle pair counts in the radial histogram bins.
**Shape**: :math:`(N_\\mathrm{bins},)`.
results.rdf : `numpy.ndarray`
.. container::
One of
* :code:`norm="rdf"`: the radial distribution function
:math:`g_{\\alpha\\beta}(r)`,
* :code:`norm="density"`: the single particle density
:math:`n_{\\alpha\\beta}(r)`, or
* :code:`norm=None`: the raw particle pair count in the
radial histogram bins.
**Shape**: :math:`(N_\\mathrm{bins},)`.
results.coordination_numbers : `numpy.ndarray`
Coordination numbers :math:`n_k`. Only available after running
:meth:`calculate_coordination_numbers`.
results.pmf : `numpy.ndarray`
Potential of mean force :math:`w_{\\alpha\\beta}(r)`. Only
available after running :meth:`calculate_pmf`.
**Shape**: :math:`(N_\\mathrm{bins},)`.
**Reference unit**: :math:`\\mathrm{kJ/mol}`.
results.wavenumbers : `numpy.ndarray`
Wavenumbers :math:`q`. Only available after running
:meth:`calculate_structure_factor`.
**Reference unit**: :math:`\\mathrm{Å}^{-1}`.
results.ssf : `numpy.ndarray`
Static or partial structure factor. Only available after running
:meth:`calculate_structure_factor`.
**Shape**: Same as `results.wavenumbers`.
Example
--------
First, this analysis class must be imported:
>>> from mdcraft.analysis.structure import RadialDistributionFunction
Then, after loading a simulation trajectory:
>>> universe = mda.Universe("simulation.nc", "topology.cif")
We must then select the atom-group(s) to be analyzed:
>>> ag1 = universe.select_atoms("type 1")
>>> ag2 = universe.select_atoms("type 2")
The `RadialDistributionFunction` class can then be built and run (if only one atom group is provided, the RDF is calculated for the same atom group):
>>> rdf = RadialDistributionFunction(ag1, ag2)
>>> rdf.run()
The results can be obtained under the `results` attribute:
>>> rdf.results.rdf
One can also obtain the coordination numbers by running (specifying a reasonable threshold is recommended):
>>> rdf.calculate_coordination_numbers(rho=0.1,threshold=0.1)
Similarly, the potential of mean force can be calculated by running:
>>> rdf.calculate_pmf()
Finally, the static or partial structure factor can be calculated by running:
>>> rdf.calculate_structure_factor()
All of these results can be saved by doing the following:
>>> rdf.save("rdf")
"""
def __init__(
self,
group_a: mda.AtomGroup,
group_b: mda.AtomGroup = None,
n_bins: int = 201,
range_: tuple[float] = (0.0, 15.0),
*,
exclusion: tuple[int] = None,
groupings: Union[str, tuple[str]] = "atoms",
drop_axis: Union[int, str] = None,
norm: str = "rdf",
n_batches: int = None,
reduced: bool = False,
parallel: bool = False,
verbose: bool = True,
**kwargs,
) -> None:
self._groups = (group_a, group_a if group_b is None else group_b)
self.universe = self._groups[0].universe
if self.universe.dimensions is None:
raise ValueError(
"Trajectory does not contain system " "dimension information."
)
super().__init__(self.universe.trajectory, parallel, verbose, **kwargs)
GROUPINGS = {"atoms", "residues", "segments"}
if isinstance(groupings, str):
if groupings not in GROUPINGS:
emsg = (
f"Invalid grouping '{groupings}'. Valid values: "
"'" + "', '".join(GROUPINGS) + "'."
)
raise ValueError(emsg)
self._groupings = 2 * [groupings]
else:
if len(groupings) != 2:
emsg = "'groupings' must be a string or an array with length 2."
raise ValueError(emsg)
for gr in groupings:
if gr not in GROUPINGS:
emsg = (
f"Invalid grouping '{gr}'. Valid values: "
"'" + "', '".join(GROUPINGS) + "'."
)
raise ValueError(emsg)
self._groupings = groupings
self._drop_axis = (
ord(drop_axis) - 120 if isinstance(drop_axis, str) else drop_axis
)
if self._drop_axis not in {0, 1, 2, None}:
raise ValueError("Invalid axis in 'drop_axis'.")
self._n_bins = n_bins
self._range = range_
self._norm = norm
self._exclusion = exclusion
self._reduced = reduced
self._n_batches = n_batches
self._verbose = verbose
def _prepare(self) -> None:
# Preallocate floating-point number for system volume or area
if not self._parallel and self._norm == "rdf":
self._area_or_volume = 0.0
# Preallocate arrays to store results
self.results.bin_edges = accelerated.numba_histogram_bin_edges(
np.asarray(self._range, dtype=float), self._n_bins
)
self.results.bins = (
self.results.bin_edges[:-1] + self.results.bin_edges[1:]
) / 2
if not self._parallel:
self.results.counts = np.zeros(self._n_bins, dtype=int)
# Store reference units
self.results.units = Hash({"bins": ureg.angstrom, "edges": ureg.angstrom})
def _single_frame(self) -> None:
# Get system dimensions and orthogonality
dimensions = self._ts.dimensions
# Get positions or centers of mass of entities in the groups
positions = [
ag.positions if gr == "atoms" else center_of_mass(ag, gr)
for ag, gr in zip(self._groups, self._groupings)
]
# Add system volume or area analyzed in current frame
if self._drop_axis is None:
if self._norm == "rdf":
self._area_or_volume += self._ts.volume
else:
# Apply corrections to avoid including periodic images in
# the dimension to exclude
positions[0][:, self._drop_axis] = positions[1][:, self._drop_axis] = 0
dimensions[self._drop_axis] = dimensions[:3].max()
if self._norm == "rdf":
self._area_or_volume += np.delete(
dimensions[:3], self._drop_axis
).prod()
# Tally counts in each pair separation distance bin
if self._n_batches:
edges = np.array_split(self.results.bin_edges, self._n_batches)
ranges_indices = {
e: np.where((self.results.bins > e[0]) & (self.results.bins < e[1]))[0]
for e in [
(self._range[0], edges[0][-1]),
*((a[-1], b[-1]) for a, b in zip(edges[:-1], edges[1:])),
]
}
for r, i in ranges_indices.items():
self.results.counts[i] += radial_histogram(
*positions,
n_bins=i.shape[0],
range_=r,
dimensions=dimensions,
exclusion=self._exclusion,
)
else:
self.results.counts += radial_histogram(
*positions,
n_bins=self._n_bins,
range_=self._range,
dimensions=dimensions,
exclusion=self._exclusion,
)
def _single_frame_parallel(self, index: int) -> np.ndarray[float]:
# Set current trajectory frame
ts = self._sliced_trajectory[index]
# Preallocate array to store entity counts followed by system
# volume
results = np.empty(1 + self._n_bins)
# Get system dimensions and orthogonality
dimensions = ts.dimensions
# Get positions or centers of mass of entities in the groups
positions = [
ag.positions if gr == "atoms" else center_of_mass(ag, gr)
for ag, gr in zip(self._groups, self._groupings)
]
# Store system volume or area analyzed in current frame
if self._drop_axis is None:
results[self._n_bins] = ts.volume
else:
# Apply corrections to avoid including periodic images in
# the dimension to exclude
positions[0][:, self._drop_axis] = positions[1][:, self._drop_axis] = 0
dimensions[self._drop_axis] = dimensions[:3].max()
results[self._n_bins] = np.delete(dimensions[:3], self._drop_axis).prod()
# Compute radial histogram for current frame
if self._n_batches:
edges = np.array_split(self.results.bin_edges, self._n_batches)
ranges_indices = {
e: np.where((self.results.bins > e[0]) & (self.results.bins < e[1]))[0]
for e in [
(self._range[0], edges[0][-1]),
*((a[-1], b[-1]) for a, b in zip(edges[:-1], edges[1:])),
]
}
for r, i in ranges_indices.items():
results[i] = radial_histogram(
*positions,
n_bins=i.shape[0],
range_=r,
dimensions=dimensions,
exclusion=self._exclusion,
)
else:
results[: self._n_bins] = radial_histogram(
*positions,
n_bins=self._n_bins,
range_=self._range,
dimensions=dimensions,
exclusion=self._exclusion,
)
return results
def _conclude(self) -> None:
# Tally counts in each pair separation distance bin over all
# frames
if self._parallel:
self._results = np.vstack(self._results).sum(axis=0)
self.results.counts = self._results[: self._n_bins]
self._area_or_volume = self._results[self._n_bins]
# Compute the normalization factor
norm = self.n_frames
if self._norm is not None:
if self._drop_axis is None:
norm *= 4 * np.pi * np.diff(self.results.bin_edges**3) / 3
else:
norm *= np.pi * np.diff(self.results.bin_edges**2)
if self._norm == "rdf":
N_2 = getattr(self._groups[1], f"n_{self._groupings[1]}")
if self._exclusion:
N_2 -= self._exclusion[1]
norm *= (
getattr(self._groups[0], f"n_{self._groupings[0]}")
* N_2
* self.n_frames
/ self._area_or_volume
)
# Compute and store the radial distribution function, the single
# particle density, or the raw radial pair counts
self.results.rdf = self.results.counts / norm
def _get_rdf(self) -> np.ndarray[float]:
r"""
Returns the existing radial distribution function (RDF) if
:code:`norm="rdf"` was passed to the :class:`RDF` constructor.
Otherwise, the RDF is calculated and returned.
Returns
-------
rdf : `numpy.ndarray`
Radial distribution function :math:`g_{\alpha\beta}(r)`.
"""
if self._norm == "rdf":
return self.results.rdf
N_2 = getattr(self._groups[1], f"n_{self._groupings[1]}")
if self._exclusion:
N_2 -= self._exclusion[1]
if self._drop_axis is None:
norm = 4 * np.diff(self.results.bin_edges**3) / 3
else:
norm = np.diff(self.results.bin_edges**2)
return (
self._area_or_volume
* self.results.counts
/ (
np.pi
* self.n_frames**2
* N_2
* norm
* getattr(self._groups[0], f"n_{self._groupings[0]}")
)
)
[docs]
def calculate_coordination_numbers(
self, rho: float, *, n_coord_nums: int = 2, threshold: float = 0.1
) -> None:
r"""
Calculates the coordination numbers :math:`n_k`.
If the radial distribution function :math:`g_{\alpha\beta}(r)`
does not contain :math:`k` local minima, this method will return
`numpy.nan` for the coordination numbers that could not be
calculated.
Parameters
----------
rho : `float`
Number density :math:`\rho_\beta` of species :math:`\beta`.
**Reference unit**: :math:`\mathrm{nm}^{-3}`.
n_coord_nums : `int`, keyword-only, default: :code:`2`
Number of coordination numbers :math:`n_\mathrm{coord}` to
calculate.
threshold : `float`, keyword-only, default: :code:`0.1`
Minimum :math:`g_{\alpha\beta}(r)` value that must be
reached before local minima are found.
"""
self.results.coordination_numbers = calculate_coordination_numbers(
self.results.bins,
self._get_rdf(),
rho,
n_coord_nums=n_coord_nums,
n_dims=2 + (self._drop_axis is None),
threshold=threshold,
)
[docs]
def calculate_pmf(self, temperature: Union[float, "unit.Quantity", Q_]) -> None:
r"""
Calculates the potential of mean force
:math:`w_{\alpha\beta}(r)`.
Parameters
----------
temperature : `float` or `openmm.unit.Quantity`
System temperature :math:`T`.
.. note::
If :code:`reduced=True` was set in the :class:`RDF`
constructor, `temperature` should be equal to the energy
scale. When the Lennard-Jones potential is used, it
generally means that :math:`T^*=1`, or `temperature=1`.
**Reference unit**: :math:`\mathrm{K}`.
"""
if self._reduced and not is_unitless(temperature):
emsg = "'temperature' cannot have units when reduced=True."
raise ValueError(emsg)
self.results.units["pmf"] = ureg.kilojoule / ureg.mole
kBT = strip_unit(temperature, "K")[0]
if not self._reduced:
kBT *= (
ureg.avogadro_constant * ureg.boltzmann_constant * ureg.kelvin
).m_as(self.results.units["pmf"])
self.results.pmf = -kBT * np.log(self._get_rdf())
[docs]
def calculate_structure_factor(
self,
rho: float,
x_a: float = None,
x_b: float = None,
q: np.ndarray[float] = None,
*,
q_lower: float = None,
q_upper: float = None,
n_q: int = 1_000,
formalism: str = "FZ",
) -> None:
r"""
Computes the static or partial structure factor
:math:`S_{\alpha\beta}(q)` using the radial histogram bins
:math:`r` and the radial distribution function
:math:`g_{\alpha\beta}(r)` for an isotropic fluid.
Parameters
----------
rho : `float`
Bulk number density :math:`\rho`.
**Reference unit**: :math:`\mathrm{Å}^{-3}`.
x_a : `float`, default: :code:`1`
Number concentration of species :math:`\alpha`. Required if
the two atom groups are not identical.
x_b : `float`, optional
Number concentration of species :math:`\beta`. Required if
the two atom groups are not identical.
q : `numpy.ndarray`, optional
Wavenumbers :math:`q`.
**Reference unit**: :math:`\mathrm{Å}^{-1}`.
q_lower : `float`, keyword-only, optional
Lower bound for the wavenumbers :math:`q`. Has no effect if `q`
is specified.
**Reference unit**: :math:`\mathrm{Å}^{-1}`.
q_upper : `float`, keyword-only, optional
Upper bound for the wavenumbers :math:`q`. Has no effect if `q`
is specified.
**Reference unit**: :math:`\mathrm{Å}^{-1}`.
n_q : `int`, keyword-only, default: :code:`1_000`
Number of wavenumbers :math:`q` to generate. Has no effect if `q`
is specified.
formalism :`str`, keyword-only, default: :code:`"FZ"`
Formalism to use for the partial structure factor. Has no effect
if the two atom groups are the same.
.. container::
**Valid values**:
* :code:`"general"`: A general formalism similar to that
of the static structure factor, except the second term
is multiplied by :math:`x_ix_j`.
* :code:`"FZ"`: Faber–Ziman formalism.
* :code:`"AL"`: Ashcroft–Langreth formalism.
.. seealso::
For more information, see
:func:`calculate_structure_factor`.
"""
self.results.wavenumbers, self.results.ssf = calculate_structure_factor(
self.results.bins,
self._get_rdf(),
self._groups[0] == self._groups[1],
rho,
x_a,
x_b,
q=q,
q_lower=q_lower,
q_upper=q_upper,
n_q=n_q,
n_dims=2 + (self._drop_axis is None),
formalism=formalism,
)
[docs]
@numba.njit("f8(f8[:])", fastmath=True)
def numba_pythagorean_trigonometric_identity(r: np.ndarray[float]) -> float:
r"""
Serial Numba-accelerated evaluation of the Pythagorean trigonometric
identity for a one-dimensional NumPy array :math:`\mathbf{r}`.
.. math::
\left(\sum_{i=1}^3\cos(r_i)\right)^2
+\left(\sum_{i=1}^3\sin(r_i)\right)^2
Parameters
----------
r : `np.ndarray`
Vector :math:`\mathbf{r}`.
**Shape**: :math:`(N_r,)`.
Returns
-------
c2_s2 : `float`
Pythagorean trigonometric identity for the vector
:math:`\mathbf{r}`.
"""
c = s = 0
for i in range(r.shape[0]):
c += np.cos(r[i])
s += np.sin(r[i])
return c**2 + s**2
[docs]
@numba.njit("f8(f8[:],f8[:])", fastmath=True)
def numba_cross_pythagorean_trigonometric_identity(
r: np.ndarray[float], s: np.ndarray[float]
) -> float:
r"""
Serial Numba-accelerated evaluation of the cross Pythagorean
trigonometric identity for two one-dimensional NumPy arrays
:math:`\mathbf{r}` and :math:`\mathbf{s}`.
.. math::
2\left(\sum_{i=1}^3\cos(r_i)\sum_{j=1}^3\cos(s_j)
+\sum_{i=1}^3\sin(r_i)\sum_{j=1}^3\sin(s_j)\right)
Parameters
----------
r : `np.ndarray`
First vector :math:`\mathbf{r}`.
**Shape**: :math:`(N_r,)`.
s : `np.ndarray`
Second vector :math:`\mathbf{s}`.
**Shape**: :math:`(N_s,)`.
Returns
-------
c2_s2 : `float`
Cross Pythagorean trigonometric identity for the vectors
:math:`\mathbf{r}` and :math:`\mathbf{s}`.
"""
c1 = c2 = s1 = s2 = 0
for i in range(r.shape[0]):
c1 += np.cos(r[i])
s1 += np.sin(r[i])
for j in range(s.shape[0]):
c2 += np.cos(s[j])
s2 += np.sin(s[j])
return 2 * (c1 * c2 + s1 * s2)
[docs]
def ssf_trigonometric(qrs: np.ndarray[float]) -> np.ndarray[float]:
r"""
Computes the static structure factors using a two-dimensional NumPy
array containing :math:`\mathbf{q}\cdot\mathbf{r}` using the
trigonometric form.
.. math::
NS(q)=\left\langle\left(\sum_{j=1}^N
\cos{(\mathbf{q}\cdot\mathbf{r}_j)}\right)^2+\left(
\sum_{j=1}^N\sin{(\mathbf{q}\cdot\mathbf{r}_j)}
\right)^2\right\rangle
Parameters
----------
qrs : `np.ndarray`
Inner products :math:`\mathbf{q}\cdot\mathbf{r}_j`.
**Shape**: :math:`(N_q,\,N_r)`.
Returns
-------
ssf : `np.ndarray`
Static structure factors.
**Shape**: :math:`(N_q,)`.
"""
ssf = np.empty(qrs.shape[0])
for i in numba.prange(qrs.shape[0]):
ssf[i] = numba_pythagorean_trigonometric_identity(qrs[i])
return ssf
[docs]
def psf_trigonometric(
qrs1: np.ndarray[float], qrs2: np.ndarray[float]
) -> np.ndarray[float]:
r"""
Computes the partial structure factors given two two-dimensional
NumPy arrays, each containing :math:`\mathbf{q}\cdot\mathbf{r}`,
using the trigonometric form.
.. math::
\frac{NS_{\alpha\beta}(q)}{2-\delta_{\alpha\beta}}
=\left\langle
\sum_{j=1}^{N_\alpha}\cos{(\mathbf{q}\cdot\mathbf{r}_j)}
\sum_{k=1}^{N_\beta}\cos{(\mathbf{q}\cdot\mathbf{r}_k)}
+\sum_{j=1}^{N_\alpha}\sin{(\mathbf{q}\cdot\mathbf{r}_j)}
\sum_{k=1}^{N_\beta}\sin{(\mathbf{q}\cdot\mathbf{r}_k)}
\right\rangle
Parameters
----------
qrs1 : `np.ndarray`
First set of inner products :math:`\mathbf{q}\cdot\mathbf{r}_j`.
**Shape**: :math:`(N_q,\,N_r)`.
qrs2 : `np.ndarray`
Second set of inner products :math:`\mathbf{q}\cdot\mathbf{r}_k`.
**Shape**: :math:`(N_q,\,N_r)`.
Returns
-------
ssf : `np.ndarray`
Partial structure factors.
**Shape**: :math:`(N_q,)`.
"""
ssf = np.empty(qrs1.shape[0])
for i in numba.prange(qrs1.shape[0]):
ssf[i] = numba_cross_pythagorean_trigonometric_identity(qrs1[i], qrs2[i])
return ssf
[docs]
def generate_spherical_wavevectors(
dimensions: np.ndarray[float],
*,
n_points: Union[int, tuple[int]] = 32,
q_max: float = None,
wavenumbers: bool = False,
) -> np.ndarray[float]:
r"""
Generates spherical wavevectors :math:`\mathbf{q}` for a triclinic
simulation box.
Parameters
----------
dimensions : `np.ndarray`
Simulation box dimensions.
**Shape**: :math:`(3,\,3)`.
n_points : `int` or `tuple`, default: :code:`32`
Number of points along each axis. If an `int` is provided, the
same value is used for all axes.
q_max : `float`, optional
Maximum wavenumber :math:`q`.
**Reference unit**: :math:`\mathrm{Å}^{-1}`.
wavenumbers : `bool`, default: :code:`False`
Determines whether the wavenumbers are returned.
Returns
-------
wavevectors : `np.ndarray`
Wavevectors :math:`\mathbf{q}`.
**Shape**: :math:`(N_q,\,3)`.
wavenumbers : `np.ndarray`
Wavenumbers :math:`q`.
**Shape**: :math:`(N_q,\,)`.
"""
dimensions = np.asarray(dimensions)
if dimensions.ndim == 1:
dimensions = convert_cell_representation(parameters=dimensions)
two_pi_over_Ls = 2 * np.pi * np.linalg.inv(dimensions.T)
if q_max is None:
if isinstance(n_points, int):
n_points = 3 * [n_points]
elif len(n_points) != 3:
emsg = "'n_points' must be an integer or a tuple with length 3."
raise ValueError(emsg)
else:
n_points = np.floor(
q_max * np.linalg.norm(np.linalg.inv(two_pi_over_Ls.T), axis=1)
).astype(int)
wavevectors = (
np.stack(np.meshgrid(*[np.arange(n) for n in n_points]), axis=-1).reshape(-1, 3)
[docs]
@ two_pi_over_Ls
)
wavenumbers_ = np.linalg.norm(wavevectors, axis=1)
if q_max is not None:
keep_indices = wavenumbers_ <= q_max
wavevectors = wavevectors[keep_indices]
wavenumbers_ = wavenumbers_[keep_indices]
if wavenumbers:
return wavevectors, wavenumbers_
return wavevectors
class StructureFactor(NumbaAnalysisBase):
"""
Serial and parallel implementations to calculate the static
structure factor :math:`S(q)` or partial structure factor
:math:`S_{\\alpha\\beta}(q)` for species :math:`\\alpha` and
:math:`\\beta`.
The static structure factor is a measure of how a material scatters
incident radiation, and can be computed directly from a molecular
dynamics simulation trajectory using
.. math::
S(q)&=\\frac{1}{N}\\left\\langle\\sum_{j=1}^N\\sum_{k=1}^N
\\exp{[-i\\mathbf{q}\\cdot(\\mathbf{r}_j-\\mathbf{r}_k)]}
\\right\\rangle\\\\&=\\frac{1}{N}\\left\\langle\\left(
\\sum_{j=1}^N\\sin{(\\mathbf{q}\\cdot\\mathbf{r}_j)}\\right)^2
+\\left(\\sum_{j=1}^N\\cos{(\\mathbf{q}\\cdot\\mathbf{r}_j)}
\\right)^2\\right\\rangle
where :math:`N` is the total number of particles or centers of mass,
:math:`\\mathbf{q}` is the scattering wavevector, and
:math:`\\mathbf{r}_i` is the position of the :math:`i`-th particle.
For multicomponent systems, the equation above can be generalized to
get the partial structure factors
.. math::
S_{\\alpha\\beta}(q)&=\\frac{2-\\delta_{\\alpha\\beta}}{N}
\\left\\langle\\sum_{j=1}^{N_\\alpha}\\sum_{k=1}^{N_\\beta}
\\exp{[-i\\mathbf{q}\\cdot(\\mathbf{r}_j-\\mathbf{r}_k)]}
\\right\\rangle\\\\
&=\\frac{2-\\delta_{\\alpha\\beta}}{N}\\left\\langle
\\sum_{j=1}^{N_\\alpha}\\cos{(\\mathbf{q}\\cdot\\mathbf{r}_j)}
\\sum_{k=1}^{N_\\beta}\\cos{(\\mathbf{q}\\cdot\\mathbf{r}_k)}
+\\sum_{j=1}^{N_\\alpha}\\sin{(\\mathbf{q}\\cdot\\mathbf{r}_j)}
\\sum_{k=1}^{N_\\beta}\\sin{(\\mathbf{q}\\cdot\\mathbf{r}_k)}
\\right\\rangle
where :math:`\\delta_{ij}` is the Kronecker delta, :math:`N_\\alpha`
and :math:`N_\\beta` are the numbers of particles or centers of mass
for species :math:`\\alpha` and :math:`\\beta`.
The partial structure factors :math:`S_{\\alpha\\beta}(q)` and the
static structure factor :math:`S(q)` are related via
.. math::
S(q)=\\sum_{\\alpha=1}^{N_\\mathrm{g}}
\\sum_{\\beta=\\alpha}^{N_\\mathrm{g}}S_{\\alpha\\beta}(q)
Parameters
----------
groups : `MDAnalysis.AtomGroup` or array-like
Groups of atoms that share the same grouping type. If
:code:`mode=None`, all atoms in the universe must be assigned to
a group. If :code:`mode="pair"`, there must be exactly one or
two groups.
groupings : `str` or array-like, default: :code:`"atoms"`
Determines whether the centers of mass are used in lieu of
individual atom positions. If `groupings` is a `str`, the same
value is used for all `groups`.
.. note::
If the desired grouping is not :code:`"atoms"`,
* the trajectory file should have segments (or chains)
containing residues (or molecules) and residues containing
atoms, and
* residues and segments should be locally unwrapped at the
simulation box edges, if not already, using
:class:`MDAnalysis.transformations.wrap.unwrap`,
:meth:`MDAnalysis.core.groups.AtomGroup.unwrap`, or
:func:`MDAnalysis.lib.mdamath.make_whole`.
.. container::
**Valid values**:
* :code:`"atoms"`: Atom positions (generally or for
coarse-grained simulations).
* :code:`"residues"`: Residues' centers of mass (for
atomistic simulations).
* :code:`"segments"`: Segments' centers of mass (for
atomistic polymer simulations).
mode : `str`, keyword-only, optional
Evaluation mode.
.. container::
**Valid values**:
* :code:`None`: The static structure factor is computed.
* :code:`"pair"`: The partial structure factor is computed
between the groups in `groups`.
* :code:`"partial"`: The partial structure factors for all
unique pairs in `groups` is computed.
form : `str`, keyword-only, default: :code:`"exp"`
Expression used to evaluate the structure factors.
.. container::
**Valid values**:
* :code:`"exp"`: Exponential form. Slightly faster due to
fewer mathematical operations.
* :code:`"trig"`: Trigonometric form. Slightly slower but
doesn't have overflow issues.
dimensions : array-like, `openmm.unit.Quantity`, or \
`pint.Quantity`, keyword-only, optional
System dimensions :math:`(L_x,\\,L_y,\\,L_z)`. Only used to
determine the scattering wavevectors when `wavevectors` is not
specified.
**Shape**: :math:`(3,)`.
**Reference unit**: :math:`\\mathrm{Å}`.
n_points : `int`, keyword-only, default: :code:`32`
Number of points :math:`n_\\mathrm{points}` in the scattering
wavevector grid to generate
.. math::
\\mathbf{q}=2\\pi\\left(\\frac{a}{L_x},\\,\\frac{b}{L_y},\\,
\\frac{c}{L_z}\\right)
where :math:`a`, :math:`b`, and :math:`c` are integers from
:math:`0` up to :math:`n_\\mathrm{points}-1`.
Additional wavevectors can be introduced via `n_surfaces` and
`n_surface_points` for more accurate structure factors at small
wavenumbers. Alternatively, the desired wavevectors can be
specified directly in `wavevectors`.
n_surfaces : `int`, keyword-only, optional
Number of spherical surfaces in the first octant that intersect
with the grid wavevectors along the three coordinate axes for
which to introduce extra wavevectors for more accurate structure
factor values. Only available if the system is perfectly cubic.
n_surface_points : `int`, keyword-only, default: :code:`8`
Number of extra wavevectors to introduce per spherical surface.
Has no effect if `n_surfaces` is not specified.
q_max : `float`, `openmm.unit.Quantity`, or `pint.Quantity`, \
keyword-only, optional
Maximum wavenumber :math:`q_\\mathrm{max}`.
**Reference unit**: :math:`\\mathrm{Å}^{-1}`.
wavevectors : array-like, `openmm.unit.Quantity`, or `pint.Quantity`, \
keyword-only, optional
Scattering wavevectors for which to compute structure factors.
Has precedence over `n_points`, `n_surfaces`, and
`n_surface_points` if specified.
**Shape**: :math:`(N_q,\\,3)`.
**Reference unit**: :math:`\\mathrm{Å}^{-1}`.
sort : `bool`, keyword-only, default: :code:`True`
Determines whether the results are sorted by the wavenumbers.
unique : `bool`, keyword-only, default: :code:`True`
Determines whether structure factors for the same wavenumber
are grouped and averaged.
parallel : `bool`, keyword-only, default: :code:`False`
Determines whether the analysis is performed in parallel.
verbose : `bool`, keyword-only, default: :code:`True`
Determines whether detailed progress is shown.
**kwargs
Additional keyword arguments to pass to
:class:`MDAnalysis.analysis.base.AnalysisBase`.
Attributes
----------
universe : `MDAnalysis.Universe`
:class:`MDAnalysis.core.universe.Universe` object containing all
information describing the system.
results.units : `dict`
Reference units for the results. For example, to get the
reference units for :code:`results.wavenumbers`, call
:code:`results.units["wavenumbers"]`.
results.pairs : `tuple`
All unique pairs of indices of the groups of atoms in `groups`.
The ordering coincides with the column indices in
`results.ssf`.
results.wavenumbers : `numpy.ndarray`
Wavenumbers :math:`q`.
**Shape**: :math:`(N_q,)`.
**Reference unit**: :math:`\\mathrm{Å}^{-1}`.
results.ssf : `numpy.ndarray`
Static structure factor :math:`S(q)` or partial structure
factors :math:`S_{\\alpha\\beta}(q)`.
**Shape**: :math:`(1,\\,N_q)` or
:math:`(C(N_\\mathrm{g}+1,\\,2),\\,N_q)`.
Example
--------
First, this analysis class must be imported:
>>> from mdcraft.analysis.structure import StructureFactor
Then, after loading a simulation trajectory:
>>> universe = mda.Universe("simulation.nc", "topology.cif")
We must then select the atom-groups to be analyzed:
>>> ag = universe.select_atoms("all")
Note that if one wishes to obtain the total structure factor, all particles must be included in `ag`. The `StructureFactor` class can then be built and run:
>>> ssf = StructureFactor(ag)
>>> ssf.run()
.. note::
If `parallel=True`, the analysis will be performed in parallel using all available cores.
The results can be obtained under the `results` attribute:
>>> ssf.results.ssf
To save the results to a file, use the `save` method:
>>> ssf.save("ssf")
"""
def __init__(
self,
groups: Union[mda.AtomGroup, tuple[mda.AtomGroup]],
groupings: Union[str, tuple[str]] = "atoms",
*,
mode: str = None,
form: str = "exp",
dimensions: Union[np.ndarray[float], "unit.Quantity", Q_] = None,
n_points: int = 32,
n_surfaces: int = None,
n_surface_points: int = 8,
q_max: Union[float, "unit.Quantity", Q_] = None,
wavevectors: np.ndarray[float] = None,
sort: bool = True,
unique: bool = True,
parallel: bool = False,
verbose: bool = True,
**kwargs,
) -> None:
self._groups = [groups] if isinstance(groups, mda.AtomGroup) else groups
self.universe = self._groups[0].universe
super().__init__(self.universe.trajectory, verbose, **kwargs)
GROUPINGS = {"atoms", "residues"}
self._n_groups = len(self._groups)
if isinstance(groupings, str):
if groupings not in GROUPINGS:
emsg = (
f"Invalid grouping '{groupings}'. Valid values: "
"'" + "', '".join(GROUPINGS) + "'."
)
raise ValueError(emsg)
self._groupings = self._n_groups * [groupings]
else:
if self._n_groups != len(groupings):
emsg = (
"The shape of 'groupings' is incompatible with " "that of 'groups'."
)
raise ValueError(emsg)
for gr in groupings:
if gr not in GROUPINGS:
emsg = (
f"Invalid grouping '{gr}'. Valid values: "
"'" + "', '".join(GROUPINGS) + "'."
)
raise ValueError(emsg)
self._groupings = groupings
self._mode = mode
if self._mode == "pair" and not 1 <= len(self._groups) <= 2:
emsg = (
"There must be exactly one or two atom groups in "
"'groups' when mode='pair'."
)
raise ValueError(emsg)
elif self._mode is None:
if sum(g.n_atoms for g in self._groups) != self.universe.atoms.n_atoms:
emsg = (
"The atom groups in 'groups' must contain all "
"atoms in the simulation when 'mode=None'."
)
raise ValueError(emsg)
if wavevectors is not None:
self._wavevectors = wavevectors
self._wavenumbers = np.linalg.norm(wavevectors, axis=1)
else:
dimensions = strip_unit(dimensions or self.universe.dimensions[:3], "Å")[0]
if dimensions is None:
emsg = (
"System dimensions were not found, but are "
"required when 'wavevectors' is not specified."
)
raise ValueError(emsg)
elif len(dimensions) != 3:
raise ValueError("'dimensions' must have length 3.")
dimensions = np.asarray(dimensions)
if np.allclose(dimensions, dimensions[0]):
grid = 2 * np.pi * np.arange(n_points) / dimensions[0]
self._wavevectors = np.stack(np.meshgrid(grid, grid, grid), -1).reshape(
-1, 3
)
if n_surfaces:
n_theta, n_phi = get_closest_factors(
n_surface_points, 2, reverse=True
)
theta = np.linspace(
np.pi / (2 * n_theta + 4),
np.pi / 2 - np.pi / (2 * n_theta + 4),
n_theta,
)
phi = np.linspace(
np.pi / (2 * n_phi + 4),
np.pi / 2 - np.pi / (2 * n_phi + 4),
n_phi,
)
self._wavevectors = np.vstack(
(
self._wavevectors,
np.einsum(
"o,tpd->otpd",
grid[1 : n_surfaces + 1],
np.stack(
(
np.sin(theta) * np.cos(phi)[:, None],
np.sin(theta) * np.sin(phi)[:, None],
np.tile(np.cos(theta)[None, :], (n_phi, 1)),
),
axis=-1,
),
).reshape((n_surfaces * n_surface_points, 3)),
)
)
self._wavenumbers = np.linalg.norm(self._wavevectors, axis=1)
else:
self._wavevectors = np.stack(
np.meshgrid(
*[2 * np.pi * np.arange(n_points) / L for L in dimensions]
),
axis=-1,
).reshape(-1, 3)
self._wavenumbers = np.linalg.norm(self._wavevectors, axis=1)
if q_max is not None:
q_max = strip_unit(q_max, "Å^-1")[0]
keep = self._wavenumbers <= q_max
self._wavevectors = self._wavevectors[keep]
self._wavenumbers = self._wavenumbers[keep]
self._Ns = np.fromiter(
(getattr(a, f"n_{g}") for (a, g) in zip(self._groups, self._groupings)),
dtype=int,
count=self._n_groups,
)
self._N = self._Ns.sum()
self._slices = []
_ = 0
for N in self._Ns:
self._slices.append(slice(_, _ + N))
_ += N
self._njit = lambda s: numba.njit(s, fastmath=True, parallel=parallel)
self._delta_fourier_transform_sum = self._njit("c16[:](f8[:,:],f8[:,:])")(
delta_fourier_transform_sum
)
self._ssf_trigonometric = self._njit("f8[:](f8[:,:])")(ssf_trigonometric)
self._psf_trigonometric = self._njit("f8[:](f8[:,:],f8[:,:])")(
psf_trigonometric
)
self._inner = (
accelerated.numba_inner_parallel if parallel else accelerated.numba_inner
)
self._form = form
self._sort = sort
self._unique = unique
self._verbose = verbose
def _prepare(self) -> None:
# Determine all unique pairs
self.results.pairs = (
tuple(combinations_with_replacement(range(self._n_groups), 2))
if self._mode == "partial"
else ((0, self._n_groups - 1),) if self._mode == "pair" else ((None, None),)
)
# Create a persisting array to hold atom positions for a single
# frame so that it doesn't have to be recreated every frame
self._positions = np.empty((self._N, 3))
# Preallocate arrays to store structure factors
self.results.ssf = np.zeros((len(self.results.pairs), len(self._wavenumbers)))
# Determine the unique wavenumbers, if desired
self.results.wavenumbers = (
np.unique(self._wavenumbers.round(11))
if self._unique
else self._wavenumbers
)
# Store reference units
self.results.units = Hash({"wavenumbers": ureg.angstrom**-1})
def _single_frame(self) -> None:
# Store positions or centers of mass in the current frame
for g, gr, s in zip(self._groups, self._groupings, self._slices):
self._positions[s] = g.positions if gr == "atoms" else center_of_mass(g, gr)
# Compute the structure factor by multiplying exp(iqr) by its
# conjugates
if self._form == "exp":
if self._mode is None:
rhos = self._delta_fourier_transform_sum(
self._wavevectors, self._positions
)
self.results.ssf += (rhos * rhos.conj()).real
else:
for i, (j, k) in enumerate(self.results.pairs):
rhos_j = self._delta_fourier_transform_sum(
self._wavevectors, self._positions[self._slices[j]]
)
if j == k:
self.results.ssf[i] += (rhos_j * rhos_j.conj()).real
else:
self.results.ssf[i] += (
2
* (
rhos_j
* self._delta_fourier_transform_sum(
self._wavevectors, self._positions[self._slices[k]]
).conj()
).real
)
# Compute the structure factor by summing cos(qr)^2 and sin(qr)^2
elif self._form == "trig":
if self._mode is None:
self.results.ssf += self._ssf_trigonometric(
self._inner(self._wavevectors, self._positions)
)
else:
for i, (j, k) in enumerate(self.results.pairs):
qrs_j = self._inner(
self._wavevectors, self._positions[self._slices[j]]
)
if j == k:
self.results.ssf[i] += self._ssf_trigonometric(qrs_j)
else:
self.results.ssf[i] += self._psf_trigonometric(
qrs_j,
self._inner(
self._wavevectors, self._positions[self._slices[k]]
),
)
def _conclude(self) -> None:
# Normalize the structure factors by the number of entities and
# timesteps
self.results.ssf /= self.n_frames * self._N
# Combine values sharing the same wavenumber, if desired
if self._unique:
self.results.ssf = np.hstack(
[
self.results.ssf[:, np.isclose(q, self._wavenumbers)].mean(
axis=1, keepdims=True
)
for q in self.results.wavenumbers
]
)
# Sort the results by wavenumber, if desired
if self._sort:
order = np.argsort(self.results.wavenumbers)
self.results.wavenumbers = self.results.wavenumbers[order]
self.results.ssf = self.results.ssf[:, order]
# Clean up memory by deleting arrays that will not be reused
del self._positions
[docs]
@numba.njit("f8(f8[:])", fastmath=True)
def numba_cosine_sum(x: np.ndarray[float]) -> float:
r"""
Serial Numba-accelerated sum of the cosines of the elements in a
one-dimensional NumPy array :math:`\mathbf{x}`.
.. math::
\sum_{i=1}^N\cos(x_i)
Parameters
----------
x : `np.ndarray`
Vector :math:`\mathbf{x}`.
**Shape**: :math:`(N,)`.
Returns
-------
s : `float`
Sum of the cosines of the elements in the vector
:math:`\mathbf{x}`.
"""
s = 0
for i in range(x.shape[0]):
s += np.cos(x[i])
return s
[docs]
def cosine_column_sum(xs: np.ndarray[float]) -> np.ndarray[float]:
r"""
Evaluates the column-wise sum of the cosines of the elements in a
two-dimensional NumPy array :math:`\mathbf{x}`.
.. math::
\sum_{i=1}^N\cos(x_{ij})
Parameters
----------
xs : `np.ndarray`
Matrix :math:`\mathbf{x}`.
**Shape**: :math:`(M,\,N)`.
Returns
-------
s : `np.ndarray`
Column-wise sums of the cosines of the elements in the matrix
:math:`\mathbf{x}`.
**Shape**: :math:`(M,)`.
"""
s = np.empty(xs.shape[0])
for i in numba.prange(xs.shape[0]):
s[i] = numba_cosine_sum(xs[i])
return s
[docs]
def cosine_column_sum_inplace(xs: np.ndarray[float], s: np.ndarray[float]) -> None:
r"""
Evaluates the column-wise sum of the cosines of the elements in a
two-dimensional NumPy array :math:`\mathbf{x}`.
.. math::
\sum_{i=1}^N\cos(x_{ij})
Parameters
----------
xs : `np.ndarray`
Matrix :math:`\mathbf{x}`.
**Shape**: :math:`(M,\,N)`.
s : `np.ndarray`
Array to hold the column-wise sums of the cosines of the
elements in the matrix :math:`\mathbf{x}`.
**Shape**: :math:`(N,)`.
"""
assert s.shape[0] == xs.shape[0]
for i in numba.prange(xs.shape[0]):
s[i] = numba_cosine_sum(xs[i])
[docs]
@numba.njit("f8(f8[:])", fastmath=True)
def numba_sine_sum(x: np.ndarray[float]) -> float:
r"""
Serial Numba-accelerated sum of the sines of the elements in a
one-dimensional NumPy array :math:`\mathbf{x}`.
.. math::
\sum_{i=1}^N\sin(x_i)
Parameters
----------
x : `np.ndarray`
Vector :math:`\mathbf{x}`.
**Shape**: :math:`(N,)`.
Returns
-------
s : `float`
Sum of the sines of the elements in the vector
:math:`\mathbf{x}`.
"""
s = 0
for i in range(x.shape[0]):
s += np.sin(x[i])
return s
[docs]
def sine_column_sum(xs: np.ndarray[float]) -> None:
r"""
Evaluates the column-wise sum of the sines of the elements in a
two-dimensional NumPy array :math:`\mathbf{x}`.
.. math::
\sum_{i=1}^N\sin(x_{ij})
Parameters
----------
xs : `np.ndarray`
Matrix :math:`\mathbf{x}`.
**Shape**: :math:`(M,\,N)`.
Returns
-------
s : `np.ndarray`
Column-wise sums of the sines of the elements in the matrix
:math:`\mathbf{x}`.
**Shape**: :math:`(M,)`.
"""
s = np.empty(xs.shape[0])
for i in numba.prange(xs.shape[0]):
s[i] = numba_sine_sum(xs[i])
return s
[docs]
def sine_column_sum_inplace(xs: np.ndarray[float], s: np.ndarray[float]) -> None:
r"""
Evaluates the column-wise sum of the sines of the elements in a
two-dimensional NumPy array :math:`\mathbf{x}`.
.. math::
\sum_{i=1}^N\sin(x_{ij})
Parameters
----------
xs : `np.ndarray`
Matrix :math:`\mathbf{x}`.
**Shape**: :math:`(M,\,N)`.
s : `np.ndarray`
Array to hold the column-wise sums of the sines of the elements
in the matrix :math:`\mathbf{x}`.
**Shape**: :math:`(M,)`.
"""
assert s.shape[0] == xs.shape[0]
for i in numba.prange(xs.shape[0]):
s[i] = numba_sine_sum(xs[i])