441 lines
15 KiB
Python
441 lines
15 KiB
Python
from typing import Callable, Optional, Union, Tuple, List
|
|
|
|
import numpy as np
|
|
from numpy.typing import ArrayLike, NDArray
|
|
from scipy import spatial
|
|
from scipy.spatial import KDTree
|
|
from scipy.sparse.csgraph import connected_components
|
|
|
|
from .coordinates import (
|
|
rotate_axis,
|
|
polar_coordinates,
|
|
Coordinates,
|
|
CoordinateFrame,
|
|
next_neighbors,
|
|
number_of_neighbors,
|
|
)
|
|
from .autosave import autosave_data
|
|
from .pbc import pbc_diff, pbc_points
|
|
|
|
|
|
@autosave_data(nargs=2, kwargs_keys=("coordinates_b",))
|
|
def time_average(
|
|
function: Callable,
|
|
coordinates: Coordinates,
|
|
coordinates_b: Optional[Coordinates] = None,
|
|
skip: float = 0.1,
|
|
segments: int = 100,
|
|
) -> NDArray:
|
|
"""
|
|
Compute the time average of a function.
|
|
|
|
Args:
|
|
function:
|
|
The function that will be averaged, it has to accept exactly one argument
|
|
which is the current atom set (or two if coordinates_b is provided)
|
|
coordinates: The coordinates object of the simulation
|
|
coordinates_b: Additional coordinates object of the simulation
|
|
skip:
|
|
segments:
|
|
"""
|
|
frame_indices = np.unique(
|
|
np.int_(
|
|
np.linspace(len(coordinates) * skip, len(coordinates) - 1, num=segments)
|
|
)
|
|
)
|
|
if coordinates_b is None:
|
|
result = [function(coordinates[frame_index]) for frame_index in frame_indices]
|
|
else:
|
|
result = [
|
|
function(coordinates[frame_index], coordinates_b[frame_index])
|
|
for frame_index in frame_indices
|
|
]
|
|
return np.mean(result, axis=0)
|
|
|
|
|
|
@autosave_data(nargs=2, kwargs_keys=("coordinates_b",))
|
|
def time_distribution(
|
|
function: Callable,
|
|
coordinates: Coordinates,
|
|
coordinates_b: Optional[Coordinates] = None,
|
|
skip: float = 0,
|
|
segments: int = 100,
|
|
) -> Tuple[NDArray, List]:
|
|
"""
|
|
Compute the time distribution of a function.
|
|
|
|
Args:
|
|
function:
|
|
The function that will be averaged, it has to accept exactly one argument
|
|
which is the current atom set (or two if coordinates_b is provided)
|
|
coordinates: The coordinates object of the simulation
|
|
coordinates_b: Additional coordinates object of the simulation
|
|
skip:
|
|
segments:
|
|
"""
|
|
frame_indices = np.unique(
|
|
np.int_(
|
|
np.linspace(len(coordinates) * skip, len(coordinates) - 1, num=segments)
|
|
)
|
|
)
|
|
times = np.array([coordinates[frame_index].time for frame_index in frame_indices])
|
|
if coordinates_b is None:
|
|
result = [function(coordinates[frame_index]) for frame_index in frame_indices]
|
|
else:
|
|
result = [
|
|
function(coordinates[frame_index], coordinates_b[frame_index])
|
|
for frame_index in frame_indices
|
|
]
|
|
return times, result
|
|
|
|
|
|
def rdf(
|
|
atoms_a: CoordinateFrame,
|
|
atoms_b: Optional[CoordinateFrame] = None,
|
|
bins: Optional[ArrayLike] = None,
|
|
remove_intra: bool = False,
|
|
**kwargs
|
|
) -> NDArray:
|
|
r"""
|
|
Compute the radial pair distribution of one or two sets of atoms.
|
|
|
|
.. math::
|
|
g_{AB}(r) = \frac{1}{\langle \rho_B\rangle N_A}\sum\limits_{i\in A}^{N_A}
|
|
\sum\limits_{j\in B}^{N_B}\frac{\delta(r_{ij} -r)}{4\pi r^2}
|
|
|
|
For use with :func:`time_average`, define bins through the use of
|
|
:func:`~functools.partial`, the atom sets are passed to :func:`time_average`, if a
|
|
second set of atoms should be used specify it as ``coordinates_b`` and it will be
|
|
passed to this function.
|
|
|
|
Args:
|
|
atoms_a: First set of atoms, used internally
|
|
atoms_b (opt.): Second set of atoms, used internal
|
|
bins: Bins of the radial distribution function
|
|
remove_intra: removes contributions from intra molecular pairs
|
|
"""
|
|
distinct = True
|
|
if atoms_b is None:
|
|
atoms_b = atoms_a
|
|
distinct = False
|
|
elif np.array_equal(atoms_a, atoms_b):
|
|
distinct = False
|
|
if bins is None:
|
|
bins = np.arange(0, 1, 0.01)
|
|
|
|
particles_in_volume = int(
|
|
np.max(number_of_neighbors(atoms_a, query_atoms=atoms_b, r_max=bins[-1])) * 1.1
|
|
)
|
|
distances, indices = next_neighbors(
|
|
atoms_a,
|
|
atoms_b,
|
|
number_of_neighbors=particles_in_volume,
|
|
distance_upper_bound=bins[-1],
|
|
distinct=distinct,
|
|
**kwargs
|
|
)
|
|
|
|
if remove_intra:
|
|
new_distances = []
|
|
for entry in list(zip(atoms_a.residue_ids, distances, indices)):
|
|
mask = entry[1] < np.inf
|
|
new_distances.append(
|
|
entry[1][mask][atoms_b.residue_ids[entry[2][mask]] != entry[0]]
|
|
)
|
|
distances = np.concatenate(new_distances)
|
|
else:
|
|
distances = [d for dist in distances for d in dist]
|
|
|
|
hist, bins = np.histogram(distances, bins=bins, range=(0, bins[-1]), density=False)
|
|
hist = hist / len(atoms_b)
|
|
hist = hist / (4 / 3 * np.pi * bins[1:] ** 3 - 4 / 3 * np.pi * bins[:-1] ** 3)
|
|
n = len(atoms_a) / np.prod(np.diag(atoms_a.box))
|
|
hist = hist / n
|
|
|
|
return hist
|
|
|
|
|
|
def distance_distribution(
|
|
atoms: CoordinateFrame, bins: Union[int, ArrayLike]
|
|
) -> NDArray:
|
|
connection_vectors = atoms[:-1, :] - atoms[1:, :]
|
|
connection_lengths = (connection_vectors**2).sum(axis=1) ** 0.5
|
|
return np.histogram(connection_lengths, bins)[0]
|
|
|
|
|
|
def tetrahedral_order(
|
|
atoms: CoordinateFrame, reference_atoms: CoordinateFrame = None
|
|
) -> NDArray:
|
|
if reference_atoms is None:
|
|
reference_atoms = atoms
|
|
indices = next_neighbors(
|
|
reference_atoms,
|
|
query_atoms=atoms,
|
|
number_of_neighbors=4,
|
|
)[1]
|
|
neighbors = reference_atoms[indices]
|
|
neighbors_1, neighbors_2, neighbors_3, neighbors_4 = (
|
|
neighbors[:, 0, :],
|
|
neighbors[:, 1, :],
|
|
neighbors[:, 2, :],
|
|
neighbors[:, 3, :],
|
|
)
|
|
|
|
# Connection vectors
|
|
neighbors_1 = pbc_diff(neighbors_1, atoms, box=atoms.box)
|
|
neighbors_2 = pbc_diff(neighbors_2, atoms, box=atoms.box)
|
|
neighbors_3 = pbc_diff(neighbors_3, atoms, box=atoms.box)
|
|
neighbors_4 = pbc_diff(neighbors_4, atoms, box=atoms.box)
|
|
|
|
# Normed Connection vectors
|
|
neighbors_1 /= np.linalg.norm(neighbors_1, axis=-1).reshape(-1, 1)
|
|
neighbors_2 /= np.linalg.norm(neighbors_2, axis=-1).reshape(-1, 1)
|
|
neighbors_3 /= np.linalg.norm(neighbors_3, axis=-1).reshape(-1, 1)
|
|
neighbors_4 /= np.linalg.norm(neighbors_4, axis=-1).reshape(-1, 1)
|
|
|
|
a_1_2 = ((neighbors_1 * neighbors_2).sum(axis=1) + 1 / 3) ** 2
|
|
a_1_3 = ((neighbors_1 * neighbors_3).sum(axis=1) + 1 / 3) ** 2
|
|
a_1_4 = ((neighbors_1 * neighbors_4).sum(axis=1) + 1 / 3) ** 2
|
|
|
|
a_2_3 = ((neighbors_2 * neighbors_3).sum(axis=1) + 1 / 3) ** 2
|
|
a_2_4 = ((neighbors_2 * neighbors_4).sum(axis=1) + 1 / 3) ** 2
|
|
|
|
a_3_4 = ((neighbors_3 * neighbors_4).sum(axis=1) + 1 / 3) ** 2
|
|
|
|
q = 1 - 3 / 8 * (a_1_2 + a_1_3 + a_1_4 + a_2_3 + a_2_4 + a_3_4)
|
|
|
|
return q
|
|
|
|
|
|
def tetrahedral_order_distribution(
|
|
atoms: CoordinateFrame,
|
|
reference_atoms: Optional[CoordinateFrame] = None,
|
|
bins: Optional[ArrayLike] = None,
|
|
) -> NDArray:
|
|
assert bins is not None, "Bin edges of the distribution have to be specified."
|
|
Q = tetrahedral_order(atoms, reference_atoms=reference_atoms)
|
|
return np.histogram(Q, bins=bins)[0]
|
|
|
|
|
|
def radial_density(
|
|
atoms: CoordinateFrame,
|
|
bins: Optional[ArrayLike] = None,
|
|
symmetry_axis: ArrayLike = (0, 0, 1),
|
|
origin: Optional[ArrayLike] = None,
|
|
) -> NDArray:
|
|
"""
|
|
Calculate the radial density distribution.
|
|
|
|
This function is meant to be used with time_average.
|
|
|
|
Args:
|
|
atoms:
|
|
Set of coordinates.
|
|
bins (opt.):
|
|
Bin specification that is passed to numpy.histogram. This needs to be
|
|
a list of bin edges if the function is used within time_average.
|
|
symmetry_axis (opt.):
|
|
Vector of the symmetry axis, around which the radial density is calculated,
|
|
default is z-axis.
|
|
origin (opt.):
|
|
Origin of the rotational symmetry, e.g. center of the pore.
|
|
"""
|
|
if origin is None:
|
|
origin = np.diag(atoms.box) / 2
|
|
if bins is None:
|
|
bins = np.arange(0, np.min(np.diag(atoms.box) / 2), 0.01)
|
|
length = np.diag(atoms.box) * symmetry_axis
|
|
cartesian = rotate_axis(atoms - origin, symmetry_axis)
|
|
radius, _ = polar_coordinates(cartesian[:, 0], cartesian[:, 1])
|
|
hist = np.histogram(radius, bins=bins)[0]
|
|
volume = np.pi * (bins[1:] ** 2 - bins[:-1] ** 2) * length
|
|
return hist / volume
|
|
|
|
|
|
def shell_density(
|
|
atoms: CoordinateFrame,
|
|
shell_radius: float,
|
|
bins: ArrayLike,
|
|
shell_thickness: float = 0.5,
|
|
symmetry_axis: ArrayLike = (0, 0, 1),
|
|
origin: Optional[ArrayLike] = None,
|
|
) -> NDArray:
|
|
"""
|
|
Compute the density distribution on a cylindrical shell.
|
|
|
|
Args:
|
|
atoms: The coordinates of the atoms
|
|
shell_radius: Inner radius of the shell
|
|
bins: Histogramm bins, this has to be a two-dimensional list of bins: [angle, z]
|
|
shell_thickness (opt.): Thicknes of the shell, default is 0.5
|
|
symmetry_axis (opt.): The symmtery axis of the pore, the coordinates will be
|
|
rotated such that this axis is the z-axis
|
|
origin (opt.): Origin of the pore, the coordinates will be moved such that this
|
|
is the new origin.
|
|
|
|
Returns:
|
|
Two-dimensional density distribution of the atoms in the defined shell.
|
|
"""
|
|
if origin is None:
|
|
origin = np.diag(atoms.box) / 2
|
|
cartesian = rotate_axis(atoms - origin, symmetry_axis)
|
|
radius, theta = polar_coordinates(cartesian[:, 0], cartesian[:, 1])
|
|
shell_indices = (shell_radius <= radius) & (
|
|
radius <= shell_radius + shell_thickness
|
|
)
|
|
hist = np.histogram2d(theta[shell_indices], cartesian[shell_indices, 2], bins)[0]
|
|
|
|
return hist
|
|
|
|
|
|
def next_neighbor_distribution(
|
|
atoms: CoordinateFrame,
|
|
reference: Optional[CoordinateFrame] = None,
|
|
number_of_neighbors: int = 4,
|
|
bins: Optional[ArrayLike] = None,
|
|
normed: bool = True,
|
|
) -> NDArray:
|
|
"""
|
|
Compute the distribution of next neighbors with the same residue name.
|
|
"""
|
|
assert bins is not None, "Bins have to be specified."
|
|
if reference is None:
|
|
reference = atoms
|
|
nn = next_neighbors(
|
|
reference, query_atoms=atoms, number_of_neighbors=number_of_neighbors
|
|
)[1]
|
|
resname_nn = reference.residue_names[nn]
|
|
count_nn = (resname_nn == atoms.residue_names.reshape(-1, 1)).sum(axis=1)
|
|
return np.histogram(count_nn, bins=bins, density=normed)[0]
|
|
|
|
|
|
def hbonds(
|
|
atoms: CoordinateFrame,
|
|
donator_indices: ArrayLike,
|
|
hydrogen_indices: ArrayLike,
|
|
acceptor_indices: ArrayLike,
|
|
DA_lim: float = 0.35,
|
|
HA_lim: float = 0.35,
|
|
max_angle_deg: float = 30,
|
|
full_output: bool = False,
|
|
) -> Union[NDArray, tuple[NDArray, NDArray, NDArray]]:
|
|
"""
|
|
Compute h-bond pairs
|
|
|
|
Args:
|
|
atoms: Set of all coordinates for a frame.
|
|
donator_indices: Set of indices for donators.
|
|
hydrogen_indices: Set of indices for hydrogen atoms. Should have the same
|
|
length as D.
|
|
acceptor_indices: Set of indices for acceptors.
|
|
DA_lim (opt.): Minimum distance between donator and acceptor.
|
|
HA_lim (opt.): Minimum distance between hydrogen and acceptor.
|
|
max_angle_deg (opt.): Maximum angle in degree for the HDA angle.
|
|
full_output (opt.): Returns additionally the cosine of the
|
|
angles and the DA distances
|
|
|
|
Return:
|
|
List of (D,A)-pairs in hbonds.
|
|
"""
|
|
|
|
def dist_DltA(
|
|
D: CoordinateFrame, A: CoordinateFrame, max_dist: float = 0.35
|
|
) -> NDArray:
|
|
ppoints, pind = pbc_points(D, thickness=max_dist + 0.1, index=True)
|
|
Dtree = spatial.cKDTree(ppoints)
|
|
Atree = spatial.cKDTree(A)
|
|
pairs = Dtree.sparse_distance_matrix(Atree, max_dist, output_type="ndarray")
|
|
pairs = np.asarray(pairs.tolist())
|
|
pairs = np.int_(pairs[pairs[:, 2] > 0][:, :2])
|
|
pairs[:, 0] = pind[pairs[:, 0]]
|
|
return pairs
|
|
|
|
def dist_AltD(
|
|
D: CoordinateFrame, A: CoordinateFrame, max_dist: float = 0.35
|
|
) -> NDArray:
|
|
ppoints, pind = pbc_points(A, thickness=max_dist + 0.1, index=True)
|
|
Atree = spatial.cKDTree(ppoints)
|
|
Dtree = spatial.cKDTree(D)
|
|
pairs = Atree.sparse_distance_matrix(Dtree, max_dist, output_type="ndarray")
|
|
pairs = np.asarray(pairs.tolist())
|
|
pairs = np.int_(pairs[pairs[:, 2] > 0][:, :2])
|
|
pairs = pairs[:, ::-1]
|
|
pairs[:, 1] = pind[pairs[:, 1]]
|
|
return pairs
|
|
|
|
D = atoms[donator_indices]
|
|
H = atoms[hydrogen_indices]
|
|
A = atoms[acceptor_indices]
|
|
min_cos = np.cos(max_angle_deg * np.pi / 180)
|
|
box = D.box
|
|
if len(D) <= len(A):
|
|
pairs = dist_DltA(D, A, DA_lim)
|
|
else:
|
|
pairs = dist_AltD(D, A, DA_lim)
|
|
|
|
vDH = pbc_diff(D[pairs[:, 0]], H[pairs[:, 0]], box)
|
|
vDA = pbc_diff(D[pairs[:, 0]], A[pairs[:, 1]], box)
|
|
vHA = pbc_diff(H[pairs[:, 0]], A[pairs[:, 1]], box)
|
|
angles_cos = np.clip(
|
|
np.einsum("ij,ij->i", vDH, vDA)
|
|
/ np.linalg.norm(vDH, axis=1)
|
|
/ np.linalg.norm(vDA, axis=1),
|
|
-1,
|
|
1,
|
|
)
|
|
is_bond = (
|
|
(angles_cos >= min_cos)
|
|
& (np.sum(vHA**2, axis=-1) <= HA_lim**2)
|
|
& (np.sum(vDA**2, axis=-1) <= DA_lim**2)
|
|
)
|
|
if full_output:
|
|
return (
|
|
pairs[is_bond],
|
|
angles_cos[is_bond],
|
|
np.sum(vDA[is_bond] ** 2, axis=-1) ** 0.5,
|
|
)
|
|
else:
|
|
return pairs[is_bond]
|
|
|
|
|
|
def calc_cluster_sizes(atoms: CoordinateFrame, r_max: float = 0.35) -> NDArray:
|
|
frame_PBC, indices_PBC = pbc_points(atoms, thickness=r_max + 0.1, index=True)
|
|
tree = KDTree(frame_PBC)
|
|
matrix = tree.sparse_distance_matrix(tree, r_max, output_type="ndarray")
|
|
new_matrix = np.zeros((len(atoms), len(atoms)))
|
|
for entry in matrix:
|
|
if entry[2] > 0:
|
|
new_matrix[indices_PBC[entry[0]], indices_PBC[entry[1]]] = 1
|
|
n_components, labels = connected_components(new_matrix, directed=False)
|
|
cluster_sizes = []
|
|
for i in range(0, np.max(labels) + 1):
|
|
cluster_sizes.append(np.sum(labels == i))
|
|
return np.array(cluster_sizes).flatten()
|
|
|
|
|
|
def gyration_radius(position: CoordinateFrame) -> NDArray:
|
|
r"""
|
|
Calculates a list of all radii of gyration of all molecules given in the coordinate
|
|
frame, weighted with the masses of the individual atoms.
|
|
|
|
Args:
|
|
position: Coordinate frame object
|
|
|
|
..math::
|
|
R_G = \left(\frac{\sum_{i=1}^{n} m_i |\vec{r_i}
|
|
- \vec{r_{COM}}|^2 }{\sum_{i=1}^{n} m_i }
|
|
\rigth)^{\frac{1}{2}}
|
|
"""
|
|
gyration_radii = np.array([])
|
|
|
|
for resid in np.unique(position.residue_ids):
|
|
pos = position.whole[position.residue_ids == resid]
|
|
mass = position.masses[position.residue_ids == resid][:, np.newaxis]
|
|
COM = 1 / mass.sum() * (mass * position).sum(axis=0)
|
|
r_sq = ((pbc_diff(pos, COM, pos.box.diagonal())) ** 2).sum(1)[:, np.newaxis]
|
|
g_radius = ((r_sq * mass).sum() / mass.sum()) ** 0.5
|
|
|
|
gyration_radii = np.append(gyration_radii, g_radius)
|
|
|
|
return gyration_radii
|