diff --git a/src/mdevaluate/pbc.py b/src/mdevaluate/pbc.py index b44f083..785d9d2 100644 --- a/src/mdevaluate/pbc.py +++ b/src/mdevaluate/pbc.py @@ -149,32 +149,21 @@ def nojump(frame: CoordinateFrame, usecache: bool = True) -> CoordinateFrame: i0 = 0 delta = 0 - delta = ( - delta - + np.array( - np.vstack( - [m[i0 : abstep + 1].sum(axis=0) for m in reader.nojump_matrices] - ).T - ) - @ frame.box - ) + delta = (delta + + np.vstack( + [m[i0 : abstep + 1].sum(axis=0) for m in reader.nojump_matrices] + ).T) reader._nojump_cache[abstep] = delta while len(reader._nojump_cache) > NOJUMP_CACHESIZE: reader._nojump_cache.popitem(last=False) - delta = delta[selection, :] else: - delta = ( - np.array( - np.vstack( - [ - m[: frame.step + 1, selection].sum(axis=0) - for m in reader.nojump_matrices - ] + delta = np.vstack( + [m[: frame.step + 1, selection].sum(axis=0) for m in reader.nojump_matrices] ).T - ) - @ frame.box - ) + + delta = delta[selection, :] + delta = np.array(delta @ frame.box) return frame - delta