"""
Correlation functions
=====================
.. moduleauthor:: Benjamin Ye <GitHub: @bbye98>
This module contains algorithms for evaluating the spatial and temporal
correlation between data sets, such as the autocorrelation and
cross-correlation functions and the closely related mean squared
displacement.
"""
from typing import Union
import warnings
import numpy as np
from scipy import fft
[docs]
def correlation_fft(
x: np.ndarray[Union[float, complex]],
y: np.ndarray[Union[float, complex]] = None,
/,
axis: int = None,
*,
average: bool = False,
double: bool = False,
vector: bool = False,
) -> np.ndarray[Union[float, complex]]:
r"""
Evaluates the autocorrelation functions (ACF)
:math:`\mathrm{R_\mathbf{XX}}(\tau)` or cross-correlation functions
(CCF) :math:`\mathrm{R_\mathbf{XY}}(\tau)` of time series
:math:`\mathbf{X}(t)` and :math:`\mathbf{Y}(t)` using fast Fourier
transforms (FFT).
The fast convolution algorithm (FCA) [1]_ [2]_ is associated wtih
the Wiener–Khinchin theorem and has a time complexity of
:math:`\mathcal{O}(N\log{N})`.
The ACF for a data set :math:`\mathbf{X}(t)` can be computed using
.. math::
\begin{gather*}
\hat{\mathbf{X}}(\omega)=\mathcal{F}[\mathbf{X}(t)]\\
\mathrm{R}_{\mathbf{XX}}(\tau)=\mathcal{F}^{-1}
[\hat{\mathbf{X}}(\omega)\hat{\mathbf{X}}^*(\omega)]
\end{gather*}
where :math:`\tau` is the time lag and the asterisk (:math:`^*`)
denotes the complex conjugate.
Similarly, the CCF for data sets :math:`\mathbf{X}(t)` and
:math:`\mathbf{Y}(t)` can be computed using
.. math::
\mathrm{R}_{\mathbf{XY}}(\tau)=\mathcal{F}^{-1}[\mathcal{F}
[\mathbf{X}(t)]\cdot\mathcal{F}[\mathbf{Y}(t)]]
Parameters
----------
x : `numpy.ndarray`, positional-only
Time evolution of :math:`d`-dimensional data for :math:`N`
entities over :math:`N_\mathrm{b}` blocks of :math:`N_t` times
each.
.. container::
**Shape**:
* Scalar data: :math:`(N_t,)`, :math:`(N_t,\,N)`,
:math:`(N_\mathrm{b},\,N_t)`, or
:math:`(N_\mathrm{b},\,N_t,\,N)`.
* Vector data: :math:`(N_t,\,d)`, :math:`(N_t,\,N,\,d)`,
:math:`(N_\mathrm{b},\,N_t,\,d)`, or
:math:`(N_\mathrm{b},\,N_t,\,N,\,d)`.
y : `numpy.ndarray`, positional-only, optional
Time evolution of :math:`d`-dimensional data for another
:math:`N` entities over :math:`N_\mathrm{b}` blocks of
:math:`N_t` times each. If provided, the CCF for `x` and `y` is
calculated. Otherwise, the ACF for `x` is calculated.
**Shape**: Same as `x`.
axis : `int`, optional
Axis along which time evolves. If not specified, the axis is
determined automatically using the shape of `x`.
average : `bool`, keyword-only, default: :code:`True`
Determines whether the ACF/CCF is averaged over all entities.
Only available if `x` and `y` contain information for multiple
entities.
double : `bool`, keyword-only, default: :code:`False`
Determines whether the ACF is doubled or the negative and
positive time lags are combined for the CCF.
vector : `bool`, keyword-only, default: :code:`False`
Specifies whether `x` and `y` contain vectors. If :code:`True`,
the ACF/CCF is summed over the last axis.
Returns
-------
corr : `numpy.ndarray`
ACF or CCF.
.. container::
**Shape**:
For ACF, the shape is that of `x` but with the following
modifications:
* If :code:`average=True`, the axis indexing the :math:`N`
entities is removed.
* If :code:`vector=True`, the last axis is removed.
For CCF, the shape is that of `x` but with the following
modifications:
* If :code:`average=True`, the axis indexing the :math:`N`
entities is removed.
* If :code:`double=False`, the axis indexing the
:math:`N_t` times now has a length of :math:`2N_t-1` to
accomodate negative and positive time lags.
* If :code:`vector=True`, the last axis is removed.
References
----------
.. [1] Kneller, G. R.; Keiner, V.; Kneller, M.; Schiller, M.
NMOLDYN: A Program Package for a Neutron Scattering Oriented
Analysis of Molecular Dynamics Simulations. *Computer Physics
Communications* **1995**, *91* (1–3), 191–214.
https://doi.org/10.1016/0010-4655(95)00048-K.
.. [2] Calandrini, V.; Pellegrini, E.; Calligari, P.; Hinsen, K.;
Kneller, G. R. NMoldyn - Interfacing Spectroscopic Experiments,
Molecular Dynamics Simulations and Models for Time Correlation
Functions. *JDN* **2011**, *12*, 201–232.
https://doi.org/10.1051/sfn/201112010.
"""
# Ensure arrays have valid shapes
x = np.asarray(x)
if x.size == 0:
raise ValueError("The arrays cannot be empty.")
ndim = x.ndim
if not 1 <= ndim <= 4:
emsg = "The arrays must be one-, two-, three-, or four-" "dimensional."
raise ValueError(emsg)
if y is not None:
y = np.asarray(y)
if x.shape != y.shape:
raise ValueError("The arrays must have the same shape.")
# Check or set axis along which to compute the ACF/CCF
if axis is None:
if ndim == 4:
axis = 1
else:
axis = 0
if ndim > 1:
wmsg = (
"The axis along which to compute the ACF/CCF "
"was not specified and is ambiguous for a "
"multidimensional array. As such, it has been "
"set to the first axis by default."
)
warnings.warn(wmsg)
elif axis not in {0, 1}:
emsg = "The ACF/CCF can only be computed along the first or " "second axis."
raise ValueError(emsg)
# Determine whether faster real-valued FFTs can be used
if real := np.isrealobj(x) and (y is None or np.isrealobj(y)):
fft_ = fft.rfft
ifft = fft.irfft
else:
fft_ = fft.fft
ifft = fft.ifft
# Compute the power spectral density by first zero-padding the
# arrays for linear convolution and then inverting it to get the
# ACF/CCF
N_t = x.shape[axis]
n = 2 * fft.next_fast_len(N_t, real=real)
if y is None:
f = fft_(x, n=n, axis=axis)
r = ifft(f * f.conj(), axis=axis)
r = (double + 1) * (r[:, :N_t] if axis else r[:N_t])
else:
fx = fft_(x, n=n, axis=axis)
fy = fft_(y, n=n, axis=axis)
f = fx.conj() * fy
if double:
r = ifft(f + fx * fy.conj(), axis=axis)
r = r[:, :N_t] if axis else r[:N_t]
else:
r = ifft(f, axis=axis)
# Sum over the last axis if the arrays contain vectors
if vector:
r = r.sum(axis=-1)
# Determine the axes over which to expand the dimensions of the
# reversed time array for correct matrix division
axes = list(range(ndim - vector))
axes.remove(axis)
# Normalize the ACF/CCF
if axis:
r[:, :N_t] /= np.expand_dims(np.arange(N_t, 0, -1), axes)
if r.shape[axis] != N_t:
r[:, 1 - N_t :] /= np.expand_dims(np.arange(1, N_t), axes)
r = np.hstack((r[:, 1 - N_t :], r[:, :N_t]))
else:
r[:N_t] /= np.expand_dims(np.arange(N_t, 0, -1), axes)
if r.shape[axis] != N_t:
r[1 - N_t :] /= np.expand_dims(np.arange(1, N_t), axes)
r = np.concatenate((r[1 - N_t :], r[:N_t]))
# Average over all entities, if desired
if average:
axis_avg = ndim - vector - 1
if axis != axis_avg:
return r.mean(axis=axis_avg)
return r
[docs]
def correlation_shift(
x: np.ndarray[Union[float, complex]],
y: np.ndarray[Union[float, complex]] = None,
/,
axis: int = None,
*,
average: bool = False,
double: bool = False,
vector: bool = False,
) -> np.ndarray[Union[float, complex]]:
r"""
Evaluates the autocorrelation functions (ACF)
:math:`\mathrm{R_\mathbf{XX}}(\tau)` or cross-correlation functions
(CCF) :math:`\mathrm{R_\mathbf{XY}}(\tau)` of time series
:math:`\mathbf{X}(t)` and :math:`\mathbf{Y}(t)` directly by using
sliding windows.
The ACF for a data set :math:`\mathbf{X}(t)` can be computed using
.. math::
\mathrm{R}_{\mathbf{XX}}(\tau)
=\langle\mathbf{X}(t_0+\tau)\cdot\mathbf{X}^*(t_0)\rangle
=\dfrac{1}{N_\tau}\sum_{j=1}^{N_\tau}
\textbf{X}(t_j+\tau)\cdot\textbf{X}^*(t_j)
where :math:`\tau` is the time lag, :math:`t_j` is an arbitrary
reference time, :math:`N_\tau` is the number of possible reference
times, and the asterisk (:math:`^*`) denotes the complex conjugate.
Similarly, the CCF for data sets :math:`\mathbf{X}(t)` and
:math:`\mathbf{Y}(t)` can be computed using
.. math::
\mathrm{R}_{\mathbf{XY}}(\tau)
=\langle\mathbf{X}(t_0+\tau)\cdot\mathbf{Y}^*(t_0)\rangle
=\dfrac{1}{N_\tau}\sum_{j=1}^{N_\tau}
\textbf{X}(t_j+\tau)\cdot\textbf{Y}^*(t_j)
To minimize statistical noise, the ACF/CCF is calculated for and
averaged over all possible reference times :math:`t_0`. As such,
this algorithm has a time complexity of :math:`\mathcal{O}(N^2)`.
With large data sets, this approach is too slow to be useful. If
your machine supports fast Fourier transforms (FFT), use the much
more performant FFT-based algorithm implemented in
:func:`mdcraft.algorithm.correlation.correlation_fft` instead.
Parameters
----------
x : `numpy.ndarray`, positional-only
Time evolution of :math:`d`-dimensional data for :math:`N`
entities over :math:`N_\mathrm{b}` blocks of :math:`N_t` times
each.
.. container::
**Shape**:
* Scalar data: :math:`(N_t,)`, :math:`(N_t,\,N)`,
:math:`(N_\mathrm{b},\,N_t)`, or
:math:`(N_\mathrm{b},\,N_t,\,N)`.
* Vector data: :math:`(N_t,\,d)`, :math:`(N_t,\,N,\,d)`,
:math:`(N_\mathrm{b},\,N_t,\,d)`, or
:math:`(N_\mathrm{b},\,N_t,\,N,\,d)`.
y : `numpy.ndarray`, positional-only, optional
Time evolution of :math:`d`-dimensional data for another
:math:`N` entities over :math:`N_\mathrm{b}` blocks of
:math:`N_t` times each. If provided, the CCF for `x` and `y` is
calculated. Otherwise, the ACF for `x` is calculated.
**Shape**: Same as `x`.
axis : `int`, optional
Axis along which time evolves. If not specified, the axis is
determined automatically using the shape of `x`.
average : `bool`, keyword-only, default: :code:`True`
Determines whether the ACF/CCF is averaged over all entities.
Only available if `x` and `y` contain information for multiple
entities.
double : `bool`, keyword-only, default: :code:`False`
Determines whether the ACF is doubled or the negative and
positive time lags are combined for the CCF.
vector : `bool`, keyword-only, default: :code:`False`
Specifies whether `x` and `y` contain vectors. If :code:`True`,
the ACF/CCF is summed over the last axis.
Returns
-------
corr : `numpy.ndarray`
ACF or CCF.
.. container::
**Shape**:
For ACF, the shape is that of `x` but with the following
modifications:
* If :code:`average=True`, the axis indexing the :math:`N`
entities is removed.
* If :code:`vector=True`, the last axis is removed.
For CCF, the shape is that of `x` but with the following
modifications:
* If :code:`average=True`, the axis indexing the :math:`N`
entities is removed.
* If :code:`double=False`, the axis indexing the
:math:`N_t` times now has a length of :math:`2N_t-1` to
accomodate negative and positive time lags.
* If :code:`vector=True`, the last axis is removed.
"""
# Ensure arrays have valid shapes
x = np.asarray(x)
if x.size == 0:
raise ValueError("The arrays cannot be empty.")
ndim = x.ndim
if not 1 <= ndim <= 4:
emsg = "The arrays must be one-, two-, three-, or four-" "dimensional."
raise ValueError(emsg)
if y is not None:
y = np.asarray(y)
if x.shape != y.shape:
raise ValueError("The arrays must have the same shape.")
# Check or set axis along which to compute the ACF/CCF
if axis is None:
if ndim == 4:
axis = 1
else:
axis = 0
if ndim > 1:
wmsg = (
"The axis along which to compute the ACF/CCF "
"was not specified and is ambiguous for a "
"multidimensional array. As such, it has been "
"set to the first axis by default."
)
warnings.warn(wmsg)
elif axis not in {0, 1}:
emsg = "The ACF/CCF can only be computed along the first or " "second axis."
raise ValueError(emsg)
# Calculate the ACF/CCF
N_t = x.shape[axis]
if y is None:
if ndim == 1:
r = np.fromiter(
(np.dot(x[i:], x[: -i if i else None]) for i in range(N_t)),
dtype=float,
count=N_t,
)
elif axis:
axes = f"bt...{'d' * vector}"
r = np.stack(
[
np.einsum(
f"{axes},{axes}->b...", x[:, i:], x[:, : -i if i else None]
)
for i in range(N_t)
],
axis=1,
)
else:
axes = f"t...{'d' * vector}"
r = np.stack(
[
np.einsum(f"{axes},{axes}->...", x[i:], x[: -i if i else None])
for i in range(N_t)
]
)
else:
start = np.r_[np.zeros(N_t - 1, dtype=int), 0:N_t]
stop = np.r_[1 : N_t + 1, N_t * np.ones(N_t - 1, dtype=int)]
if ndim == 1:
r = np.fromiter(
(
np.dot(x[i:j], y[k:m])
for i, j, k, m in zip(start[::-1], stop[::-1], start, stop)
),
dtype=float,
count=2 * N_t - 1,
)
elif axis:
axes = f"bt...{'d' * vector}"
r = np.stack(
[
np.einsum(f"{axes},{axes}->b...", x[:, i:j], y[:, k:m])
for i, j, k, m in zip(start[::-1], stop[::-1], start, stop)
],
axis=1,
)
else:
axes = f"t...{'d' * vector}"
r = np.stack(
[
np.einsum(f"{axes},{axes}->...", x[i:j], y[k:m])
for i, j, k, m in zip(start[::-1], stop[::-1], start, stop)
]
)
# Double the ACF or overlap the negative and positive time lags for
# the CCF, if desired
if double:
if y is None:
r *= 2
elif axis:
r = r[:, N_t - 1 :] + r[:, N_t - 1 :: -1]
else:
r = r[N_t - 1 :] + r[N_t - 1 :: -1]
# Determine the axes over which to expand the dimensions of the
# reversed time array for correct matrix division
axes = list(range(ndim - vector))
axes.remove(axis)
# Normalize the ACF/CCF
if axis:
r[:, -N_t:] /= np.expand_dims(np.arange(N_t, 0, -1), axes)
if r.shape[axis] != N_t:
r[:, : N_t - 1] /= np.expand_dims(np.arange(1, N_t), axes)
else:
r[-N_t:] /= np.expand_dims(np.arange(N_t, 0, -1), axes)
if r.shape[axis] != N_t:
r[: N_t - 1] /= np.expand_dims(np.arange(1, N_t), axes)
# Average over all entities, if desired
if average:
axis_avg = ndim - 1 - vector
if axis != axis_avg:
return r.mean(axis=axis_avg)
return r
[docs]
def msd_fft(
r_i: np.ndarray[float],
r_j: np.ndarray[float] = None,
/,
axis: int = None,
*,
average: bool = True,
) -> np.ndarray[float]:
r"""
Evaluates the mean squared displacements (MSD) or the analogous
cross displacements (CD) of positions :math:`\mathbf{r}_i(t)` and
:math:`\mathbf{r}_j(t)` using fast Fourier transforms (FFT).
For a set of positions :math:`\mathbf{r}_i(t)`, the MSD is computed
using the algorithm [1]_ [2]_
.. math::
\mathrm{MSD}_{i,\,m}&=\frac{1}{N_t-m}\sum_{k=0}^{N_t-m-1}
[\textbf{r}_{i,\,k+m}-\textbf{r}_{i,\,k}]^2\\
&=\frac{1}{N_t-m}\sum_{k=0}^{N_t-m-1}
\left[\textbf{r}_{i,\,k+m}^2+\textbf{r}_{i,\,k}^2\right]
-\frac{2}{N_t-m}\sum_{k=0}^{N_t-m-1}
\textbf{r}_{i,\,k}\cdot\textbf{r}_{i,\,k+m}\\
&=S_{ii,\,m}-2\mathrm{R}_{ii,\,m}
where :math:`i` is the species index, :math:`m` is the index
corresponding to time lag :math:`\tau`, :math:`\mathrm{R}_{ii,\,m}`
is the autocorrelation of :math:`\mathbf{r}_i(t)`, and :math:`S_m`
is evaluated using the recursive relation
.. math::
\begin{gather*}
D_{ii,\,m}=\textbf{r}_{i,\,m}^2\\
Q_{ii,\,-1}=2\sum_{k=0}^{N_t-1}D_{ii,\,k}\\
Q_{ii,\,m}=Q_{ii,\,m-1}-D_{ii,\,m-1}-D_{ii,\,N_t-m}\\
S_{ii,\,m}=\frac{Q_{ii,\,m}}{N_t-m}
\end{gather*}
Similarly, the CD for two sets of positions :math:`\mathbf{r}_i(t)`
and :math:`\mathbf{r}_j(t)` is computed using
.. math::
\mathrm{CD}_{ij,m}=S_{ij,\,m}-2\mathrm{R}_{ij,\,m}
where :math:`\mathrm{R}_{ij,\,m}` is the cross-correlation of
:math:`\mathbf{r}_i(t)` and :math:`\mathbf{r}_j(t)`, and
:math:`S_{ij,\,m}` is evaluated using the recursive relation
.. math::
\begin{gather*}
D_{ij,\,m}=\textbf{r}_{i,\,m}\cdot\textbf{r}_{j,\,m}\\
Q_{ij,\,-1}=2\sum_{k=0}^{N_t-1}D_{ij,\,k}\\
Q_{ij,\,m}=Q_{ij,\,m-1}-D_{ij,\,m-1}-D_{ij,\,N_t-m}\\
S_{ij,\,m}=\frac{Q_{ij,\,m}}{N_t-m}
\end{gather*}
.. note::
To evaluate the sum in the expression used to calculate the
Onsager transport coefficients [3]_
.. math::
L_{ij}=\frac{1}{6k_\mathrm{B}T}\lim_{t\rightarrow\infty}
\frac{d}{d\tau}\left\langle
\sum_{\alpha=1}^{N_i}[\mathbf{r}_{i,\,\alpha}(t_0+\tau)
-\mathbf{r}_{i,\,\alpha}(t_0)]\cdot
\sum_{\beta=1}^{N_j}[\mathbf{r}_{j,\,\beta}(t_0+\tau)
-\mathbf{r}_{j,\,\beta(t_0)}]\right\rangle
`r_i` and `r_j` should be summed over all entities before being
passed to this function.
Parameters
----------
r_i : `numpy.ndarray`, positional-only
Time evolution of individual or averaged positions for :math:`N`
entities over :math:`N_\mathrm{b}` blocks of :math:`N_t` times
each.
**Shape**: :math:`(N_t,\,3)`, :math:`(N_t,\,N,\,3)`,
:math:`(N_\mathrm{b},\,N_t,\,3)`, or
:math:`(N_\mathrm{b},\,N_t,\,N,\,3)`.
**Reference unit**: :math:`\mathrm{Å}`.
r_j : `numpy.ndarray`, positional-only, optional
Time evolution of individual or averaged positions for another
:math:`N` entities over :math:`N_\mathrm{b}` blocks of
:math:`N_t` times each.
**Shape**: Same as `r_i`.
**Reference unit**: :math:`\mathrm{Å}`.
axis : `int`, optional
Axis along which time evolves. If not specified, the axis is
determined automatically using the shape of `r_i`.
average : `bool`, keyword-only, default: :code:`True`
Determines whether the MSD/CD is averaged over all entities.
Only available if `r_i` and `r_j` contain information for
multiple entities.
Returns
-------
disp : `numpy.ndarray`
MSD or CD.
**Shape**: Same as the shape of `r_i`, except with the last axis
removed. If :code:`average=True`, the axis indexing the
:math:`N` entities is also removed.
**Reference unit**: :math:`\text{Å}^2`.
References
----------
.. [1] Kneller, G. R.; Keiner, V.; Kneller, M.; Schiller, M.
NMOLDYN: A Program Package for a Neutron Scattering Oriented
Analysis of Molecular Dynamics Simulations. *Computer Physics
Communications* **1995**, *91* (1–3), 191–214.
https://doi.org/10.1016/0010-4655(95)00048-K.
.. [2] Calandrini, V.; Pellegrini, E.; Calligari, P.; Hinsen, K.;
Kneller, G. R. NMoldyn - Interfacing Spectroscopic Experiments,
Molecular Dynamics Simulations and Models for Time Correlation
Functions. *JDN* **2011**, *12*, 201–232.
https://doi.org/10.1051/sfn/201112010.
.. [3] Fong, K. D.; Self, J.; McCloskey, B. D.; Persson, K. A.
Onsager Transport Coefficients and Transference Numbers in
Polyelectrolyte Solutions and Polymerized Ionic Liquids.
*Macromolecules* **2020**, *53* (21), 9503–9512.
https://doi.org/10.1021/acs.macromol.0c02001.
"""
# Ensure arrays have valid shapes
r_i = np.asarray(r_i)
if r_i.size == 0:
raise ValueError("The position arrays cannot be empty.")
if r_i.shape[-1] != 3:
emsg = "The position arrays must have three components in " "the last axis."
raise ValueError(emsg)
ndim = r_i.ndim
if not 2 <= ndim <= 4:
emsg = "The position arrays must be two-, three-, or four-" "dimensional."
raise ValueError(emsg)
if r_j is not None:
r_j = np.asarray(r_j)
if r_i.shape != r_j.shape:
raise ValueError("The position arrays must have the same shape.")
# Check or set axis along which to compute the MSD/CD
if axis is None:
if ndim == 4:
axis = 1
else:
axis = 0
if ndim == 3:
emsg = (
"The axis along which to compute the MSD/CD "
"was not specified and is ambiguous for a "
"three-dimensional array. As such, it has been "
"set to the first axis by default."
)
warnings.warn(emsg)
elif axis not in {0, 1}:
emsg = "The MSD/CD can only be computed along the first or " "second axis."
raise ValueError(emsg)
# Get intermediate quantities required for the MSD/CD calculation
R_ij = correlation_fft(r_i, r_j, axis, average=False, double=True, vector=True)
D_ij = (r_i * (r_i if r_j is None else r_j)).sum(axis=-1)
N_t = r_i.shape[axis]
if ndim - axis == 3:
# Calculate the MSD/CD for each entity
if not average:
stack = np.vstack if ndim == 3 else np.hstack
shape = np.asarray(r_i.shape[:-1])
mask = np.ones_like(shape, dtype=bool)
mask[axis] = False
D_k = stack((D_ij, np.expand_dims(np.zeros(shape[mask]), axis)))
if axis:
Q_ij = 2 * D_k.sum(axis=axis, keepdims=True) * np.ones(
(1, N_t, 1)
) - np.cumsum(
D_k[:, np.arange(-1, N_t - 1)] + D_k[:, N_t:0:-1], axis=axis
)
else:
Q_ij = 2 * D_k.sum(axis=axis) * np.ones((N_t, 1)) - np.cumsum(
D_k[np.arange(-1, N_t - 1)] + D_k[N_t:0:-1], axis=axis
)
return Q_ij / np.arange(N_t, 0, -1)[:, None] - R_ij
# Average the intermediate quantities over all particles
R_ij = R_ij.mean(axis=ndim - 2)
D_ij = D_ij.mean(axis=ndim - 2)
# Calculate the averaged MSD/CD
if axis:
Q_ij = 2 * D_ij.sum(axis=axis, keepdims=True) * np.ones((1, N_t)) - np.insert(
np.cumsum(D_ij[:, : N_t - 1] + D_ij[:, N_t - 1 : 0 : -1], axis=axis),
0,
0,
axis=axis,
)
else:
Q_ij = 2 * D_ij.sum() * np.ones(N_t) - np.insert(
np.cumsum(D_ij[: N_t - 1] + D_ij[N_t - 1 : 0 : -1]), 0, 0
)
return Q_ij / np.arange(N_t, 0, -1) - R_ij
[docs]
def msd_shift(
r_i: np.ndarray[float],
r_j: np.ndarray[float] = None,
/,
axis: int = None,
*,
average: bool = True,
) -> np.ndarray[float]:
r"""
Evaluates the mean squared displacements (MSD) or the analogous
cross displacements (CD) of positions :math:`\mathbf{r}_i(t)` and
:math:`\mathbf{r}_j(t)` using the Einstein relation.
For a set of positions :math:`\mathbf{r}_i(t)`, the MSD is defined
as
.. math::
\mathrm{MSD}_i(\tau)=\left\langle[\textbf{r}_i(t_0+\tau)
-\textbf{r}_i(t_0)]^2\right\rangle
=\dfrac{1}{N_\tau}\sum_{k=1}^{N_\tau}[\textbf{r}_i(t_k+\tau)
-\textbf{r}_i(t_k)]^2
where :math:`\tau` is the time lag, :math:`t_j` is an arbitrary
reference time, and :math:`N_\tau` is the number of possible
reference times.
Similarly, the CD for two sets of positions :math:`\mathbf{r}_i(t)`
and :math:`\mathbf{r}_j(t)` is defined as
.. math::
\mathrm{CD}_{ij}(\tau)&=\langle
[\textbf{r}_i(t_0+\tau)-\textbf{r}_i(t_0)]\cdot
[\textbf{r}_j(t_0+\tau)-\textbf{r}_j(t_0)]\rangle\\
&=\dfrac{1}{N_\tau}\sum_{k=1}^{N_\tau}
[\textbf{r}_i(t_k+\tau)-\textbf{r}_i(t_k))\cdot
(\textbf{r}_j(t_k+\tau)-\textbf{r}_j(t_k)]
To minimize statistical noise, the MSD/CD is calculated for and
averaged over all possible reference times :math:`t_0`.
.. note::
To evaluate the sum in the expression used to calculate the
Onsager transport coefficients [1]_
.. math::
L_{ij}=\frac{1}{6k_\mathrm{B}T}\lim_{t\rightarrow\infty}
\frac{d}{d\tau}\left\langle
\sum_{\alpha=1}^{N_i}[\mathbf{r}_{i,\,\alpha}(t_0+\tau)
-\mathbf{r}_{i,\,\alpha}(t_0)]\cdot
\sum_{\beta=1}^{N_j}[\mathbf{r}_{j,\,\beta}(t_0+\tau)
-\mathbf{r}_{j,\,\beta(t_0)}]\right\rangle
`r_i` and `r_j` should be summed over all entities before being
passed to this function.
Parameters
----------
r_i : `numpy.ndarray`, positional-only
Time evolution of individual or averaged positions for :math:`N`
entities over :math:`N_\mathrm{b}` blocks of :math:`N_t` times
each.
**Shape**: :math:`(N_t,\,3)`, :math:`(N_t,\,N,\,3)`,
:math:`(N_\mathrm{b},\,N_t,\,3)`, or
:math:`(N_\mathrm{b},\,N_t,\,N,\,3)`.
**Reference unit**: :math:`\mathrm{Å}`.
r_j : `numpy.ndarray`, positional-only, optional
Time evolution of individual or averaged positions for another
:math:`N` entities over :math:`N_\mathrm{b}` blocks of
:math:`N_t` times each.
**Shape**: Same as `r_i`.
**Reference unit**: :math:`\mathrm{Å}`.
axis : `int`, optional
Axis along which time evolves. If not specified, the axis is
determined automatically using the shape of `r_i`.
average : `bool`, keyword-only, default: :code:`True`
Determines whether the MSD/CD is averaged over all entities.
Only available if `r_i` and `r_j` contain information for
multiple entities.
Returns
-------
disp : `numpy.ndarray`
MSD or CD.
**Shape**: Same as the shape of `r_i`, except with the last axis
removed. If :code:`average=True`, the axis indexing the
:math:`N` entities is also removed.
**Reference unit**: :math:`\text{Å}^2`.
References
----------
.. [1] Fong, K. D.; Self, J.; McCloskey, B. D.; Persson, K. A.
Onsager Transport Coefficients and Transference Numbers in
Polyelectrolyte Solutions and Polymerized Ionic Liquids.
*Macromolecules* **2020**, *53* (21), 9503–9512.
https://doi.org/10.1021/acs.macromol.0c02001.
"""
# Ensure arrays have valid shapes
r_i = np.asarray(r_i)
if r_i.size == 0:
raise ValueError("The position arrays cannot be empty.")
if r_i.shape[-1] != 3:
emsg = "The position arrays must have three components in " "the last axis."
raise ValueError(emsg)
ndim = r_i.ndim
if not 2 <= ndim <= 4:
emsg = "The position arrays must be two-, three-, or four-" "dimensional."
raise ValueError(emsg)
if r_j is not None:
r_j = np.asarray(r_j)
if r_i.shape != r_j.shape:
raise ValueError("The position arrays must have the same shape.")
# Check or set axis along which to compute the MSD/CD
if axis is None:
if ndim == 4:
axis = 1
else:
axis = 0
if ndim == 3:
emsg = (
"The axis along which to compute the MSD/CD "
"was not specified and is ambiguous for a "
"three-dimensional array. As such, it has been "
"set to the first axis by default."
)
warnings.warn(emsg)
elif axis not in {0, 1}:
emsg = "The MSD/CD can only be computed along the first or " "second axis."
raise ValueError(emsg)
# Calculate the MSD/CD for each entity
N_t = r_i.shape[axis]
if r_j is None:
if axis:
disp = np.stack(
[
((r_i[:, : -i if i else None] - r_i[:, i:]) ** 2)
.sum(axis=-1)
.mean(axis=axis)
for i in range(N_t)
],
axis=1,
)
else:
disp = np.stack(
[
((r_i[: -i if i else None] - r_i[i:]) ** 2)
.sum(axis=-1)
.mean(axis=axis)
for i in range(N_t)
]
)
else:
if axis:
disp = np.stack(
[
np.einsum(
"bt...d,bt...d->bt...",
r_i[:, : -i if i else None] - r_i[:, i:],
r_j[:, : -i if i else None] - r_j[:, i:],
).mean(axis=axis)
for i in range(N_t)
],
axis=1,
)
else:
disp = np.stack(
[
np.einsum(
"t...d,t...d->t...",
r_i[: -i if i else None] - r_i[i:],
r_j[: -i if i else None] - r_j[i:],
).mean(axis=axis)
for i in range(N_t)
]
)
# Average over all entities, if desired
if ndim - axis == 3 and average:
disp = disp.mean(axis=ndim - 2)
return disp