"""
Viscosity --- :mod:`transport_analysis.analysis.viscosity`
==========================================================
This module offers a class for the lightweight computation of shear
viscosity via the Einstein-Helfand method. It outputs the "viscosity
function," the product of viscosity and time as a function of time, from
which the slope is taken to calculate the shear viscosity. This is
described in eq. 5 of E M Kirova and G E Norman 2015 J. Phys.: Conf.
Ser. 653 012106.
"""
from typing import TYPE_CHECKING
from MDAnalysis.analysis.base import AnalysisBase
from MDAnalysis.core.groups import UpdatingAtomGroup
from MDAnalysis.exceptions import NoDataError
from MDAnalysis.units import constants
import numpy as np
if TYPE_CHECKING:
from MDAnalysis.core.universe import AtomGroup
[docs]class ViscosityHelfand(AnalysisBase):
"""
Class to calculate viscosity using the Einstein-Helfand approach.
Note that the slope of the viscosity function, the product of viscosity
and time as a function of time, must be taken to obtain the viscosity.
Particle velocities are required to calculate the viscosity function. Thus
you are limited to MDA trajectories that contain velocity information, e.g.
GROMACS .trr files, H5MD files, etc. See the MDAnalysis documentation:
https://docs.mdanalysis.org/stable/documentation_pages/coordinates/init.html#writers.
Parameters
----------
atomgroup : AtomGroup
An MDAnalysis :class:`~MDAnalysis.core.groups.AtomGroup`.
Note that :class:`UpdatingAtomGroup` instances are not accepted.
temp_avg : float (optional, default 300)
Average temperature over the course of the simulation, in Kelvin.
dim_type : {'xyz', 'xy', 'yz', 'xz', 'x', 'y', 'z'}
Desired dimensions to be included in the viscosity calculation.
Defaults to 'xyz'.
linear_fit_window : tuple of int (optional)
A tuple of two integers specifying the start and end lag-time for
the linear fit of the viscosity function. If not provided, the
linear fit is not performed and viscosity is not calculated.
Attributes
----------
atomgroup : :class:`~MDAnalysis.core.groups.AtomGroup`
The atoms to which this analysis is applied
dim_fac : int
Dimensionality :math:`d` of the viscosity computation.
results.timeseries : :class:`numpy.ndarray`
The averaged viscosity function over all the particles with respect
to lag-time. Obtained after calling :meth:`ViscosityHelfand.run`
results.visc_by_particle : :class:`numpy.ndarray`
The viscosity function of each individual particle with respect to
lag-time.
results.viscosity : float
The viscosity coefficient of the solvent. The argument
`linear_fit_window` must be provided to calculate this to
avoid misinterpretation of the viscosity function.
start : Optional[int]
The first frame of the trajectory used to compute the analysis.
stop : Optional[int]
The frame to stop at for the analysis.
step : Optional[int]
Number of frames to skip between each analyzed frame.
n_frames : int
Number of frames analysed in the trajectory.
n_particles : int
Number of particles viscosity was calculated over.
"""
def __init__(
self,
atomgroup: "AtomGroup",
temp_avg: float = 300.0,
dim_type: str = "xyz",
linear_fit_window: tuple[int, int] = None,
**kwargs,
) -> None:
# the below line must be kept to initialize the AnalysisBase class!
super().__init__(atomgroup.universe.trajectory, **kwargs)
# after this you will be able to access `self.results`
# `self.results` is a dictionary-like object
# that can should used to store and retrieve results
# See more at the MDAnalysis documentation:
# https://docs.mdanalysis.org/stable/documentation_pages/analysis/base.html?highlight=results#MDAnalysis.analysis.base.Results
if isinstance(atomgroup, UpdatingAtomGroup):
raise TypeError(
"UpdatingAtomGroups are not valid for viscosity computation"
)
# args
self.temp_avg = temp_avg
self.dim_type = dim_type.lower()
self.linear_fit_window = linear_fit_window
self._dim, self.dim_fac = self._parse_dim_type(self.dim_type)
# local
self.atomgroup = atomgroup
self.n_particles = len(self.atomgroup)
def _prepare(self):
"""
Set up viscosity, mass, velocity, and position arrays
before the analysis loop begins
"""
# 2D viscosity array of frames x particles
self.results.visc_by_particle = np.zeros(
(self.n_frames, self.n_particles)
)
self._volumes = np.zeros((self.n_frames))
self._masses = self.atomgroup.masses
self._masses_rs = self._masses.reshape((1, len(self._masses), 1))
# 3D arrays of frames x particles x dimensionality
# for velocities and positions
self._velocities = np.zeros(
(self.n_frames, self.n_particles, self.dim_fac)
)
self._positions = np.zeros(
(self.n_frames, self.n_particles, self.dim_fac)
)
# self.results.timeseries not set here
# update when mda 2.6.0 releases with typo fix
# (MDAnalysis Issue #4213)
try:
self.boltzmann = constants["Boltzmann_constant"]
except KeyError:
self.boltzmann = constants["Boltzman_constant"]
@staticmethod
def _parse_dim_type(dim_str):
r"""Sets up the desired dimensionality of the calculation."""
keys = {
"x": [0],
"y": [1],
"z": [2],
"xy": [0, 1],
"xz": [0, 2],
"yz": [1, 2],
"xyz": [0, 1, 2],
}
try:
_dim = keys[dim_str]
except KeyError:
raise ValueError(
"invalid dim_type: {} specified, please specify one of xyz, "
"xy, xz, yz, x, y, z".format(dim_str)
)
return _dim, len(_dim)
def _single_frame(self):
"""
Constructs arrays of velocities and positions
for viscosity computation.
"""
# This runs once for each frame of the trajectory
# The trajectory positions update automatically
# You can access the frame number using self._frame_index
# trajectory must have velocity and position information
if not (
self._ts.has_velocities
and self._ts.has_positions
and self._ts.volume != 0
):
raise NoDataError(
"Helfand viscosity computation requires "
"velocities, positions, and box volume in the trajectory"
)
# fill volume array
self._volumes[self._frame_index] = self._ts.volume
# set shape of velocity array
self._velocities[self._frame_index] = self.atomgroup.velocities[
:, self._dim
]
# set shape of position array
self._positions[self._frame_index] = self.atomgroup.positions[
:, self._dim
]
def _conclude(self):
"""
Calculates the viscosity function via the simple "windowed" algorithm.
"""
self._vol_avg = np.average(self._volumes)
lagtimes = np.arange(1, self.n_frames)
# iterate through all possible lagtimes from 1 to number of frames
for lag in lagtimes:
# get difference of momentum * position shifted by "lag" frames
diff = (
self._masses_rs
* self._velocities[:-lag, :, :]
* self._positions[:-lag, :, :]
- self._masses_rs
* self._velocities[lag:, :, :]
* self._positions[lag:, :, :]
)
# square and sum each x(, y, z) diff per particle
sq_diff = np.square(diff).mean(axis=-1)
# average over # frames
# update viscosity by particle array
self.results.visc_by_particle[lag, :] = np.mean(sq_diff, axis=0)
# divide by 2, boltzmann constant, vol_avg, and temp_avg
self.results.visc_by_particle = self.results.visc_by_particle / (
2 * self.boltzmann * self._vol_avg * self.temp_avg
)
# average over # particles and update results array
self.results.timeseries = self.results.visc_by_particle.mean(axis=1)
if self.linear_fit_window is not None:
fit_start, fit_end = (
self.linear_fit_window[0],
self.linear_fit_window[1],
)
linear_fit = np.polyfit(
lagtimes[fit_start:fit_end],
self.results.timeseries[fit_start:fit_end],
1,
)
self.results.viscosity = linear_fit[0]
[docs] def plot_viscosity_function(self):
"""
Plot the viscosity function as a function of lag-time.
If a linear fit window is provided, the window is highlighted.
"""
import matplotlib.pyplot as plt
lagtimes = np.arange(0, self.n_frames)
plt.plot(lagtimes, self.results.timeseries, label="Viscosity Function")
if self.linear_fit_window is not None:
fit_start, fit_end = (
self.linear_fit_window[0],
self.linear_fit_window[1],
)
plt.axvline(
fit_start, color="red", linestyle="--", label="Fit Start"
)
plt.axvline(fit_end, color="blue", linestyle="--", label="Fit End")
plt.xlabel("Lag-time")
plt.ylabel("Viscosity Function")
plt.title("Viscosity Function vs Lag-time")
plt.legend()
plt.show()