diff --git a/mdevaluate/correlation.py b/mdevaluate/correlation.py index 510ff43..03c0cfc 100644 --- a/mdevaluate/correlation.py +++ b/mdevaluate/correlation.py @@ -79,6 +79,64 @@ def multi_subensemble_correlation(selector_function): return map(cc, chain([start_frame], iterator)), count return cmulti + +@autosave_data(2) +def shifted_correlation(function, frames, selector=None, segments=100, + skip=0.1, window=0.5, average=True, points=100, + nodes=8): + + + + def get_correlation(start_frame, idx, selector=None): + shifted_idx = idx + start_frame + if selector: + index = selector(frames[start_frame]) + else: + index = np.arange(len(frames[start_frame])) + if len(index) == 0: + return np.zeros(len(idx)) + else: + start = frames[start_frame][index] + correlation = np.array([ function(start, frames[frame][index]) + for frame in shifted_idx ]) + return correlation + + if 1-skip < window: + window = 1-skip + + 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) + ls = np.logspace(0, np.log10(num_frames + 1), num=points) + idx = np.unique(np.int_(ls) - 1) + t = np.array([frames[i].time for i in idx]) - frames[0].time + + pool = ProcessPool(nodes=nodes) + try: + result = np.array(pool.map(partial(get_correlation, idx=idx, + selector=selector), + start_frames)) + except Exception: + print("Something went wrong!") + finally: + pool.terminate() + pool.restart() + + if average == True: + clean_result = [] + for entry in result: + if np.all(entry==0): + continue + else: + clean_result.append(entry) + result = np.array(clean_result) + result = np.average(result, axis=0) + return t, result + + @autosave_data(nargs=2, kwargs_keys=( 'index_distribution', 'correlation', 'segments', 'window', 'skip', 'average' ), version='shifted_correlation-1') diff --git a/mdevaluate/pbc.py b/mdevaluate/pbc.py index 54244e7..90165f7 100644 --- a/mdevaluate/pbc.py +++ b/mdevaluate/pbc.py @@ -138,17 +138,6 @@ def pbc_backfold_compact(act_frame,box_matrix): return sc[range(shape[0]),w] -# Parameter used to switch reference position for whole molecules -# 'com': use center of mass of each molecule -# 'simple': use first atom in each molecule -WHOLEMODE = 'com' - -fname = os.path.expanduser('~/.mdevaluate/WHOLEMODE') -if os.path.exists(fname): - with open(fname) as f: - WHOLEMODE = f.read().strip() - logger.info('Setting WHOLEMODE to %s, according to file ~/.mdevaluate/WHOLEMODE', WHOLEMODE) - def whole(frame): """ Apply ``-pbc whole`` to a CoordinateFrame. @@ -156,22 +145,13 @@ def whole(frame): residue_ids = frame.coordinates.atom_subset.residue_ids box = frame.box.diagonal() - if WHOLEMODE == 'com': - logger.debug('Using COM as reference for whole.') - coms = np.array([ - np.bincount(residue_ids, weights=c * frame.masses)[1:] / np.bincount(residue_ids, weights=frame.masses)[1:] - for c in frame.T - ]).T[residue_ids - 1] - - else: - # make sure, residue_ids are sorted, then determine indices at which the res id changes - # kind='stable' assures that any existent ordering is preserved - logger.debug('Using first atom as reference for whole.') - sort_ind = residue_ids.argsort(kind='stable') - i = np.concatenate([[0], np.where(np.diff(residue_ids[sort_ind]) > 0)[0] + 1]) - coms = frame[sort_ind[i]][residue_ids - 1] + # make sure, residue_ids are sorted, then determine indices at which the res id changes + # kind='stable' assures that any existent ordering is preserved + logger.debug('Using first atom as reference for whole.') + sort_ind = residue_ids.argsort(kind='stable') + i = np.concatenate([[0], np.where(np.diff(residue_ids[sort_ind]) > 0)[0] + 1]) + coms = frame[sort_ind[i]][residue_ids - 1] - cor = np.zeros_like(frame) cd = frame - coms n, d = np.where(cd > box / 2 * 0.9) @@ -179,13 +159,6 @@ def whole(frame): n, d = np.where(cd < -box / 2 * 0.9) cor[n, d] = box[d] - # this fix is only necessary when COM is the reference - if WHOLEMODE == 'com': - duomask = np.bincount(residue_ids)[1:][residue_ids - 1] == 2 - if np.any(duomask): - duomask[::2] = False - cor[duomask] = 0 - return frame + cor