diff --git a/src/mdevaluate/coordinates.py b/src/mdevaluate/coordinates.py index 519c6df..083a16a 100755 --- a/src/mdevaluate/coordinates.py +++ b/src/mdevaluate/coordinates.py @@ -1,13 +1,14 @@ -from functools import partial, lru_cache, wraps +from functools import partial, wraps from copy import copy from .logging import logger +from typing import Optional, Callable import numpy as np -from scipy.spatial import cKDTree, KDTree +import numpy.typing as npt from .atoms import AtomSubset from .pbc import whole, nojump, pbc_diff -from .utils import mask2indices, singledispatchmethod +from .utils import singledispatchmethod from .checksum import checksum @@ -15,104 +16,6 @@ class UnknownCoordinatesMode(Exception): pass -def rotate_axis(coords, axis): - """ - Rotate a set of coordinates to a given axis. - """ - axis = np.array(axis) / np.linalg.norm(axis) - zaxis = np.array([0, 0, 1]) - if (axis == zaxis).sum() == 3: - return coords - rotation_axis = np.cross(axis, zaxis) - rotation_axis = rotation_axis / np.linalg.norm(rotation_axis) - - theta = np.arccos(axis @ zaxis / np.linalg.norm(axis)) - - # return theta/pi, rotation_axis - - ux, uy, uz = rotation_axis - cross_matrix = np.array([[0, -uz, uy], [uz, 0, -ux], [-uy, ux, 0]]) - rotation_matrix = ( - np.cos(theta) * np.identity(len(axis)) - + (1 - np.cos(theta)) - * rotation_axis.reshape(-1, 1) - @ rotation_axis.reshape(1, -1) - + np.sin(theta) * cross_matrix - ) - - if len(coords.shape) == 2: - rotated = np.array([rotation_matrix @ xyz for xyz in coords]) - else: - rotated = rotation_matrix @ coords - return rotated - - -def spherical_radius(frame, origin=None): - """ - Transform a frame of cartesian coordinates into the sperical radius. - If origin=None the center of the box is taken as the coordinates origin. - """ - if origin is None: - origin = frame.box.diagonal() / 2 - return ((frame - origin) ** 2).sum(axis=-1) ** 0.5 - - -def polar_coordinates(x, y): - """Convert cartesian to polar coordinates.""" - radius = (x**2 + y**2) ** 0.5 - phi = np.arctan2(y, x) - return radius, phi - - -def spherical_coordinates(x, y, z): - """Convert cartesian to spherical coordinates.""" - xy, phi = polar_coordinates(x, y) - radius = (x**2 + y**2 + z**2) ** 0.5 - theta = np.arccos(z / radius) - return radius, phi, theta - - -def radial_selector(frame, coordinates, rmin, rmax): - """ - Return a selection of all atoms with radius in the interval [rmin, rmax]. - """ - crd = coordinates[frame.step] - rad, _ = polar_coordinates(crd[:, 0], crd[:, 1]) - selector = (rad >= rmin) & (rad <= 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. - Coordinates are filtered by the condition:: - - rmin <= transform(frame) <= rmax - - Args: - frame: The coordinates of the actual trajectory - transform: - A function that transforms the coordinates of the frames into - the one-dimensional spatial coordinate (e.g. radius). - rmin: Minimum value of the radius - rmax: Maximum value of the radius - """ - r = transform(frame) - selector = (rmin <= r) & (rmax >= r) - return mask2indices(selector) - - class CoordinateFrame(np.ndarray): _known_modes = ("pbc", "whole", "nojump") @@ -250,7 +153,8 @@ class Coordinates: def mode(self, val): if val in CoordinateFrame._known_modes: logger.warn( - "Changing the Coordinates mode directly is deprecated. Use Coordinates.%s instead, which returns a copy.", + "Changing the Coordinates mode directly is deprecated. " + "Use Coordinates.%s instead, which returns a copy.", val, ) self._mode = val @@ -414,30 +318,111 @@ class CoordinatesMap: return CoordinatesMap(self.coordinates.pbc, self.function) -def map_coordinates(func): +def rotate_axis(coords: npt.ArrayLike, axis: npt.ArrayLike) -> np.ndarray: + """ + Rotate a set of coordinates to a given axis. + """ + axis = np.array(axis) / np.linalg.norm(axis) + zaxis = np.array([0, 0, 1]) + if (axis == zaxis).sum() == 3: + return coords + rotation_axis = np.cross(axis, zaxis) + rotation_axis = rotation_axis / np.linalg.norm(rotation_axis) + + theta = np.arccos(axis @ zaxis / np.linalg.norm(axis)) + + # return theta/pi, rotation_axis + + ux, uy, uz = rotation_axis + cross_matrix = np.array([[0, -uz, uy], [uz, 0, -ux], [-uy, ux, 0]]) + rotation_matrix = ( + np.cos(theta) * np.identity(len(axis)) + + (1 - np.cos(theta)) + * rotation_axis.reshape(-1, 1) + @ rotation_axis.reshape(1, -1) + + np.sin(theta) * cross_matrix + ) + + if len(coords.shape) == 2: + rotated = np.array([rotation_matrix @ xyz for xyz in coords]) + else: + rotated = rotation_matrix @ coords + return rotated + + +def spherical_radius( + frame: CoordinateFrame, origin: Optional[npt.ArrayLike] = None +) -> np.ndarray: + """ + Transform a frame of cartesian coordinates into the spherical radius. + If origin=None, the center of the box is taken as the coordinates' origin. + """ + if origin is None: + origin = frame.box.diagonal() / 2 + return ((frame - origin) ** 2).sum(axis=-1) ** 0.5 + + +def polar_coordinates(x: npt.ArrayLike, y: npt.ArrayLike) -> (np.ndarray, np.ndarray): + """Convert cartesian to polar coordinates.""" + radius = (x**2 + y**2) ** 0.5 + phi = np.arctan2(y, x) + return radius, phi + + +def spherical_coordinates( + x: npt.ArrayLike, y: npt.ArrayLike, z: npt.ArrayLike +) -> (np.ndarray, np.ndarray, np.ndarray): + """Convert cartesian to spherical coordinates.""" + xy, phi = polar_coordinates(x, y) + radius = (x**2 + y**2 + z**2) ** 0.5 + theta = np.arccos(z / radius) + return radius, phi, theta + + +def selector_radial_cylindrical( + atoms: CoordinateFrame, + r_min: float, + r_max: float, + origin: Optional[npt.ArrayLike] = None, +) -> np.ndarray: + 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 map_coordinates( + func: Callable[[CoordinateFrame, ...], np.ndarray] +) -> Callable[..., CoordinatesMap]: @wraps(func) - def wrapped(coordinates, **kwargs): + def wrapped(coordinates: Coordinates, **kwargs) -> CoordinatesMap: return CoordinatesMap(coordinates, partial(func, **kwargs)) return wrapped @map_coordinates -def center_of_masses(coordinates, atoms=None, shear: bool = False): +def center_of_masses( + frame: CoordinateFrame, atoms=None, shear: bool = False +) -> np.ndarray: if atoms is None: - atoms = list(range(len(coordinates))) - res_ids = coordinates.residue_ids[atoms] - masses = coordinates.masses[atoms] + atoms = list(range(len(frame))) + res_ids = frame.residue_ids[atoms] + masses = frame.masses[atoms] if shear: - coords = coordinates[atoms] - box = coords.box + coords = frame[atoms] + 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 = coordinates.whole[atoms] + coords = frame.whole[atoms] mask = np.bincount(res_ids)[1:] != 0 positions = np.array( [ @@ -450,12 +435,14 @@ def center_of_masses(coordinates, atoms=None, shear: bool = False): @map_coordinates -def pore_coordinates(coordinates, origin, sym_axis="z"): +def pore_coordinates( + frame: CoordinateFrame, origin: npt.ArrayLike, sym_axis: str = "z" +) -> np.ndarray: """ Map coordinates of a pore simulation so the pore has cylindrical symmetry. Args: - coordinates: Coordinates of the simulation + frame: Coordinates of the simulation origin: Origin of the pore which will be the coordinates origin after mapping sym_axis (opt.): Symmtery axis of the pore, may be a literal direction 'x', 'y' or 'z' or an array of shape (3,) @@ -465,30 +452,34 @@ def pore_coordinates(coordinates, origin, sym_axis="z"): rot_axis[["x", "y", "z"].index(sym_axis)] = 1 else: rot_axis = sym_axis - - return rotate_axis(coordinates - origin, rot_axis) + return rotate_axis(frame - origin, rot_axis) @map_coordinates -def vectors(coordinates, atoms_a, atoms_b, normed=False): +def vectors( + frame: CoordinateFrame, + atoms_indices_a: npt.ArrayLike[int], + atoms_indices_b: npt.ArrayLike[int], + normed: bool = False, +) -> np.ndarray: """ Compute the vectors between the atoms of two subsets. Args: - coordinates: The Coordinates object the atoms will be taken from - atoms_a: Mask or indices of the first atom subset - atoms_b: Mask or indices of the second atom subset + frame: The Coordinates object the atoms will be taken from + atoms_indices_a: Mask or indices of the first atom subset + atoms_indices_b: Mask or indices of the second atom subset normed (opt.): If the vectors should be normed - box (opt.): If not None, the vectors are calcualte with PBC - The defintion of atoms_a/b can be any possible subript of a numpy array. + The definition of atoms_a/b can be any possible subript of a numpy array. They can, for example, be given as a masking array of bool values with the - same length as the frames of the coordinates. Or they can be a list of - indices selecting the atoms of these indices from each frame. + same length as the frames of the coordinates. + Or there can be a list of indices selecting the atoms of these indices from each + frame. - It is possible to compute the mean of several atoms before calculating the vectors, + It is possible to compute the means of several atoms before calculating the vectors, by using a two-dimensional list of indices. The following code computes the vectors - between atoms 0, 3, 6 and the mean coordinate of atoms 1, 4, 7 and 2, 5, 8:: + between atoms 0, 3, 6 and the mean coordinate of atoms 1, 4, 7 and 2, 5, 8: >>> inds_a = [0, 3, 6] >>> inds_b = [[1, 4, 7], [2, 5, 8]] @@ -499,11 +490,11 @@ def vectors(coordinates, atoms_a, atoms_b, normed=False): coords[6] - (coords[7] + coords[8])/2, ]) """ - box = coordinates.box - coords_a = coordinates[atoms_a] + box = frame.box + coords_a = frame[atoms_indices_a] if len(coords_a.shape) > 2: coords_a = coords_a.mean(axis=0) - coords_b = coordinates[atoms_b] + coords_b = frame[atoms_indices_b] if len(coords_b.shape) > 2: coords_b = coords_b.mean(axis=0) vec = pbc_diff(coords_a, coords_b, box=box) @@ -514,10 +505,12 @@ def vectors(coordinates, atoms_a, atoms_b, normed=False): @map_coordinates -def dipole_vector(coordinates, atoms, normed=None): - coords = coordinates.whole[atoms] - res_ids = coordinates.residue_ids[atoms] - charges = coordinates.charges[atoms] +def dipole_vector( + frame: CoordinateFrame, atom_indices: npt.ArrayLike[int], normed: bool = None +) -> np.ndarray: + coords = frame.whole[atom_indices] + res_ids = frame.residue_ids[atom_indices] + charges = frame.charges[atom_indices] mask = np.bincount(res_ids)[1:] != 0 dipoles = np.array( [np.bincount(res_ids, weights=c * charges)[1:] for c in coords.T] @@ -531,11 +524,11 @@ def dipole_vector(coordinates, atoms, normed=None): @map_coordinates def sum_dipole_vector( coordinates: CoordinateFrame, - atoms: list[int], + atom_indices: npt.ArrayLike[int], normed: bool = True, -): - coords = coordinates.whole[atoms] - charges = coordinates.charges[atoms] +) -> np.ndarray: + coords = coordinates.whole[atom_indices] + charges = coordinates.charges[atom_indices] dipole = np.array([c * charges for c in coords.T]).T if normed: dipole /= np.linalg.norm(dipole) @@ -544,16 +537,16 @@ def sum_dipole_vector( @map_coordinates def normal_vectors( - coordinates: CoordinateFrame, - atoms_a: list[int], - atoms_b: list[int], - atoms_c: list[int], + frame: CoordinateFrame, + atom_indices_a: npt.ArrayLike[int], + atom_indices_b: npt.ArrayLike[int], + atom_indices_c: npt.ArrayLike[int], normed: bool = True, -): - coords_a = coordinates[atoms_a] - coords_b = coordinates[atoms_b] - coords_c = coordinates[atoms_c] - box = coordinates.box +) -> np.ndarray: + coords_a = frame[atom_indices_a] + coords_b = frame[atom_indices_b] + coords_c = frame[atom_indices_c] + box = frame.box vectors_a = pbc_diff(coords_a, coords_b, box=box) vectors_b = pbc_diff(coords_a, coords_c, box=box) vec = np.cross(vectors_a, vectors_b) @@ -576,7 +569,9 @@ def displacements_without_drift( @map_coordinates -def cylindrical_coordinates(frame, origin=None): +def cylindrical_coordinates( + frame: CoordinateFrame, origin: npt.ArrayLike = None +) -> np.ndarray: if origin is None: origin = np.diag(frame.box) / 2 x = frame[:, 0] - origin[0]