diff --git a/src/mdevaluate/coordinates.py b/src/mdevaluate/coordinates.py index ae683f4..bf8e5d9 100755 --- a/src/mdevaluate/coordinates.py +++ b/src/mdevaluate/coordinates.py @@ -413,6 +413,7 @@ def center_of_masses( ) -> NDArray: if atom_indices is None: atom_indices = list(range(len(frame))) + print(type(frame)) res_ids = frame.residue_ids[atom_indices] masses = frame.masses[atom_indices] if shear: @@ -434,6 +435,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