IPKM/mdevaluate
6
1
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Changed shifted_correlation to support parallelization and easy input of selectors

Browse Source
This commit is contained in:
sebastiankloth 2022-11-03 15:42:58 +01:00
parent 9628053d26
commit 3992dd4e07
1 changed files with 44 additions and 124 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

168
mdevaluate/correlation.py
View File

@ -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