diff --git a/mdevaluate/correlation.py b/mdevaluate/correlation.py index 03c0cfc..5f75072 100644 --- a/mdevaluate/correlation.py +++ b/mdevaluate/correlation.py @@ -81,12 +81,52 @@ def multi_subensemble_correlation(selector_function): @autosave_data(2) -def shifted_correlation(function, frames, selector=None, segments=100, +def shifted_correlation(function, frames, selector=None, segments=10, skip=0.1, window=0.5, average=True, points=100, nodes=8): - - - + """ + 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 + nodes (int, opt.): + Number of nodes used for parallelization + + 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 coordinates object named ``coords``: + + >>> time, data = shifted_correlation(msd, coords) + """ def get_correlation(start_frame, idx, selector=None): shifted_idx = idx + start_frame if selector: @@ -137,126 +177,6 @@ def shifted_correlation(function, frames, selector=None, segments=100, return t, result -@autosave_data(nargs=2, kwargs_keys=( - 'index_distribution', 'correlation', 'segments', 'window', 'skip', 'average' -), version='shifted_correlation-1') -def shifted_correlation(function, frames, - index_distribution=log_indices, correlation=correlation, - segments=10, window=0.5, skip=None, - average=False, ): - - """ - Calculate the time series for a correlation function. - - The times at which the correlation is calculated are determined automatically by the - function given as ``index_distribution``. The default is a logarithmic distribution. - - Args: - function: The function that should be correlated - frames: The coordinates of the simulation data - index_distribution (opt.): - A function that returns the indices for which the timeseries - will be calculated - correlation (function, opt.): - The correlation function - segments (int, opt.): - The number of segments the time window will be shifted - window (float, opt.): - The fraction of the simulation the time series will cover - 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. - counter (bool, opt.): - If True, returns length of frames (in general number of particles specified) - average (bool, opt.): - If True, returns averaged correlation function - Returns: - tuple: - A list of length N that contains the indices 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 - - if has_counter == True: adds number of counts to output tupel. - if average is returned it will be weighted. - - Example: - Calculating the mean square displacement of a coordinates object named ``coords``: - - >>> indices, data = shifted_correlation(msd, coords) - """ - if skip is None: - try: - skip = frames._slice.start / len(frames) - except (TypeError, AttributeError): - skip = 0 - assert window + skip < 1 - - start_frames = np.unique(np.linspace( - len(frames) * skip, len(frames) * (1 - window), - num=segments, endpoint=False, dtype=int - )) - num_frames = int(len(frames) * (window)) - - idx = index_distribution(0, num_frames) - - - def correlate(start_frame): - shifted_idx = idx + start_frame - return correlation(function, map(frames.__getitem__, shifted_idx)) - - times = np.array([frames[i].time for i in idx]) - frames[0].time - - if getattr(correlation, "has_counter", False): - if average: - for i, start_frame in enumerate(start_frames): - act_result, act_count = correlate(start_frame) - act_result = np.array(list(act_result)) - act_count = np.array(act_count) - if i == 0: - count = act_count - cdim = act_count.ndim - rdim = act_result.ndim - bt = np.newaxis, - for i in range(rdim - 1): - if i >= cdim: - bt += np.newaxis, - else: - bt += slice(None), - result = act_result * act_count[bt] - else: - result += act_result * act_count[bt] - count += act_count - np.divide(result, count[bt], out = result, where = count[bt] != 0) - result = np.moveaxis(result,0,cdim) - count = count / len(start_frames) - output = times, result, count - else: - count = [] - result = [] - for i, start_frame in enumerate(start_frames): - act_result, act_count = correlate(start_frame) - act_result = list(act_result) - result.append(act_result) - count.append(act_count) - count = np.asarray(count) - cdim = count.ndim - result = np.asarray(result) - result = np.moveaxis(result,1,cdim) - output = times, result, count - else: - result = 0 if average else [] - for i, start_frame in enumerate(start_frames): - if average: - result += np.array(list(correlate(start_frame))) - else: - result.append(list(correlate(start_frame))) - result = np.array(result) - if average: - result = result / len(start_frames) - output = times, result - return output - def msd(start, frame): """ Mean square displacement