# Copyright (c) 2026 Authors and contributors
# (see the AUTHORS.rst file for the full list of names)
#
# Released under the GNU Public Licence, v3 or any higher version
# SPDX-License-Identifier: GPL-3.0-or-later
"""Small helper and utilities functions that don't fit anywhere else."""
import functools
import inspect
import logging
import os
import re
import sys
import warnings
from collections.abc import Callable
from itertools import combinations
from pathlib import Path
from typing import Protocol, TypeAlias
import MDAnalysis as mda
import numpy as np
from mdacli.utils import convert_str_time, split_time_unit
from MDAnalysis.analysis.base import Results
from scipy.signal import find_peaks
from maicos.lib.math import (
combine_subsample_covariance,
combine_subsample_variance,
correlation_time,
)
logger = logging.getLogger(__name__)
#: Type for arguments accepting a file system path.
PathLike: TypeAlias = str | os.PathLike[str]
[docs]
def triclinic_to_orthorhombic(dimensions: np.ndarray) -> np.ndarray:
"""Convert triclinic box dimensions to an orthorhombic bounding box.
Computes the orthorhombic box that corresponds to the GROMACS representation
of a triclinic cell. For orthorhombic input, the box is returned unchanged
(with angles set to 90).
Parameters
----------
dimensions : numpy.ndarray
Box dimensions in MDAnalysis format ``[a, b, c, alpha, beta, gamma]``,
where lengths are in Angstrom and angles in degrees.
Returns
-------
numpy.ndarray
Orthorhombic box dimensions ``[xx, yy, zz, 90, 90, 90]``.
"""
a, b, c, alpha, beta, gamma = dimensions
if np.allclose([alpha, beta, gamma], 90.0):
return np.array([a, b, c, 90.0, 90.0, 90.0])
alpha, beta, gamma = np.radians([alpha, beta, gamma])
xx = a
yy = b * np.sin(gamma)
xz = c * np.cos(beta)
yz = c * (np.cos(alpha) - np.cos(beta) * np.cos(gamma)) / np.sin(gamma)
zz = np.sqrt(c**2 - xz**2 - yz**2)
return np.array([xx, yy, zz, 90.0, 90.0, 90.0])
[docs]
def check_file_extension(filename: PathLike, extension: str) -> str:
"""Ensure ``filename`` ends with ``extension``, appending it with a warning.
If ``filename`` does not end with ``extension``, the extension is appended
and a :class:`UserWarning` is issued.
Parameters
----------
filename : str or os.PathLike
Name of the file to be checked.
extension : str
Expected file extension, including the leading dot (e.g. ``".dat"``).
Returns
-------
str
``filename`` as a string, with ``extension`` appended if necessary.
"""
filename = os.fspath(filename)
if not filename.endswith(extension):
warnings.warn(
f"The file name should have a {extension!r} file extension. The user "
f"requested the file with name {filename!r}, but it will be saved as "
f"{filename + extension!r}.",
stacklevel=2,
)
filename = filename + extension
return filename
DOC_REGEX_PATTERN = re.compile(r"\$\{([^\}]+)\}")
DOC_DICT = {
#####################
# DESCRIPTION SECTION
#####################
"SAVE_METHOD_DESCRIPTION": (
"Save results of analysis to file specified by ``output``."
),
"SAVE_METHOD_PREFIX_DESCRIPTION": r"""Save results of analysis to files specified by
``output_prefix``.""",
"DENSITY_DESCRIPTION_1": r"""Calculations are carried out for ``mass``
:math:`(\rm u \cdot Å^{-3})`, ``number`` :math:`(\rm Å^{-3})`, partial ``charge``
:math:`(\rm e \cdot Å^{-3})` or electron :math:`(\rm e \cdot Å^{-3})` density
profiles """,
"DENSITY_DESCRIPTION_2": """Cell dimensions are allowed to fluctuate in time.
For grouping with respect to ``molecules``, ``residues`` etc., the corresponding
centers (i.e., center of mass), taking into account periodic boundary conditions,
are calculated. For these calculations molecules will be unwrapped/made whole.
Trajectories containing already whole molecules can be run with ``unwrap=False`` to
gain a speedup. For grouping with respect to atoms, the ``unwrap`` option is always
ignored.""",
"PLANAR_DESCRIPTION": """along certain cartesian axes ``[x, y, z]`` of the
simulation cell.""",
"CYLINDRICAL_DESCRIPTION": """along the radial axis in a cylindrical coordinate
system, with principal axis along ``[x, y, z]`` axes of the simulation cell. The
origin of the coordinate system defaults to the box center, but can be set to
dynamically follow the center of mass of a reference group via ``refgroup``""",
"SPHERICAL_DESCRIPTION": """along radial axis in spherical coordinate system. The
origin of the coordinate system defaults to the box center, but can be set to
dynamically follow the center of mass of a reference group via `refgroup`.""",
"DENSITY_PLANAR_DESCRIPTION": r"""${DENSITY_DESCRIPTION_1}
${PLANAR_DESCRIPTION}
${DENSITY_DESCRIPTION_2}""",
"DENSITY_CYLINDER_DESCRIPTION": r"""${DENSITY_DESCRIPTION_1}
${CYLINDRICAL_DESCRIPTION}
${DENSITY_DESCRIPTION_2}""",
"DENSITY_SPHERE_DESCRIPTION": r"""${DENSITY_DESCRIPTION_1}
${SPHERICAL_DESCRIPTION}
${DENSITY_DESCRIPTION_2}""",
"DIPORDER_DESCRIPTION": r"""Calculations include the first- and second-rank
orientational order parameters,
which are the Legendre polynomials :math:`P_1(x) = x` and :math:`P_2(x) = \frac{3
x^2 - 1}{2}` with :math:`x = \cos(θ[z])`.""",
"CORRELATION_INFO": r"""For further information on the correlation analysis please
refer to :class:`AnalysisBase <maicos.core.base.AnalysisBase>` or the
:ref:`general-design` section.""",
"CORRELATION_INFO_PLANAR": r"""For the correlation analysis the central bin
(:math:`N / 2`) of the 0th's group profile is used. ${CORRELATION_INFO}""",
"CORRELATION_INFO_RADIAL": """For the correlation analysis the 0th bin of the 0th's
group profile is used. ${CORRELATION_INFO}""",
"RUN_METHOD_DESCRIPTION": """Iterate over the trajectory.
Parameters
----------
start : int
start frame of analysis
stop : int
stop frame of analysis
step : int
number of frames to skip between each analysed frame
frames : array_like
array of integers or booleans to slice trajectory; ``frames`` can only be
used *instead* of ``start``, ``stop``, and ``step``. Setting *both*
``frames`` and at least one of ``start``, ``stop``, ``step`` to a
non-default value will raise a :exc:`ValueError`.
verbose : bool
Turn on verbosity
progressbar_kwargs : dict
ProgressBar keywords with custom parameters regarding progress bar position,
etc; see :class:`MDAnalysis.lib.log.ProgressBar` for full list.
Returns
-------
self : object
analysis object
""",
##########################
# SINGLE PARAMETER SECTION
##########################
"ATOMGROUP_PARAMETER": """atomgroup : MDAnalysis.core.groups.AtomGroup
A :class:`~MDAnalysis.core.groups.AtomGroup` for which the calculations are
performed.""",
"WRAP_COMPOUND_PARAMETER": """wrap_compound : str
The group which will be kept together through the wrap processes. Allowed values
are: ``"atoms"``, ``"group"``, ``"residues"``, ``"segments"``, ``"molecules"``, or
``"fragments"``.""",
"DENS_PARAMETER": """dens : {``"mass"``, ``"number"``, ``"charge"``, ``"electron"``}
density type to be calculated.""",
"TEMPERATURE_PARAMETER": """temperature : float
Reference temperature (K)""",
"BIN_WIDTH_PARAMETER": """bin_width : float
Width of the bins (in Å).""",
"DIM_PARAMETER": """dim : {0, 1, 2}
Dimension for binning (``x=0``, ``y=1``, ``z=1``).""",
"VDIM_PARAMETER": """vdim : {0, 1, 2}
Dimension for velocity binning (``x=0``, ``y=1``, ``z=1``).""",
"PDIM_PLANAR_PARAMETER": """pdim : {0, 1, 2}
direction of the projection""",
"PDIM_RADIAL_PARAMETER": """pdim : {``"r"``, ``"z"``}
direction of the projection""",
"FLUX_PARAMETER": r"""flux : bool
Calculate the flux (:math:`[Å^2/\mathrm{ps}]`) instead of the velocity.""",
"GROUPING_PARAMETER": """grouping : {``"atoms"``, ``"residues"``, \
``"segments"``, ``"molecules"``, ``"fragments"``}
Atom grouping for the calculations.
The possible grouping options are the atom positions (in the case where
``grouping="atoms"``) or the center of mass of the specified grouping unit (in the
case where ``grouping="residues"``, ``"segments"``, ``"molecules"`` or
``"fragments"``).""",
"OUTPUT_PARAMETER": """output : str or os.PathLike
Output filename.""",
"OUTPUT_PREFIX_PARAMETER": """output_prefix : str or os.PathLike
Prefix for output files.""",
"SYM_PARAMETER": """sym : bool
Symmetrize the profile. Only works in combination with ``refgroup``.""",
"BIN_METHOD_PARAMETER": """bin_method : {``"com"``, ``"cog"``, ``"coc"``}
Method for the position binning.
The possible options are center of mass (``"com"``), center of geometry (``"cog"``),
and center of charge (``"coc"``).""",
"ORDER_PARAMETER_PARAMETER": r"""order_parameter : {``"P1"``, ``"P2"``}
Order parameter to be calculated:
- ``"P1"``: Legendre polynomial :math:`P_1(cos[θ]) = cos[θ]`
- ``"P2"``: Legendre polynomial
:math:`P_2(cos[θ]) = \frac{3cos^2[θ] - 1}{2}`""",
###################################
# MULTI/COMBINES PARAMETERS SECTION
###################################
"BASE_CLASS_PARAMETERS": """refgroup : MDAnalysis.core.groups.AtomGroup
Reference :class:`~MDAnalysis.core.groups.AtomGroup` used for the calculation. If
``refgroup`` is provided, the calculation is performed relative to the center of
mass of the AtomGroup. If ``refgroup`` is :obj:`None` the calculations are performed
with respect to the center of the (changing) box.
unwrap : bool
When :obj:`True`, molecules that are broken due to the periodic boundary conditions
are made whole.
If the input contains molecules that are already whole, speed up the calculation by
disabling unwrap. To do so, use the flag ``-no-unwrap`` when using MAICoS from the
command line, or use ``unwrap=False`` when using MAICoS from the Python interpreter.
Note: Molecules containing virtual sites (e.g. TIP4P water models) are not currently
supported in MDAnalysis. In this case, you need to provide unwrapped trajectory
files directly, and disable unwrap. Trajectories can be unwrapped, for example,
using the ``trjconv`` command of GROMACS.
pack : bool
When :obj:`True`, molecules are put back into the unit cell. This is required
because MAICoS only takes into account molecules that are inside the unit cell.
If the input contains molecules that are already packed, speed up the calculation by
disabling packing with ``pack=False``.
jitter : float
Magnitude of the random noise to add to the atomic positions.
A jitter can be used to stabilize the aliasing effects sometimes appearing when
histogramming data. The jitter value should be about the precision of the
trajectory. In that case, using jitter will not alter the results of the histogram.
If ``jitter = 0.0`` (default), the original atomic positions are kept unchanged.
You can estimate the precision of the positions in your trajectory with
:func:`maicos.lib.util.trajectory_precision`. Note that if the precision is not the
same for all frames, the smallest precision should be used.
concfreq : int
When concfreq (for conclude frequency) is larger than ``0``, the conclude function
is called and the output files are written every ``concfreq`` frames.""",
"PROFILE_CLASS_PARAMETERS_PRIVATE": """weighting_function : callable
The function calculating the array weights for the histogram analysis. It must take
an :py:class:`AtomGroup<MDAnalysis.AtomGroup>` as first argument and a grouping
(``"atoms"``, ``"residues"``, ``"segments"``, ``"molecules"``, ``"fragments"``) as
second. Additional parameters can be given as ``weighting_function_kwargs``. The
function must return a numpy.ndarray with the same length as the number of group
members.
weighting_function_kwargs : dict
Additional keyword arguments for ``weighting_function``
normalization : {``"none"``, ``"number"``, ``"volume"``}
The normalization of the profile performed in every frame. If :obj:`None`, no
normalization is performed. If `number`, the histogram is divided by the number of
occurences in each bin. If `volume`, the profile is divided by the volume of each
bin.""",
"PLANAR_Z_PARAMETERS": """zmin : float
Minimal coordinate for evaluation (in Å) with respect to the center of mass of
the refgroup.
If ``zmin=None``, all coordinates down to the lower cell boundary are taken into
account.
zmax : float
Maximal coordinate for evaluation (in Å) with respect to the center of mass of the
refgroup.
If ``zmax = None``, all coordinates up to the upper cell boundary are taken into
account.""",
"PLANAR_CLASS_PARAMETERS": """${DIM_PARAMETER}
${PLANAR_Z_PARAMETERS}
${BIN_WIDTH_PARAMETER}""",
"RADIAL_CLASS_PARAMETERS": """rmin : float
Minimal radial coordinate relative to the center of mass of the refgroup for
evaluation (in Å).
rmax : float
Maximal radial coordinate relative to the center of mass of the refgroup for
evaluation (in Å).
If ``rmax=None``, the box extension is taken.""",
"PDF_PARAMETERS": """g1 : MDAnalysis.core.groups.AtomGroup
First AtomGroup.
g2 : MDAnalysis.core.groups.AtomGroup
Second AtomGroup.""",
"PROFILE_CLASS_PARAMETERS": """${BIN_METHOD_PARAMETER}
${GROUPING_PARAMETER}""",
"CYLINDER_CLASS_PARAMETERS": """${DIM_PARAMETER}
${PLANAR_Z_PARAMETERS}
${RADIAL_CLASS_PARAMETERS}
${BIN_WIDTH_PARAMETER}""",
"SPHERE_CLASS_PARAMETERS": """${RADIAL_CLASS_PARAMETERS}
${BIN_WIDTH_PARAMETER}""",
"PROFILE_PLANAR_CLASS_PARAMETERS": """
${PLANAR_CLASS_PARAMETERS}
${PROFILE_CLASS_PARAMETERS}
${SYM_PARAMETER}
${BASE_CLASS_PARAMETERS}""",
"PROFILE_CYLINDER_CLASS_PARAMETERS": """${CYLINDER_CLASS_PARAMETERS}
${PROFILE_CLASS_PARAMETERS}
${BASE_CLASS_PARAMETERS}""",
"PROFILE_SPHERE_CLASS_PARAMETERS": """${SPHERE_CLASS_PARAMETERS}
${PROFILE_CLASS_PARAMETERS}
${BASE_CLASS_PARAMETERS}""",
###################
# ATTRIBUTE SECTION
###################
"PLANAR_CLASS_ATTRIBUTES": """results.bin_pos : numpy.ndarray
Bin positions (in Å) ranging from ``zmin`` to ``zmax``.""",
"RADIAL_CLASS_ATTRIBUTES": """results.bin_pos : numpy.ndarray
Bin positions (in Å) ranging from ``rmin`` to ``rmax``.""",
"PROFILE_CLASS_ATTRIBUTES": """results.profile : numpy.ndarray
Calculated profile.
results.dprofile : numpy.ndarray
Estimated profile's uncertainity.""",
"CYLINDER_CLASS_ATTRIBUTES": "${RADIAL_CLASS_ATTRIBUTES}",
"SPHERE_CLASS_ATTRIBUTES": "${RADIAL_CLASS_ATTRIBUTES}",
"PROFILE_PLANAR_CLASS_ATTRIBUTES": """${PLANAR_CLASS_ATTRIBUTES}
${PROFILE_CLASS_ATTRIBUTES}""",
"PROFILE_CYLINDER_CLASS_ATTRIBUTES": """${RADIAL_CLASS_ATTRIBUTES}
${PROFILE_CLASS_ATTRIBUTES}""",
"PROFILE_SPHERE_CLASS_ATTRIBUTES": """${RADIAL_CLASS_ATTRIBUTES}
${PROFILE_CLASS_ATTRIBUTES}""",
}
"""Dictionary containing the keys and the actual docstring used by :func:`maicos.lib.util.render_docs`.
:meta hide-value:
""" # noqa: E501
def _render_docs(func: Callable, doc_dict: dict = DOC_DICT) -> Callable:
if func.__doc__ is not None:
while True:
keys = DOC_REGEX_PATTERN.findall(func.__doc__)
if not keys:
break # Exit the loop if no more patterns are found
for key in keys:
func.__doc__ = func.__doc__.replace(f"${{{key}}}", doc_dict[key])
return func
[docs]
def render_docs(func: Callable) -> Callable:
"""Replace all template phrases in the functions docstring.
Keys for the replacement are taken from in :attr:`maicos.lib.util.DOC_DICT`.
Parameters
----------
func : callable
The callable (function, class) where the phrase old should be replaced.
Returns
-------
Callable
callable with replaced phrase
"""
return _render_docs(func, doc_dict=DOC_DICT)
[docs]
def correlation_analysis(timeseries: np.ndarray) -> float:
"""Timeseries correlation analysis.
Analyses a timeseries for correlation and prints a warning if the correlation time
is larger than the step size.
Parameters
----------
timeseries : numpy.ndarray
Array of (possibly) correlated data.
Returns
-------
corrtime: float
Estimated correlation time of ``timeseries``.
"""
if np.any(np.isnan(timeseries)):
# Fail silently if there are NaNs in the timeseries. This is the case if the
# feature is not implemented for the given analysis. It could also be because of
# a bug, but that is not our business.
return -1
if len(timeseries) <= 4:
warnings.warn(
"Your trajectory is too short to estimate a correlation time. Use the "
"calculated error estimates with caution.",
stacklevel=2,
)
return -1
corrtime = correlation_time(timeseries)
if corrtime == -1:
warnings.warn(
"Your trajectory does not provide sufficient statistics to estimate a "
"correlation time. Use the calculated error estimates with caution.",
stacklevel=2,
)
if corrtime > 0.5:
warnings.warn(
"Your data seems to be correlated with a correlation time which is "
f"{corrtime + 1:.2f} times larger than your step size. Consider increasing "
f"your step size by a factor of {int(np.ceil(2 * corrtime + 1)):d} to get "
"a reasonable error estimate.",
stacklevel=2,
)
return corrtime
[docs]
@render_docs
def get_compound(atomgroup: mda.AtomGroup) -> str:
"""Returns the highest order topology attribute.
The order is "molecules", "fragments", "residues". If the topology contains none of
those attributes, an AttributeError is raised.
Parameters
----------
${ATOMGROUP_PARAMETER}
Returns
-------
str
Name of the topology attribute.
Raises
------
AttributeError
``atomgroup`` is missing any connection information"
"""
if hasattr(atomgroup, "molnums"):
return "molecules"
if hasattr(atomgroup, "fragments"):
logger.info("Cannot use 'molecules'. Falling back to 'fragments'")
return "fragments"
if hasattr(atomgroup, "residues"):
logger.info("Cannot use 'fragments'. Falling back to 'residues'")
return "residues"
raise AttributeError("Missing any connection information in `atomgroup`.")
[docs]
def charge_neutral(filter: str) -> Callable:
"""Raise a Warning when AtomGroup is not charge neutral.
Class Decorator to raise an Error/Warning when AtomGroup in an AnalysisBase class is
not charge neutral. The behaviour of the warning can be controlled with the filter
attribute. If the AtomGroup's corresponding universe is non-neutral an ValueError is
raised.
Parameters
----------
filter : str
Filter type to control warning filter. Common values are: "error" or "default"
See :func:`warnings.simplefilter` for more options.
"""
def inner(original_class):
def charge_check(function):
@functools.wraps(function)
def wrapped(self):
if not np.allclose(
self.atomgroup.total_charge(compound=get_compound(self.atomgroup)),
0,
atol=1e-4,
):
with warnings.catch_warnings():
warnings.simplefilter(filter)
warnings.warn(
"At least one AtomGroup has free charges. Analysis for "
"systems with free charges could lead to severe "
"artifacts!",
stacklevel=1,
)
if not np.allclose(
self.atomgroup.universe.atoms.total_charge(), 0, atol=1e-4
):
raise ValueError(
"Analysis for non-neutral systems is not supported."
)
return function(self)
return wrapped
original_class._prepare = charge_check(original_class._prepare)
return original_class
return inner
[docs]
def unwrap_refgroup(original_class):
"""Class decorator error if ``unwrap = False`` and ``refgroup != None``."""
def unwrap_check(function):
@functools.wraps(function)
def unwrap_check(self):
if (
hasattr(self, "unwrap")
and hasattr(self, "refgroup")
and not self.unwrap
and self.refgroup is not None
):
raise ValueError(
"Analysis using `unwrap=False` and `refgroup != None` can lead "
"to broken molecules and severe errors."
)
return function(self)
return unwrap_check
original_class._prepare = unwrap_check(original_class._prepare)
return original_class
[docs]
def trajectory_precision(
trajectory: mda.coordinates.base.ReaderBase, dim: int = 2
) -> np.ndarray:
"""Detect the precision of a trajectory.
Parameters
----------
trajectory : MDAnalysis.coordinates.base.ReaderBase
Trajectory from which the precision is detected.
dim : {2, 0, 1}
Dimension along which the precision is detected.
Returns
-------
precision : numpy.ndarray
Precision of each frame of the trajectory.
If the trajectory has a high precision, its resolution will not be detected, and
a value of 1e-4 is returned.
"""
# The threshold will limit the precision of the detection. Using a value that is too
# low will end up costing a lot of memory. 1e-4 is enough to safely detect the
# resolution of format like XTC
threshold_bin_width = 1e-4
precision = np.zeros(trajectory.n_frames)
# to be done, add range=(0, -1, 1) parameter for ts in
# trajectory[range[0]:range[1]:range[2]]:
for ts in trajectory:
n_bins = int(
np.ceil(
(
np.max(trajectory.ts.positions[:, dim])
- np.min(trajectory.ts.positions[:, dim])
)
/ threshold_bin_width
)
)
hist1, z = np.histogram(trajectory.ts.positions[:, dim], bins=n_bins)
(
hist2,
bin_edges,
) = np.histogram(np.diff(z[np.where(hist1)]), bins=1000, range=(0, 0.1))
if len(find_peaks(hist2)[0]) == 0 or bin_edges[find_peaks(hist2)[0][0]] <= 5e-4:
precision[ts.frame] = 1e-4
else:
precision[ts.frame] = bin_edges[find_peaks(hist2)[0][0]]
return precision
DOI_LIST = {
"10.1103/PhysRevLett.117.048001": "Schlaich, A. et al., Phys. Rev. Lett. 117, "
"(2016).",
"10.1021/acs.jpcb.9b09269": "Loche, P. et al., J. Phys. Chem. B 123, (2019).",
"10.1021/acs.jpca.0c04063": "Carlson, S. et al., J. Phys. Chem. A 124, (2020).",
"10.1103/PhysRevE.92.032718": "Schaaf, C. et al., Phys. Rev. E 92, (2015).",
}
"""References associated with MAICoS
:meta hide-value:
"""
[docs]
def citation_reminder(*dois: str) -> str:
"""Prints citations in order to remind users to give due credit.
Parameters
----------
dois : list
dois associated with the method which calls this. Possible dois are registered
in :attr:`maicos.lib.util.DOI_LIST`.
Returns
-------
cite : str
formatted citation reminders
"""
cite = ""
for doi in dois:
lines = [
"If you use this module in your work, please read and cite:",
DOI_LIST[doi],
f"doi: {doi}",
]
plus = f"{max([len(i) for i in lines]) * '+'}"
lines.insert(0, f"\n{plus}")
lines.append(f"{plus}\n")
cite += "\n".join(lines)
return cite
[docs]
@render_docs
def get_center(atomgroup: mda.AtomGroup, bin_method: str, compound: str) -> np.ndarray:
"""Center attribute for an :class:`MDAnalysis.core.groups.AtomGroup`.
This function acts as a wrapper for the
:meth:`MDAnalysis.core.groups.AtomGroup.center` method, providing a more
user-friendly interface by automatically determining the appropriate weights based
on the chosen binning method.
Parameters
----------
${ATOMGROUP_PARAMETER}
${BIN_METHOD_PARAMETER}
compound : {``"group"``, ``"segments"``, ``"residues"``, ``"molecules"``, ``"fragments"``}
The compound to be used in the center calculation. For example, ``"residue"``,
``"segment"``, etc.
Returns
-------
np.ndarray
The coordinates of the calculated center.
Raises
------
ValueError
If the provided ``bin_method`` is not one of {``"com"``, ``"cog"``, ``"coc"``}.
""" # noqa: E501
if bin_method == "cog":
weights = None
elif bin_method == "com":
weights = atomgroup.masses
elif bin_method == "coc":
weights = atomgroup.charges.__abs__()
else:
raise ValueError(
f"'{bin_method}' is an unknown binning method. Use 'cog', 'com' or 'coc'."
)
return atomgroup.center(weights=weights, compound=compound)
[docs]
@render_docs
def unit_vectors_planar(
atomgroup: mda.AtomGroup, # noqa: ARG001
grouping: str, # noqa: ARG001
pdim: int,
) -> np.ndarray:
"""Calculate unit vectors in planar geometry.
Parameters
----------
${ATOMGROUP_PARAMETER}
${GROUPING_PARAMETER}
${PDIM_PLANAR_PARAMETER}
Returns
-------
numpy.ndarray
the unit vector
"""
unit_vectors = np.zeros(3)
unit_vectors[pdim] += 1
return unit_vectors
[docs]
@render_docs
def unit_vectors_cylinder(
atomgroup: mda.AtomGroup,
grouping: str,
bin_method: str,
dim: int,
pdim: str,
) -> np.ndarray:
"""Calculate cylindrical unit vectors in cartesian coordinates.
Parameters
----------
${ATOMGROUP_PARAMETER}
${GROUPING_PARAMETER}
${BIN_METHOD_PARAMETER}
${DIM_PARAMETER}
${PDIM_RADIAL_PARAMETER}
Returns
-------
numpy.ndarray
Array of the calculated unit vectors with shape (3,) for ``pdim='z'`` and shape
(3,n) for ``pdim='r'``. The length of ``n`` depends on the grouping.
"""
# We do NOT transform ``unit_vectors`` into cylindrical coordinates, because all
# scalar products in ``dipolar_weights`` will be performed cartesian coordinates!
if pdim == "r":
unit_vectors = get_center(
atomgroup=atomgroup, bin_method=bin_method, compound=grouping
)
unit_vectors -= atomgroup.universe.dimensions[:3] / 2
# set z direction to zero. r in cylindrical coordinates contains only x and y.
unit_vectors[:, dim] = 0
unit_vectors /= np.linalg.norm(unit_vectors, axis=1)[:, np.newaxis]
elif pdim == "z":
unit_vectors = np.zeros(3)
unit_vectors[dim] += 1
else:
raise ValueError(
f"'{pdim}' is an unknown direction for the projection. Use 'r' or 'z'."
)
return unit_vectors
[docs]
@render_docs
def unit_vectors_sphere(
atomgroup: mda.AtomGroup, grouping: str, bin_method: str
) -> np.ndarray:
"""Calculate spherical unit vectors in cartesian coordinates.
Parameters
----------
${ATOMGROUP_PARAMETER}
${GROUPING_PARAMETER}
${BIN_METHOD_PARAMETER}
Returns
-------
numpy.ndarray
Array of the calculated unit vectors with shape (3,n). The length of ``n``
depends on the grouping.
"""
# We do NOT transform ``unit_vectors`` into spherical coordinates, because all
# scalar products in ``dipolar_weights`` will be performed cartesian coordinates!
unit_vectors = get_center(
atomgroup=atomgroup, bin_method=bin_method, compound=grouping
)
# shift origin to box center and afterwards normalize
unit_vectors -= atomgroup.universe.dimensions[:3] / 2
unit_vectors /= np.linalg.norm(unit_vectors, axis=1)[:, np.newaxis]
return unit_vectors
[docs]
def maicos_banner(version: str = "", frame_char: str = "-") -> str:
"""Prints ASCII banner resembling the MAICoS Logo with 80 chars width.
Parameters
----------
version : str
Version string to add to the banner.
frame_char : str
Character used to as framing around the banner.
Returns
-------
banner : str
formatted banner
"""
banner = r"""
@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@ __ __ _____ _____ _____ @
@ ()----() | \/ | /\ |_ _| / ____| / ____| @
@ / | \ | \ / | / \ | | | | ___ | (___ @
@ () ||| | () | |\/| | / /\ \ | | | | / _ \ \___ \ @
@ \ |||||_ / | | | | / ____ \ _| |_ | |____ | (_) | ____) | @"""
if len(version) < 8:
banner += rf"""
@ ()----() |_| |_| /_/ \_\ |_____| \_____| \___/ |_____/ {version:^8}@
@ @
@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
"""
else:
banner += rf"""
@ ()----() |_| |_| /_/ \_\ |_____| \_____| \___/ |_____/ @
@ @
@{version:^78}@
@ @
@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
"""
return banner.replace("@", frame_char)
[docs]
class Unit_vector(Protocol):
"""Protocol class for unit vector methods type hints."""
def __call__(self, atomgroup: mda.AtomGroup, grouping: str) -> np.ndarray:
"""Call for type hints."""
...
[docs]
def make_pair_key(key_i: str, key_j: str) -> tuple[str, str]:
"""Return a canonical pair key for two observable keys.
The key is ordered so that ``make_pair_key(a, b) == make_pair_key(b, a)``,
guaranteeing deterministic lookups regardless of argument order.
Parameters
----------
key_i, key_j : str
Keys of the two observables.
Returns
-------
tuple[str, str]
Canonical ``(smaller, larger)`` pair key.
"""
return (key_i, key_j) if key_i <= key_j else (key_j, key_i)
[docs]
def times_to_frames(start: str, stop: str, step: str, dt: float) -> dict:
"""Convert 'start', 'stop', 'step' into their frame number based on given dt.
Parameters
----------
start : str
start in a '12ps' format
stop : str
stop in a '12ps' format
step : str
step in a '12ps' format
dt : float
the time step in ps
Returns
-------
dict
Dictionnary containing start, stop and step frames
Example
-------
>>> times_to_frames("2ps", "12ps", "2ps", 1)
{'start': 2, 'stop': 12, 'step': 2}
>>> times_to_frames("0ns", "2ns", "10ps", 0.01)
{'start': 0, 'stop': 200000, 'step': 1000}
"""
val, unit = split_time_unit(step)
val_in_ps = mda.units.convert(val, unit, "ps")
temp_step = round(val_in_ps / dt)
if not (np.isclose(temp_step * dt, val_in_ps)):
raise ValueError("step should be a multiple of dt.")
temp_start = convert_str_time(start, dt)
temp_stop = convert_str_time(stop, dt)
timedict = {}
timedict["start"] = temp_start
timedict["stop"] = temp_stop
timedict["step"] = convert_str_time(step, dt)
if not (temp_start < temp_stop):
warnings.warn("'start' should be < to 'stop - step'.", stacklevel=1)
return timedict
return timedict
def _joint_pop(pop_x: np.ndarray, pop_y: np.ndarray) -> np.ndarray:
"""Shared sample count of two co-sampled observables.
Raises a ``ValueError`` if the populations do not broadcast or differ, since a
covariance is only defined for co-sampled observables.
"""
b_x, b_y = np.broadcast_arrays(pop_x, pop_y)
if not np.array_equal(b_x, b_y):
raise ValueError(
"cannot sample the covariance of observables with different populations"
)
return b_x
#: Observable types the accumulator knows how to merge.
_COMPATIBLE_TYPES = (
np.ndarray,
float,
int,
list,
tuple,
np.float32,
np.float64,
np.int32,
np.int64,
)
[docs]
class MomentAccumulator:
"""Streaming mean, variance and covariance of per-frame observables.
This is the statistics backend of :class:`~maicos.core.AnalysisBase`. Each
frame's observables are merged into running statistics with Welford's online
algorithm, so no per-frame data has to be stored. Observables may be scalars
or arrays (e.g. profiles), in which case every element is accumulated
independently. A frame can report a single sample per observable or, via the
population / variance / covariance containers, the mean of several samples.
After the first frame is passed to :meth:`register`, every following frame is
merged with :meth:`update`. The running statistics are exposed as
:class:`~MDAnalysis.analysis.base.Results` containers keyed by observable name:
- ``means``: mean of each observable
- ``sems``: standard error of the mean
- ``sums``: sum of the samples
- ``pop``: number of samples
- ``M2``: sum of squared deviations from the mean
- ``C``: co-moment of each requested pair, keyed by
:func:`~maicos.lib.util.make_pair_key`
:meth:`cov` returns the covariance of two observable means and
:meth:`propagate_error` the standard error of a function of several
observables, including their cross-covariances.
Parameters
----------
requested_pairs : iterable of tuple of str
Canonical observable pairs (see :func:`~maicos.lib.util.make_pair_key`)
whose off-diagonal covariance should be accumulated. The observables of
each pair must broadcast against each other and be co-sampled (same
population). Only the listed pairs are tracked. Empty (the default)
disables covariance entirely.
Examples
--------
Accumulate two correlated scalar observables over four frames
>>> import numpy as np
>>> from MDAnalysis.analysis.base import Results
>>> from maicos.lib.util import MomentAccumulator, make_pair_key
>>> acc = MomentAccumulator({make_pair_key("x", "y")})
>>> acc.register(Results(x=0.0, y=0.0), Results(), Results(), Results())
>>> for x, y in [(0.0, 1.0), (0.0, 1.0), (6.0, 6.0)]:
... acc.update(Results(x=x, y=y), Results(), Results(), Results())
>>> print(acc.means.x, acc.means.y)
1.5 2.0
Covariance of the means and the error of ``f = x - y``, which is smaller than
the uncorrelated estimate because ``x`` and ``y`` are positively correlated
>>> print(acc.cov("x", "y"))
1.5
>>> print(acc.propagate_error({"x": 1.0, "y": -1.0}))
0.25
>>> print(np.sqrt(acc.sems.x**2 + acc.sems.y**2))
1.75
"""
def __init__(self, requested_pairs=()):
self.means = Results() # mean of the observables across frames
self.sems = Results() # standard error of the mean across frames
self.sums = Results() # sum of the observables across frames
self.pop = Results() # count of samples across frames
self.M2 = Results() # sum of squared deviations from the mean
self.C = Results() # off-diagonal co-moments, keyed (i, j)
self._requested_cov_pairs = set(requested_pairs) # pair keys tracked in C
self._keys = [] # observable keys, in first-seen order
def _sanitize(self, obs, _pop, _var):
"""Return per-frame observable / population / variance arrays.
Lists and scalars become float arrays (scalars stay 0-d), and a missing
population / variance defaults to a single sample with zero within-frame
variance. The input ``Results`` containers are not mutated.
"""
s_obs, s_pop, s_var = {}, {}, {}
for key in obs:
value = obs[key]
if not isinstance(value, _COMPATIBLE_TYPES):
raise TypeError(f"Observable {key!r} has an incompatible type.")
s_obs[key] = np.asarray(value, dtype=float)
if key in _pop:
s_pop[key] = np.asarray(_pop[key])
s_var[key] = np.asarray(_var[key], dtype=float)
else:
s_pop[key] = np.ones(s_obs[key].shape, dtype=int)
s_var[key] = np.zeros(s_obs[key].shape)
return s_obs, s_pop, s_var
[docs]
def register(self, obs, _pop, _var, _cov):
"""Seed the running containers from the first frame's observables.
Parameters
----------
obs : MDAnalysis.analysis.base.Results
The current frame's observables, keyed by observable name.
_pop : MDAnalysis.analysis.base.Results
The current frame's sample count for each observable. Missing entries
default to a unit population (the observable is a single sample).
_var : MDAnalysis.analysis.base.Results
The current frame's within-frame variance for each observable.
_cov : MDAnalysis.analysis.base.Results
The current frame's within-frame covariance, keyed by canonical
observable pair (see :func:`maicos.lib.util.make_pair_key`).
"""
s_obs, s_pop, s_var = self._sanitize(obs, _pop, _var)
for key in s_obs:
# np.array keeps scalars as writable 0-d arrays for the in-place merge.
self.means[key] = np.array(s_obs[key], dtype=float)
self.pop[key] = np.array(s_pop[key])
self.M2[key] = np.array(s_var[key] * s_pop[key], dtype=float)
with np.errstate(divide="ignore", invalid="ignore"):
self.sems[key] = np.array(
np.sqrt(self.M2[key] / self.pop[key] ** 2), dtype=float
)
# Empty (NaN) entries contribute nothing, as in :meth:`_merge_var`.
obs = s_obs[key]
self.sums[key] = np.array(
np.where(np.isnan(obs), 0.0, obs) * s_pop[key], dtype=float
)
self._keys.append(key)
self._build_pairs(s_obs, s_pop, _cov)
def _build_pairs(self, s_obs, s_pop, _cov):
"""Seed ``C`` for the requested pairs; raise if a pair is not co-sampled."""
for pair_key in self._requested_cov_pairs:
key_i, key_j = pair_key
for key in pair_key:
if key not in s_obs:
raise KeyError(
f"requested covariance pair {pair_key} references "
f"unknown observable {key!r}; available: {list(s_obs)}"
)
try:
shape = np.broadcast_shapes(s_obs[key_i].shape, s_obs[key_j].shape)
pop = _joint_pop(s_pop[key_i], s_pop[key_j])
except ValueError as err:
raise ValueError(
f"requested covariance pair {pair_key} is invalid: the "
f"observables must broadcast against each other (shapes "
f"{s_obs[key_i].shape} and {s_obs[key_j].shape}) and be "
f"co-sampled (same population)"
) from err
# Seed with the first frame's within-frame co-moment (zero for
# single-sample observables, where _cov is absent).
cov = np.asarray(_cov.get(pair_key, 0.0), dtype=float)
seed = np.where(np.isnan(cov), 0.0, cov) * pop
self.C[pair_key] = np.broadcast_to(seed, shape).astype(float).copy()
[docs]
def update(self, obs, _pop, _var, _cov):
"""Welford merge of the current frame into all running statistics.
Covariance is merged first, while the running means still hold their
pre-frame values; the variance merge advances the means afterwards.
Parameters
----------
obs : MDAnalysis.analysis.base.Results
The current frame's observables, keyed by observable name.
_pop : MDAnalysis.analysis.base.Results
The current frame's sample count for each observable. Missing entries
default to a unit population (the observable is a single sample).
_var : MDAnalysis.analysis.base.Results
The current frame's within-frame variance for each observable.
_cov : MDAnalysis.analysis.base.Results
The current frame's within-frame covariance, keyed by canonical
observable pair (see :func:`maicos.lib.util.make_pair_key`).
"""
s_obs, s_pop, s_var = self._sanitize(obs, _pop, _var)
for pair_key in self.C:
self._merge_cov(pair_key, s_obs, s_pop, _cov)
for key in self._keys:
self._merge_var(key, s_obs, s_pop, s_var)
def _merge_var(self, key, s_obs, s_pop, s_var):
"""Merge the current frame into one observable's mean / variance."""
mean, M2, pop = self.means[key], self.M2[key], self.pop[key]
pop[...], mean[...], M2[...] = combine_subsample_variance(
s_pop[key], pop, s_obs[key], mean, s_var[key] * s_pop[key], M2
)
with np.errstate(divide="ignore", invalid="ignore"):
self.sems[key][...] = np.sqrt(M2 / pop**2)
obs = s_obs[key]
self.sums[key][...] += np.where(np.isnan(obs), 0.0, obs) * s_pop[key]
def _merge_cov(self, pair_key, s_obs, s_pop, _cov):
"""Merge the current frame into one pair's running co-moment.
Inputs of different shapes (e.g. a scalar paired with a profile) broadcast
against the running co-moment ``C``.
"""
key_i, key_j = pair_key
C = self.C[pair_key]
n_new = _joint_pop(s_pop[key_i], s_pop[key_j])
cov = np.asarray(_cov.get(pair_key, 0.0), dtype=float)
within = np.where(np.isnan(cov), 0.0, cov) * n_new
_, C[...] = combine_subsample_covariance(
n_new,
self.pop[key_i],
s_obs[key_i],
self.means[key_i],
s_obs[key_j],
self.means[key_j],
within,
C,
)
[docs]
def cov(self, key_i: str, key_j: str) -> np.ndarray:
r"""Covariance of the means of two observables.
The element-wise covariance :math:`\mathrm{Cov}(\bar x_i, \bar x_j)` of the
observable means accumulated across frames. The diagonal (``key_i == key_j``)
equals the squared standard error of the mean, :attr:`sems`.
Parameters
----------
key_i, key_j : str
Keys of the two observables.
Returns
-------
numpy.ndarray
Covariance of the means of ``key_i`` and ``key_j``.
Raises
------
KeyError
If the off-diagonal pair was not requested in ``_compute_covariance``.
"""
if key_i == key_j:
return self.sems[key_i] ** 2
if not self._requested_cov_pairs:
raise RuntimeError(
"Covariance tracking is disabled. List the observable pairs in the "
"`_compute_covariance` class attribute to use `cov`/`propagate_error`."
)
pair_key = make_pair_key(key_i, key_j)
if pair_key not in self.C:
raise KeyError(
f"covariance of {key_i!r} and {key_j!r} not tracked: the pair was not "
f"requested in `_compute_covariance`"
)
pop = _joint_pop(self.pop[key_i], self.pop[key_j])
with np.errstate(divide="ignore", invalid="ignore"):
return self.C[pair_key] / pop**2
[docs]
def propagate_error(self, grads: dict) -> np.ndarray:
r"""Propagate observable errors through an estimator.
Computes the standard error of an estimator :math:`f` from the full
covariance of the observable means,
.. math::
\sigma_f^2 = \sum_{ij}
\frac{\partial f}{\partial x_i}
\frac{\partial f}{\partial x_j}
\mathrm{Cov}(\bar x_i, \bar x_j),
where ``grads[key]`` provides :math:`\partial f / \partial x_{key}`. The
diagonal terms reproduce the independent-variable (uncorrelated) estimate;
the off-diagonal terms add the cross-covariance contributions.
Parameters
----------
grads : dict
Mapping of observable key to the gradient of the estimator with
respect to that observable's mean.
Returns
-------
numpy.ndarray
Standard error of the estimator.
Raises
------
KeyError
If two of the supplied observables have no tracked covariance (see
:meth:`cov`).
"""
keys = list(grads)
var_cross = 0.0
for key_i, key_j in combinations(keys, 2):
cov_ij = self.cov(key_i, key_j)
var_cross = var_cross + 2 * grads[key_i] * grads[key_j] * cov_ij
var_diag = 0.0
for key in keys:
var_diag = var_diag + grads[key] ** 2 * self.sems[key] ** 2
# Where the variance went negative, fall back to the uncorrelated
# estimate. Undefined (NaN) elements, e.g. empty bins, stay NaN.
var = var_diag + var_cross
return np.sqrt(np.where(var < 0, var_diag, var))