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