Added type hints, adjusted imports, and doc strings
This commit is contained in:
parent
c00fc78f23
commit
eed94cfca6
@ -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]
|
||||
|
||||
Loading…
Reference in New Issue
Block a user