11 Commits

4 changed files with 42 additions and 22 deletions

View File

@@ -166,8 +166,10 @@ def autosave_data(
@functools.wraps(function)
def autosave(*args, **kwargs):
description = kwargs.pop("description", "")
autosave_dir_overwrite = kwargs.pop("autosave_dir_overwrite", None)
autosave_dir = autosave_dir_overwrite if autosave_dir_overwrite is not None else autosave_directory
autoload = kwargs.pop("autoload", True) and load_autosave_data
if autosave_directory is not None:
if autosave_dir is not None:
relevant_args = list(args[:nargs])
if kwargs_keys is not None:
for key in [*posargs_keys, *kwargs_keys]:

View File

@@ -434,6 +434,34 @@ def center_of_masses(
]
).T[mask]
return np.array(positions)
@map_coordinates
def center_of_atoms(
frame: CoordinateFrame, atom_indices=None, shear: bool = False
) -> NDArray:
if atom_indices is None:
atom_indices = list(range(len(frame)))
res_ids = frame.residue_ids[atom_indices]
if shear:
coords = frame[atom_indices]
box = frame.box
sort_ind = res_ids.argsort(kind="stable")
i = np.concatenate([[0], np.where(np.diff(res_ids[sort_ind]) > 0)[0] + 1])
coms = coords[sort_ind[i]][res_ids - min(res_ids)]
cor = pbc_diff(coords, coms, box)
coords = coms + cor
else:
coords = frame.whole[atom_indices]
mask = np.bincount(res_ids)[1:] != 0
positions = np.array(
[
np.bincount(res_ids, weights=c)[1:]
/ np.bincount(res_ids)[1:]
for c in coords.T
]
).T[mask]
return np.array(positions)
@map_coordinates

View File

@@ -147,7 +147,8 @@ def shifted_correlation(
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
dt = round(frames[1].time - frames[0].time, 6) # round to avoid bad floats
t = idx * dt
result = np.array(
[

View File

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