Added type hints
This commit is contained in:
parent
bffcd56cdc
commit
3cfdf79777
@ -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
|
||||
|
Loading…
Reference in New Issue
Block a user