Changed time_average to work similar to shifted_correlation

This commit is contained in:
Sebastian Kloth 2023-12-26 16:57:53 +01:00
parent 32e3541c4b
commit cc3768925a

View File

@ -1,3 +1,5 @@
from typing import Callable, Optional
import numpy as np
from numpy.typing import ArrayLike
from scipy import spatial
@ -17,8 +19,14 @@ from .pbc import pbc_diff, pbc_points
from .logging import logger
@autosave_data(nargs=2, kwargs_keys=("coordinates_b",), version="time_average-1")
def time_average(function, coordinates, coordinates_b=None, pool=None):
@autosave_data(nargs=2, kwargs_keys=("coordinates_b",))
def time_average(
function: Callable,
coordinates: Coordinates,
coordinates_b: Optional[Coordinates] = None,
skip: float = 0.1,
segments: int = 100,
) -> np.ndarray:
"""
Compute the time average of a function.
@ -27,35 +35,23 @@ def time_average(function, coordinates, coordinates_b=None, pool=None):
The function that will be averaged, it has to accept exactly one argument
which is the current atom set
coordinates: The coordinates object of the simulation
pool (multiprocessing.Pool, opt.):
A multiprocessing pool which will be used for cocurrent calculation of the
averaged function
coordinates_b: Additional coordinates object of the simulation
skip:
segments:
"""
if pool is not None:
_map = pool.imap
frame_indices = np.unique(
np.int_(
np.linspace(len(coordinates) * skip, len(coordinates) - 1, num=segments)
)
)
if coordinates_b is None:
result = [function(coordinates[frame_index]) for frame_index in frame_indices]
else:
_map = map
number_of_averages = 0
result = 0
if coordinates_b is not None:
if coordinates._slice != coordinates_b._slice:
logger.warning("Different slice for coordinates and coordinates_b.")
coordinate_iter = (iter(coordinates), iter(coordinates_b))
else:
coordinate_iter = (iter(coordinates),)
evaluated = _map(function, *coordinate_iter)
for ev in evaluated:
number_of_averages += 1
result += ev
if number_of_averages % 100 == 0:
logger.debug("time_average: %d", number_of_averages)
return result / number_of_averages
result = [
function(coordinates[frame_index], coordinates_b[frame_index])
for frame_index in frame_indices
]
return np.mean(result, axis=0)
def time_histogram(function, coordinates, bins, hist_range, pool=None):