quick version of center_of_masses for equal weights

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