From 118ed8e6af321d65c5c7131cd3886c300cac9882 Mon Sep 17 00:00:00 2001
From: Sebastian Kloth <sebastian.kloth@pkm.tu-darmstadt.de>
Date: Mon, 6 Nov 2023 15:53:56 +0100
Subject: [PATCH] Added missing functions from mdeval_skloth

---
 src/mdevaluate/coordinates.py | 17 +++++++++++++++--
 src/mdevaluate/utils.py       | 14 ++++++++++++++
 2 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/src/mdevaluate/coordinates.py b/src/mdevaluate/coordinates.py
index 0418bff..3cd737f 100755
--- a/src/mdevaluate/coordinates.py
+++ b/src/mdevaluate/coordinates.py
@@ -82,6 +82,17 @@ def radial_selector(frame, coordinates, rmin, rmax):
     return mask2indices(selector)
 
 
+def selector_radial_cylindrical(atoms, r_min, r_max, origin=None):
+    box = atoms.box
+    atoms = atoms % np.diag(box)
+    if origin is None:
+        origin = [box[0, 0] / 2, box[1, 1] / 2, box[2, 2] / 2]
+    r_vec = (atoms - origin)[:, :2]
+    r = np.linalg.norm(r_vec, axis=1)
+    index = np.argwhere((r >= r_min) * (r < r_max))
+    return index.flatten()
+
+
 def spatial_selector(frame, transform, rmin, rmax):
     """
     Select a subset of atoms which have a radius between rmin and rmax.
@@ -411,7 +422,7 @@ def map_coordinates(func):
 
 
 @map_coordinates
-def center_of_masses(coordinates, atoms, shear: bool = False):
+def center_of_masses(coordinates, atoms=None, shear: bool = False):
     """
     Example:
     rd = XTCReader('t.xtc')
@@ -419,6 +430,8 @@ def center_of_masses(coordinates, atoms, shear: bool = False):
     com = centers_of_mass(coordinates, (1.0, 2.0, 1.0, 3.0))
 
     """
+    if atoms is None:
+        atoms = list(range(len(coordinates)))
     res_ids = coordinates.residue_ids[atoms]
     masses = coordinates.masses[atoms]
     if shear:
@@ -427,7 +440,7 @@ def center_of_masses(coordinates, atoms, shear: bool = False):
         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 = md.pbc.pbc_diff(coords, coms, box)
+        cor = pbc_diff(coords, coms, box)
         coords = coms + cor
     else:
         coords = coordinates.whole[atoms]
diff --git a/src/mdevaluate/utils.py b/src/mdevaluate/utils.py
index 0ad26c3..42ad521 100644
--- a/src/mdevaluate/utils.py
+++ b/src/mdevaluate/utils.py
@@ -2,6 +2,7 @@
 Collection of utility functions.
 """
 import functools
+from time import time
 from types import FunctionType
 
 import numpy as np
@@ -486,3 +487,16 @@ def fibonacci_sphere(samples=1000):
         points.append((x, y, z))
 
     return np.array(points)
+
+
+def timing(function):
+    @functools.wraps(function)
+    def wrap(*args, **kw):
+        start_time = time()
+        result = function(*args, **kw)
+        end_time = time()
+        time_needed = end_time - start_time
+        print(f"Finished in {int(time_needed // 60)} min " f"{int(time_needed % 60)} s")
+        return result
+
+    return wrap
\ No newline at end of file