Added time_distribution function

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,