diff --git a/src/mdevaluate/distribution.py b/src/mdevaluate/distribution.py index 501b71a..bebe26d 100644 --- a/src/mdevaluate/distribution.py +++ b/src/mdevaluate/distribution.py @@ -1,4 +1,4 @@ -from typing import Callable, Optional, Union +from typing import Callable, Optional, Union, Tuple, List import numpy as np from numpy.typing import ArrayLike, NDArray @@ -52,6 +52,42 @@ def time_average( return np.mean(result, axis=0) +@autosave_data(nargs=2, kwargs_keys=("coordinates_b",)) +def time_distribution( + function: Callable, + coordinates: Coordinates, + coordinates_b: Optional[Coordinates] = None, + skip: float = 0.1, + segments: int = 100, +) -> Tuple[NDArray, List]: + """ + Compute the time distribution of a function. + + Args: + function: + The function that will be averaged, it has to accept exactly one argument + which is the current atom set (or two if coordinates_b is provided) + coordinates: The coordinates object of the simulation + coordinates_b: Additional coordinates object of the simulation + skip: + segments: + """ + frame_indices = np.unique( + np.int_( + np.linspace(len(coordinates) * skip, len(coordinates) - 1, num=segments) + ) + ) + times = np.array([coordinates[frame_index].time for frame_index in frame_indices]) + if coordinates_b is None: + result = [function(coordinates[frame_index]) for frame_index in frame_indices] + else: + result = [ + function(coordinates[frame_index], coordinates_b[frame_index]) + for frame_index in frame_indices + ] + return times, result + + def gr( atoms_a: CoordinateFrame, atoms_b: Optional[CoordinateFrame] = None,