Changed whole to only support simple

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')

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