diff --git a/src/mdevaluate/correlation.py b/src/mdevaluate/correlation.py index a8b3265..eeb042e 100644 --- a/src/mdevaluate/correlation.py +++ b/src/mdevaluate/correlation.py @@ -2,7 +2,7 @@ from typing import Callable, Optional import numpy as np from numpy.typing import ArrayLike -from scipy.special import legendre +from scipy.special import legendre, jn import dask.array as darray from functools import partial from scipy.spatial import KDTree @@ -183,6 +183,12 @@ def msd( displacements = displacements_without_drift(start_frame, end_frame, trajectory) if axis == "all": return (displacements**2).sum(axis=1).mean() + elif axis == "xy" or axis == "yx": + return (displacements[:, [0, 1]]**2).sum(axis=1).mean() + elif axis == "xz" or axis == "zx": + return (displacements[:, [0, 2]]**2).sum(axis=1).mean() + elif axis == "yz" or axis == "zy": + return (displacements[:, [1, 2]]**2).sum(axis=1).mean() elif axis == "x": return (displacements[:, 0] ** 2).mean() elif axis == "y": @@ -211,6 +217,15 @@ def isf( if axis == "all": distance = (displacements**2).sum(axis=1) ** 0.5 return np.sinc(distance * q / np.pi).mean() + elif axis == "xy" or axis == "yx": + distance = (displacements[:, [0, 1]]**2).sum(axis=1) ** 0.5 + return np.real(jn(0, distance * q / np.pi)).mean() + elif axis == "xz" or axis == "zx": + distance = (displacements[:, [0, 2]]**2).sum(axis=1) ** 0.5 + return np.real(jn(0, distance * q / np.pi)).mean() + elif axis == "yz" or axis == "zy": + distance = (displacements[:, [1, 2]]**2).sum(axis=1) ** 0.5 + return np.real(jn(0, distance * q / np.pi)).mean() elif axis == "x": distance = np.abs(displacements[:, 0]) return np.mean(np.cos(np.abs(q * distance))) @@ -262,6 +277,12 @@ def van_hove_self( vectors = displacements_without_drift(start_frame, end_frame, trajectory) if axis == "all": delta_r = (vectors**2).sum(axis=1) ** 0.5 + elif axis == "xy" or axis == "yx": + return (vectors[:, [0, 1]]**2).sum(axis=1) ** 0.5 + elif axis == "xz" or axis == "zx": + return (vectors[:, [0, 2]]**2).sum(axis=1) ** 0.5 + elif axis == "yz" or axis == "zy": + return (vectors[:, [1, 2]]**2).sum(axis=1) ** 0.5 elif axis == "x": delta_r = np.abs(vectors[:, 0]) elif axis == "y": @@ -423,6 +444,15 @@ def non_gaussian_parameter( if axis == "all": r = (vectors**2).sum(axis=1) dimensions = 3 + elif axis == "xy" or axis == "yx": + r = (vectors[:, [0, 1]]**2).sum(axis=1) + dimensions = 2 + elif axis == "xz" or axis == "zx": + r = (vectors[:, [0, 2]]**2).sum(axis=1) + dimensions = 2 + elif axis == "yz" or axis == "zy": + r = (vectors[:, [1, 2]]**2).sum(axis=1) + dimensions = 2 elif axis == "x": r = vectors[:, 0] ** 2 dimensions = 1