IPKM/mdevaluate
6
1
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Added type hints

Browse Source
This commit is contained in:
Sebastian Kloth 2023-12-28 11:52:58 +01:00
parent bffcd56cdc
commit 3cfdf79777
1 changed files with 64 additions and 87 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

151
src/mdevaluate/distribution.py
View File

@ -14,7 +14,6 @@ from .coordinates import (
next_neighbors,
)
from .autosave import autosave_data
from .utils import runningmean
from .pbc import pbc_diff, pbc_points
@ -56,7 +55,7 @@ def time_average(
def gr(
atoms_a: CoordinateFrame,
atoms_b: Optional[CoordinateFrame] = None,
bins: Union[int, ArrayLike] = 100,
bins: Optional[ArrayLike] = None,
remove_intra: bool = False,
**kwargs
) -> np.ndarray:
@ -84,13 +83,10 @@ def gr(
distinct = False
elif np.array_equal(atoms_a, atoms_b):
distinct = False
if bins is None:
bins = np.arange(0, 1, 0.01)
box = atoms_b.box
if isinstance(bins, int):
upper_bound = np.min(np.diag(box))
else:
upper_bound = bins[-1]
n = len(atoms_a) / np.prod(np.diag(box))
V = 4 / 3 * np.pi * bins[-1] ** 3
particles_in_volume = int(n * V * 1.1)
@ -98,7 +94,7 @@ def gr(
atoms_a,
atoms_b,
number_of_neighbors=particles_in_volume,
distance_upper_bound=upper_bound,
distance_upper_bound=bins[-1],
distinct=distinct,
**kwargs
)
@ -114,9 +110,7 @@ def gr(
else:
distances = distances.flatten()
hist, bins = np.histogram(
distances, bins=bins, range=(0, upper_bound), density=False
)
hist, bins = np.histogram(distances, bins=bins, range=(0, bins[-1]), density=False)
hist = hist / len(atoms_a)
hist = hist / (4 / 3 * np.pi * bins[1:] ** 3 - 4 / 3 * np.pi * bins[:-1] ** 3)
n = len(atoms_b) / np.prod(np.diag(atoms_b.box))
@ -126,14 +120,16 @@ def gr(
def distance_distribution(
atoms: ArrayLike, bins: Optional[int, ArrayLike]
atoms: CoordinateFrame, bins: Optional[int, ArrayLike]
) -> np.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, reference_atoms=None):
def tetrahedral_order(
atoms: CoordinateFrame, reference_atoms: CoordinateFrame = None
) -> np.ndarray:
if reference_atoms is None:
reference_atoms = atoms
indices = next_neighbors(
@ -175,15 +171,22 @@ def tetrahedral_order(atoms, reference_atoms=None):
return q
def tetrahedral_order_distribution(atoms, reference_atoms=None, bins=None):
def tetrahedral_order_distribution(
atoms: CoordinateFrame,
reference_atoms: Optional[CoordinateFrame] = None,
bins: Optional[ArrayLike] = None,
) -> np.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, bins, symmetry_axis=(0, 0, 1), origin=(0, 0, 0), height=1, returnx=False
):
atoms: CoordinateFrame,
bins: Optional[ArrayLike] = None,
symmetry_axis: ArrayLike = (0, 0, 1),
origin: Optional[ArrayLike] = None,
) -> np.ndarray:
"""
Calculate the radial density distribution.
@ -192,7 +195,7 @@ def radial_density(
Args:
atoms:
Set of coordinates.
bins:
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.):
@ -200,30 +203,27 @@ def radial_density(
default is z-axis.
origin (opt.):
Origin of the rotational symmetry, e.g. center of the pore.
height (opt.):
Height of the pore, necessary for correct normalization of the density.
returnx (opt.):
If True, the x ordinate of the distribution is returned.
"""
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) * height
res = hist / volume
if returnx:
return np.vstack((runningmean(bins, 2), res))
else:
return res
volume = np.pi * (bins[1:] ** 2 - bins[:-1] ** 2) * length
return hist / volume
def shell_density(
atoms,
shell_radius,
bins,
shell_thickness=0.5,
symmetry_axis=(0, 0, 1),
origin=(0, 0, 0),
):
atoms: CoordinateFrame,
shell_radius: float,
bins: ArrayLike,
shell_thickness: float = 0.5,
symmetry_axis: ArrayLike = (0, 0, 1),
origin: Optional[ArrayLike] = None,
) -> np.ndarray:
"""
Compute the density distribution on a cylindrical shell.
@ -240,6 +240,8 @@ def shell_density(
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) & (
@ -250,40 +252,13 @@ def shell_density(
return hist
def spatial_density(atoms, bins, weights=None):
"""
Compute the spatial density distribution.
"""
density, _ = np.histogramdd(atoms, bins=bins, weights=weights)
return density
def mixing_ratio_distribution(
atoms_a,
atoms_b,
bins_ratio,
bins_density,
weights_a=None,
weights_b=None,
weights_ratio=None,
):
"""
Compute the distribution of the mixing ratio of two sets of atoms.
"""
density_a, _ = time_average
density_b, _ = np.histogramdd(atoms_b, bins=bins_density, weights=weights_b)
mixing_ratio = density_a / (density_a + density_b)
good_inds = (density_a != 0) & (density_b != 0)
hist, _ = np.histogram(
mixing_ratio[good_inds], bins=bins_ratio, weights=weights_ratio
)
return hist
def next_neighbor_distribution(
atoms, reference=None, number_of_neighbors=4, bins=None, normed=True
):
atoms: CoordinateFrame,
reference: Optional[CoordinateFrame] = None,
number_of_neighbors: int = 4,
bins: Optional[ArrayLike] = None,
normed: bool = True,
) -> np.ndarray:
"""
Compute the distribution of next neighbors with the same residue name.
"""
@ -299,15 +274,14 @@ def next_neighbor_distribution(
def hbonds(
D,
H,
A,
box,
DA_lim=0.35,
HA_lim=0.35,
min_cos=np.cos(30 * np.pi / 180),
full_output=False,
):
D: CoordinateFrame,
H: CoordinateFrame,
A: CoordinateFrame,
DA_lim: float = 0.35,
HA_lim: float = 0.35,
min_cos: float = np.cos(30 * np.pi / 180),
full_output: bool = False,
) -> Union[np.ndarray, tuple[np.ndarray, np.ndarray, np.ndarray]]:
"""
Compute h-bond pairs
@ -327,8 +301,10 @@ def hbonds(
List of (D,A)-pairs in hbonds.
"""
def dist_DltA(D, H, A, box, max_dist=0.35):
ppoints, pind = pbc_points(D, box, thickness=max_dist + 0.1, index=True)
def dist_DltA(
D: CoordinateFrame, A: CoordinateFrame, max_dist: float = 0.35
) -> np.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")
@ -337,8 +313,10 @@ def hbonds(
pairs[:, 0] = pind[pairs[:, 0]]
return pairs
def dist_AltD(D, H, A, box, max_dist=0.35):
ppoints, pind = pbc_points(A, box, thickness=max_dist + 0.1, index=True)
def dist_AltD(
D: CoordinateFrame, A: CoordinateFrame, max_dist: float = 0.35
) -> np.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")
@ -348,10 +326,11 @@ def hbonds(
pairs[:, 1] = pind[pairs[:, 1]]
return pairs
box = D.box
if len(D) <= len(A):
pairs = dist_DltA(D, H, A, box, DA_lim)
pairs = dist_DltA(D, A, DA_lim)
else:
pairs = dist_AltD(D, H, A, box, DA_lim)
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)
@ -378,13 +357,11 @@ def hbonds(
return pairs[is_bond]
def calc_cluster_sizes(frame, r_max=0.35):
frame_PBC, indices_PBC = pbc_points(
frame, frame.box, thickness=r_max + 0.1, index=True
)
def calc_cluster_sizes(atoms: CoordinateFrame, r_max: float = 0.35) -> np.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(frame), len(frame)))
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