Improved performance and split up shifted_correlation

src/mdevaluate/correlation.py

@@ -13,12 +13,97 @@ from .pbc import pbc_diff, pbc_points
 from .coordinates import Coordinates, CoordinateFrame, displacements_without_drift
 
 
-def log_indices(first: int, last: int, num: int = 100) -> np.ndarray:
+def _is_multi_selector(selection):
-    ls = np.logspace(0, np.log10(last - first + 1), num=num)
+    if len(selection) == 0:
-    return np.unique(np.int_(ls) - 1 + first)
+        return False
+    elif (
+        isinstance(selection[0], int)
+        or isinstance(selection[0], bool)
+        or isinstance(selection[0], np.integer)
+        or isinstance(selection[0], np.bool_)
+    ):
+        return False
+    else:
+        for indices in selection:
+            if len(indices) == 0:
+                continue
+            elif (
+                isinstance(indices[0], int)
+                or isinstance(indices[0], bool)
+                or isinstance(indices[0], np.integer)
+                or isinstance(indices[0], np.bool_)
+            ):
+                return True
+            else:
+                raise ValueError(
+                    "selector has more than two dimensions or does not "
+                    "contain int or bool types"
+                )
+
+
+def _calc_correlation(
+    frames: Coordinates,
+    start_frame: CoordinateFrame,
+    function: Callable,
+    selection: np.ndarray,
+    shifted_idx: np.ndarray,
+) -> np.ndarray:
+    if len(selection) == 0:
+        correlation = np.zeros(len(shifted_idx))
+    else:
+        start = start_frame[selection]
+        correlation = np.array(
+            [
+                function(start, frames[frame_index][selection])
+                for frame_index in shifted_idx
+            ]
+        )
+    return correlation
+
+
+def _calc_correlation_multi(
+    frames: Coordinates,
+    start_frame: CoordinateFrame,
+    function: Callable,
+    selection: np.ndarray,
+    shifted_idx: np.ndarray,
+) -> np.ndarray:
+    correlations = np.zeros((len(selection), len(shifted_idx)))
+    for i, frame_index in enumerate(shifted_idx):
+        frame = frames[frame_index]
+        for j, current_selection in enumerate(selection):
+            if len(selection) == 0:
+                correlations[j, i] = 0
+            else:
+                correlations[j, i] = function(
+                    start_frame[current_selection], frame[current_selection]
+                )
+    return correlations
+
+
+def _average_correlation(result):
+    averaged_result = []
+    for n in range(result.shape[1]):
+        clean_result = []
+        for entry in result[:, n]:
+            if np.all(entry == 0):
+                continue
+            else:
+                clean_result.append(entry)
+        averaged_result.append(np.average(np.array(clean_result), axis=0))
+    return np.array(averaged_result)
+
+
+def _average_correlation_multi(result):
+    clean_result = []
+    for entry in result:
+        if np.all(entry == 0):
+            continue
+        else:
+            clean_result.append(entry)
+    return np.average(np.array(clean_result), axis=0)
 
 
-@autosave_data(2)
 def shifted_correlation(
     function: Callable,
     frames: Coordinates,
@@ -28,113 +113,82 @@ def shifted_correlation(
     window: float = 0.5,
     average: bool = True,
     points: int = 100,
-) -> (np.ndarray, np.ndarray):
+) -> tuple[np.ndarray, np.ndarray]:
+    """Compute a time-dependent correlation function for a given trajectory.
+
+    To improve statistics, multiple (possibly overlapping) windows will be
+    layed over the whole trajectory and the correlation is computed for them separately.
+    The start frames of the windows are spaced linearly over the valid region of
+    the trajectory (skipping frames in the beginning given by skip parameter).
+
+    The points within each window are spaced logarithmically.
+
+    Only a certain subset of the given atoms may be selected for each window
+    individually using a selector function.
+
+    Note that this function is specifically optimized for multi selectors, which select
+    multiple selection sets per window, for which the correlation is to be computed
+    separately.
+
+
+    Arguments
+    ---------
+    function:
+        The (correlation) function to evaluate.
+        Should be of the form (CoordinateFrame, CoordinateFrame) -> float
+
+    frames:
+        Trajectory to evaluate on
+
+    selector: (optional)
+        Selection function so select only certain selection sets for each start frame.
+        Should be of the form
+            (CoordinateFrame) -> list[A]
+        where A is something you can index an ndarray with.
+        For example a list of indices or a bool array.
+        Must return the same number of selection sets for every frame.
+
+    segments:
+        Number of start frames
+
+    skip:
+        Percentage of trajectory to skip from the start
+
+    window:
+        Length of each segment given as percentage of trajectory
+
+    average:
+        Whether to return averaged results.
+        See below for details on the returned ndarray.
+
+    points:
+        Number of points per segment
+
+
+    Returns
+    -------
+    times: ndarray
+        1d array of time differences to start frame
+    result: ndarray
+        2d ndarray of averaged (or non-averaged) correlations.
+
+        When average==True (default) the returned array will be of the shape (S, P)
+        where S is the number of selection sets and P the number of points per window.
+        For selection sets that where empty for all start frames all data points will be
+        zero.
+
+        When average==False the returned array will be of shape (W, S) with
+        dtype=object. The elements are either ndarrays of shape (P,) containing the
+        correlation data for the specific window and selection set or None if the
+        corresponding selection set was empty.
+        W is the number of segments (windows).
+        S and P are the same as for average==True.
+
     """
-    Calculate the time series for a correlation function.
-
-    The times at which the correlation is calculated are determined by
-    a logarithmic distribution.
-
-    Args:
-        function: The function that should be correlated
-        frames: The coordinates of the simulation data
-        selector (opt.):
-                    A function that returns the indices depending on
-                    the staring frame for which particles the
-                    correlation should be calculated.
-        segments (int, opt.):
-                    The number of segments the time window will be
-                    shifted
-        skip (float, opt.):
-                    The fraction of the trajectory that will be skipped
-                    at the beginning, if this is None the start index
-                    of the frames slice will be used, which defaults
-                    to 0.1.
-        window (float, opt.):
-                    The fraction of the simulation the time series will
-                    cover
-        average (bool, opt.):
-                    If True, returns averaged correlation function
-        points (int, opt.):
-                    The number of timeshifts for which the correlation
-                    should be calculated
-    Returns:
-        tuple:
-            A list of length N that contains the timeshiftes of the frames at which
-            the time series was calculated and a numpy array of shape (segments, N)
-            that holds the (non-avaraged) correlation data
-
-    Example:
-        Calculating the mean square displacement of a coordinate object
-        named ``coords``:
-
-        >>> time, data = shifted_correlation(msd, coords)
-    """
-
-    def get_correlation(
-        frames: CoordinateFrame,
-        start_frame: CoordinateFrame,
-        index: np.ndarray,
-        shifted_idx: np.ndarray,
-    ) -> np.ndarray:
-        if len(index) == 0:
-            correlation = np.zeros(len(shifted_idx))
-        else:
-            start = frames[start_frame][index]
-            correlation = np.array(
-                [function(start, frames[frame][index]) for frame in shifted_idx]
-            )
-        return correlation
-
-    def apply_selector(
-        start_frame: CoordinateFrame,
-        frames: CoordinateFrame,
-        idx: np.ndarray,
-        selector: Optional[Callable] = None,
-    ):
-        shifted_idx = idx + start_frame
-
-        if selector is None:
-            index = np.arange(len(frames[start_frame]))
-            return get_correlation(frames, start_frame, index, shifted_idx)
-        else:
-            index = selector(frames[start_frame])
-            if len(index) == 0:
-                return np.zeros(len(shifted_idx))
-
-            elif (
-                isinstance(index[0], int)
-                or isinstance(index[0], bool)
-                or isinstance(index[0], np.integer)
-                or isinstance(index[0], np.bool_)
-            ):
-                return get_correlation(frames, start_frame, index, shifted_idx)
-            else:
-                correlations = []
-                for ind in index:
-                    if len(ind) == 0:
-                        correlations.append(np.zeros(len(shifted_idx)))
-
-                    elif (
-                        isinstance(ind[0], int)
-                        or isinstance(ind[0], bool)
-                        or isinstance(ind[0], np.integer)
-                        or isinstance(ind[0], np.bool_)
-                    ):
-                        correlations.append(
-                            get_correlation(frames, start_frame, ind, shifted_idx)
-                        )
-                    else:
-                        raise ValueError(
-                            "selector has more than two dimensions or does not "
-                            "contain int or bool types"
-                        )
-                return correlations
-
     if 1 - skip < window:
         window = 1 - skip
 
-    start_frames = np.unique(
+    start_frame_indices = np.unique(
         np.linspace(
             len(frames) * skip,
             len(frames) * (1 - window),
@@ -144,28 +198,44 @@ def shifted_correlation(
         )
     )
 
-    num_frames = int(len(frames) * window)
+    num_frames_per_window = int(len(frames) * window)
-    ls = np.logspace(0, np.log10(num_frames + 1), num=points)
+    logspaced_indices = np.logspace(0, np.log10(num_frames_per_window + 1), num=points)
-    idx = np.unique(np.int_(ls) - 1)
+    logspaced_indices = np.unique(np.int_(logspaced_indices) - 1)
-    t = np.array([frames[i].time for i in idx]) - frames[0].time
+    logspaced_time = (
-
+        np.array([frames[i].time for i in logspaced_indices]) - frames[0].time
-    result = np.array(
-        [
-            apply_selector(start_frame, frames=frames, idx=idx, selector=selector)
-            for start_frame in start_frames
-        ]
     )
 
-    if average:
+    if selector is None:
-        clean_result = []
+        multi_selector = False
-        for entry in result:
-            if np.all(entry == 0):
-                continue
     else:
-                clean_result.append(entry)
+        selection = selector(frames[0])
-        result = np.array(clean_result)
+        multi_selector = _is_multi_selector(selection)
-        result = np.average(result, axis=0)
+
-    return t, result
+    result = []
+    for start_frame_index in start_frame_indices:
+        shifted_idx = logspaced_indices + start_frame_index
+        start_frame = frames[start_frame_index]
+        if selector is None:
+            selection = np.arange(len(start_frame))
+        else:
+            selection = selector(start_frame)
+        if multi_selector:
+            result_segment = _calc_correlation_multi(
+                frames, start_frame, function, selection, shifted_idx
+            )
+        else:
+            result_segment = _calc_correlation(
+                frames, start_frame, function, selection, shifted_idx
+            )
+        result.append(result_segment)
+    result = np.array(result)
+
+    if average:
+        if multi_selector:
+            result = _average_correlation_multi(result)
+        else:
+            result = _average_correlation(result)
+    return logspaced_time, result
 
 
 def msd(
@@ -184,11 +254,11 @@ def msd(
     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()
+        return (displacements[:, [0, 1]] ** 2).sum(axis=1).mean()
     elif axis == "xz" or axis == "zx":
-        return (displacements[:, [0, 2]]**2).sum(axis=1).mean()
+        return (displacements[:, [0, 2]] ** 2).sum(axis=1).mean()
     elif axis == "yz" or axis == "zy":
-        return (displacements[:, [1, 2]]**2).sum(axis=1).mean()
+        return (displacements[:, [1, 2]] ** 2).sum(axis=1).mean()
     elif axis == "x":
         return (displacements[:, 0] ** 2).mean()
     elif axis == "y":
@@ -218,13 +288,13 @@ def isf(
         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
+        distance = (displacements[:, [0, 1]] ** 2).sum(axis=1) ** 0.5
         return np.real(jn(0, distance * q)).mean()
     elif axis == "xz" or axis == "zx":
-        distance = (displacements[:, [0, 2]]**2).sum(axis=1) ** 0.5
+        distance = (displacements[:, [0, 2]] ** 2).sum(axis=1) ** 0.5
         return np.real(jn(0, distance * q)).mean()
     elif axis == "yz" or axis == "zy":
-        distance = (displacements[:, [1, 2]]**2).sum(axis=1) ** 0.5
+        distance = (displacements[:, [1, 2]] ** 2).sum(axis=1) ** 0.5
         return np.real(jn(0, distance * q)).mean()
     elif axis == "x":
         distance = np.abs(displacements[:, 0])
@@ -278,11 +348,11 @@ def van_hove_self(
     if axis == "all":
         delta_r = (vectors**2).sum(axis=1) ** 0.5
     elif axis == "xy" or axis == "yx":
-        delta_r = (vectors[:, [0, 1]]**2).sum(axis=1) ** 0.5
+        delta_r = (vectors[:, [0, 1]] ** 2).sum(axis=1) ** 0.5
     elif axis == "xz" or axis == "zx":
-        delta_r = (vectors[:, [0, 2]]**2).sum(axis=1) ** 0.5
+        delta_r = (vectors[:, [0, 2]] ** 2).sum(axis=1) ** 0.5
     elif axis == "yz" or axis == "zy":
-        delta_r = (vectors[:, [1, 2]]**2).sum(axis=1) ** 0.5
+        delta_r = (vectors[:, [1, 2]] ** 2).sum(axis=1) ** 0.5
     elif axis == "x":
         delta_r = np.abs(vectors[:, 0])
     elif axis == "y":
@@ -445,13 +515,13 @@ def non_gaussian_parameter(
         r = (vectors**2).sum(axis=1)
         dimensions = 3
     elif axis == "xy" or axis == "yx":
-        r = (vectors[:, [0, 1]]**2).sum(axis=1)
+        r = (vectors[:, [0, 1]] ** 2).sum(axis=1)
         dimensions = 2
     elif axis == "xz" or axis == "zx":
-        r = (vectors[:, [0, 2]]**2).sum(axis=1)
+        r = (vectors[:, [0, 2]] ** 2).sum(axis=1)
         dimensions = 2
     elif axis == "yz" or axis == "zy":
-        r = (vectors[:, [1, 2]]**2).sum(axis=1)
+        r = (vectors[:, [1, 2]] ** 2).sum(axis=1)
         dimensions = 2
     elif axis == "x":
         r = vectors[:, 0] ** 2