#!/usr/bin/env python
# -*- Mode: python; tab-width: 4; indent-tabs-mode:nil; coding:utf-8 -*-
#
# Copyright (c) 2024 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
"""Base class for planar analysis."""
import logging
from typing import Callable, Dict, Optional, Union
import MDAnalysis as mda
import numpy as np
from ..lib.math import symmetrize
from ..lib.util import render_docs
from .base import AnalysisBase, ProfileBase
logger = logging.getLogger(__name__)
[docs]
@render_docs
class PlanarBase(AnalysisBase):
r"""Analysis class providing options and attributes for a planar system.
Parameters
----------
${ATOMGROUP_PARAMETER}
${PLANAR_CLASS_PARAMETERS}
${WRAP_COMPOUND_PARAMETER}
Attributes
----------
${PLANAR_CLASS_ATTRIBUTES}
zmin : float
Minimal coordinate for evaluation (Å) with in the lab frame, where 0
corresponds to the origin of the cell.
zmax : float
Maximal coordinate for evaluation (Å) with in the lab frame, where 0
corresponds to the origin of the cell.
_obs.L : float
Average length (in Å) along the chosen dimension in the current frame.
_obs.bin_pos : numpy.ndarray, (n_bins)
Central bin positions (in Å) of each bin (in Å) in the current frame.
_obs.bin_width : float
Bin width (in Å) in the current frame
_obs.bin_edges : numpy.ndarray, (n_bins + 1)
Edges of the bins (in Å) in the current frame.
_obs.bin_area : numpy.ndarray, (n_bins)
Area of the rectangle of each bin in the current frame. Calculated via
:math:`L_x \cdot L_y / N_\mathrm{bins}` where :math:`L_x` and :math:`L_y` are
the box lengths perpendicular to the dimension of evaluations given by `dim`.
:math:`N_\mathrm{bins}` is the number of bins.
results.bin_volume : numpy.ndarray, (n_bins)
Volume of an cuboid of each bin (in Å^3) in the current frame.
"""
def __init__(
self,
atomgroup: mda.AtomGroup,
unwrap: bool,
refgroup: Optional[mda.AtomGroup],
jitter: float,
concfreq: int,
dim: int,
zmin: Union[None, float],
zmax: Union[None, float],
bin_width: float,
wrap_compound: str,
):
super().__init__(
atomgroup=atomgroup,
unwrap=unwrap,
refgroup=refgroup,
jitter=jitter,
concfreq=concfreq,
wrap_compound=wrap_compound,
)
if dim not in [0, 1, 2]:
raise ValueError("Dimension can only be x=0, y=1 or z=2.")
else:
self.dim = dim
# These values are requested by the user, but the actual ones are calculated
# during runtime in the lab frame
self._zmax = zmax
self._zmin = zmin
self._bin_width = bin_width
@property
def odims(self) -> np.ndarray:
"""Other dimensions perpendicular to dim i.e. (0,2) if dim = 1."""
return np.roll(np.arange(3), -self.dim)[1:]
def _compute_lab_frame_planar(self):
"""Compute lab limits `zmin` and `zmax`."""
if self._zmin is None:
self.zmin = 0
else:
self.zmin = self.box_center[self.dim] + self._zmin
if self._zmax is None:
self.zmax = self._universe.dimensions[self.dim]
else:
self.zmax = self.box_center[self.dim] + self._zmax
def _prepare(self):
"""Prepare the planar analysis."""
self._compute_lab_frame_planar()
# TODO: There are much more wrong combinations of zmin and zmax...
if (
self._zmax is not None
and self._zmin is not None
and self._zmax <= self._zmin
):
raise ValueError("`zmax` can not be smaller or equal than `zmin`!")
try:
if self._bin_width > 0:
L = self.zmax - self.zmin
self.n_bins = int(np.ceil(L / self._bin_width))
else:
raise ValueError("Binwidth must be a positive number.")
except TypeError:
raise ValueError("Binwidth must be a number.")
def _single_frame(self):
"""Single frame for the planar analysis."""
self._compute_lab_frame_planar()
self._obs.L = self.zmax - self.zmin
self._obs.box_center = self.box_center
self._obs.bin_edges = np.linspace(self.zmin, self.zmax, self.n_bins + 1)
self._obs.bin_width = self._obs.L / self.n_bins
self._obs.bin_pos = self._obs.bin_edges[1:] - self._obs.bin_width / 2
# We define `bin_area` and `bin_volume` as array of length `n_bins` even though
# each element has the same value. With this the array shape is consistent with
# the cylindrical and spherical classes, where `bin_area` and `bin_volume` is
# different in each bin.
self._obs.bin_area = np.ones(self.n_bins) * np.prod(
self._universe.dimensions[self.odims]
)
self._obs.bin_volume = self._obs.bin_area * self._obs.bin_width
def _conclude(self):
"""Results calculations for the planar analysis."""
# Convert coordinates back from lab frame to refgroup frame.
self.results.bin_pos = self.means.bin_pos - self.means.box_center[self.dim]
[docs]
@render_docs
class ProfilePlanarBase(PlanarBase, ProfileBase):
"""Base class for computing profiles in a cartesian geometry.
${CORRELATION_INFO_RADIAL}
Parameters
----------
${PROFILE_PLANAR_CLASS_PARAMETERS}
${PROFILE_CLASS_PARAMETERS_PRIVATE}
Attributes
----------
${PROFILE_PLANAR_CLASS_ATTRIBUTES}
"""
def __init__(
self,
atomgroup: mda.AtomGroup,
unwrap: bool,
refgroup: Optional[mda.AtomGroup],
jitter: float,
concfreq: int,
dim: int,
zmin: Union[None, float],
zmax: Union[None, float],
bin_width: float,
sym: bool,
grouping: str,
bin_method: str,
output: str,
weighting_function: Callable,
weighting_function_kwargs: Union[None, Dict],
normalization: str,
):
PlanarBase.__init__(
self,
atomgroup=atomgroup,
unwrap=unwrap,
refgroup=refgroup,
jitter=jitter,
concfreq=concfreq,
dim=dim,
zmin=zmin,
zmax=zmax,
bin_width=bin_width,
wrap_compound=grouping, # same as grouping to avoid broken compounds
)
# `AnalysisBase` performs conversions on `atomgroup`.
# Take converted `atomgroup` and not the user provided ones.
ProfileBase.__init__(
self,
atomgroup=self.atomgroup,
grouping=grouping,
bin_method=bin_method,
output=output,
weighting_function=weighting_function,
weighting_function_kwargs=weighting_function_kwargs,
normalization=normalization,
)
self.sym = sym
def _prepare(self):
PlanarBase._prepare(self)
ProfileBase._prepare(self)
if self.sym and self.refgroup is None:
raise ValueError(
"For symmetrization the `refgroup` argument is " "required."
)
logger.info(
f"Computing {self.grouping} profile along " f"{'XYZ'[self.dim]}-axes."
)
def _compute_histogram(
self, positions: np.ndarray, weights: Optional[np.ndarray] = None
) -> np.ndarray:
positions = positions[:, self.dim]
hist, _ = np.histogram(
positions, bins=self.n_bins, range=(self.zmin, self.zmax), weights=weights
)
return hist
def _single_frame(self) -> float:
PlanarBase._single_frame(self)
ProfileBase._single_frame(self)
# Take the center bin of the 0th group for correlation analysis.
return self._obs.profile[self.n_bins // 2]
def _conclude(self):
PlanarBase._conclude(self)
if self.sym:
symmetrize(self.sums.profile, inplace=True)
symmetrize(self.means.profile, inplace=True)
symmetrize(self.sems.profile, inplace=True)
if self.normalization == "number":
symmetrize(self.sums.bincount, inplace=True)
# Call conclude after symmetrize since `_concude` sets empty bins to `nan` and
# this prevents symmetrizing.
ProfileBase._conclude(self)