Added type hints
This commit is contained in:
parent
bffcd56cdc
commit
3cfdf79777
@ -14,7 +14,6 @@ from .coordinates import (
|
|||||||
next_neighbors,
|
next_neighbors,
|
||||||
)
|
)
|
||||||
from .autosave import autosave_data
|
from .autosave import autosave_data
|
||||||
from .utils import runningmean
|
|
||||||
from .pbc import pbc_diff, pbc_points
|
from .pbc import pbc_diff, pbc_points
|
||||||
|
|
||||||
|
|
||||||
@ -56,7 +55,7 @@ def time_average(
|
|||||||
def gr(
|
def gr(
|
||||||
atoms_a: CoordinateFrame,
|
atoms_a: CoordinateFrame,
|
||||||
atoms_b: Optional[CoordinateFrame] = None,
|
atoms_b: Optional[CoordinateFrame] = None,
|
||||||
bins: Union[int, ArrayLike] = 100,
|
bins: Optional[ArrayLike] = None,
|
||||||
remove_intra: bool = False,
|
remove_intra: bool = False,
|
||||||
**kwargs
|
**kwargs
|
||||||
) -> np.ndarray:
|
) -> np.ndarray:
|
||||||
@ -84,13 +83,10 @@ def gr(
|
|||||||
distinct = False
|
distinct = False
|
||||||
elif np.array_equal(atoms_a, atoms_b):
|
elif np.array_equal(atoms_a, atoms_b):
|
||||||
distinct = False
|
distinct = False
|
||||||
|
if bins is None:
|
||||||
|
bins = np.arange(0, 1, 0.01)
|
||||||
|
|
||||||
box = atoms_b.box
|
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))
|
n = len(atoms_a) / np.prod(np.diag(box))
|
||||||
V = 4 / 3 * np.pi * bins[-1] ** 3
|
V = 4 / 3 * np.pi * bins[-1] ** 3
|
||||||
particles_in_volume = int(n * V * 1.1)
|
particles_in_volume = int(n * V * 1.1)
|
||||||
@ -98,7 +94,7 @@ def gr(
|
|||||||
atoms_a,
|
atoms_a,
|
||||||
atoms_b,
|
atoms_b,
|
||||||
number_of_neighbors=particles_in_volume,
|
number_of_neighbors=particles_in_volume,
|
||||||
distance_upper_bound=upper_bound,
|
distance_upper_bound=bins[-1],
|
||||||
distinct=distinct,
|
distinct=distinct,
|
||||||
**kwargs
|
**kwargs
|
||||||
)
|
)
|
||||||
@ -114,9 +110,7 @@ def gr(
|
|||||||
else:
|
else:
|
||||||
distances = distances.flatten()
|
distances = distances.flatten()
|
||||||
|
|
||||||
hist, bins = np.histogram(
|
hist, bins = np.histogram(distances, bins=bins, range=(0, bins[-1]), density=False)
|
||||||
distances, bins=bins, range=(0, upper_bound), density=False
|
|
||||||
)
|
|
||||||
hist = hist / len(atoms_a)
|
hist = hist / len(atoms_a)
|
||||||
hist = hist / (4 / 3 * np.pi * bins[1:] ** 3 - 4 / 3 * np.pi * bins[:-1] ** 3)
|
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))
|
n = len(atoms_b) / np.prod(np.diag(atoms_b.box))
|
||||||
@ -126,14 +120,16 @@ def gr(
|
|||||||
|
|
||||||
|
|
||||||
def distance_distribution(
|
def distance_distribution(
|
||||||
atoms: ArrayLike, bins: Optional[int, ArrayLike]
|
atoms: CoordinateFrame, bins: Optional[int, ArrayLike]
|
||||||
) -> np.ndarray:
|
) -> np.ndarray:
|
||||||
connection_vectors = atoms[:-1, :] - atoms[1:, :]
|
connection_vectors = atoms[:-1, :] - atoms[1:, :]
|
||||||
connection_lengths = (connection_vectors**2).sum(axis=1) ** 0.5
|
connection_lengths = (connection_vectors**2).sum(axis=1) ** 0.5
|
||||||
return np.histogram(connection_lengths, bins)[0]
|
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:
|
if reference_atoms is None:
|
||||||
reference_atoms = atoms
|
reference_atoms = atoms
|
||||||
indices = next_neighbors(
|
indices = next_neighbors(
|
||||||
@ -175,15 +171,22 @@ def tetrahedral_order(atoms, reference_atoms=None):
|
|||||||
return q
|
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."
|
assert bins is not None, "Bin edges of the distribution have to be specified."
|
||||||
Q = tetrahedral_order(atoms, reference_atoms=reference_atoms)
|
Q = tetrahedral_order(atoms, reference_atoms=reference_atoms)
|
||||||
return np.histogram(Q, bins=bins)[0]
|
return np.histogram(Q, bins=bins)[0]
|
||||||
|
|
||||||
|
|
||||||
def radial_density(
|
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.
|
Calculate the radial density distribution.
|
||||||
|
|
||||||
@ -192,7 +195,7 @@ def radial_density(
|
|||||||
Args:
|
Args:
|
||||||
atoms:
|
atoms:
|
||||||
Set of coordinates.
|
Set of coordinates.
|
||||||
bins:
|
bins (opt.):
|
||||||
Bin specification that is passed to numpy.histogram. This needs to be
|
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.
|
a list of bin edges if the function is used within time_average.
|
||||||
symmetry_axis (opt.):
|
symmetry_axis (opt.):
|
||||||
@ -200,30 +203,27 @@ def radial_density(
|
|||||||
default is z-axis.
|
default is z-axis.
|
||||||
origin (opt.):
|
origin (opt.):
|
||||||
Origin of the rotational symmetry, e.g. center of the pore.
|
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)
|
cartesian = rotate_axis(atoms - origin, symmetry_axis)
|
||||||
radius, _ = polar_coordinates(cartesian[:, 0], cartesian[:, 1])
|
radius, _ = polar_coordinates(cartesian[:, 0], cartesian[:, 1])
|
||||||
hist = np.histogram(radius, bins=bins)[0]
|
hist = np.histogram(radius, bins=bins)[0]
|
||||||
volume = np.pi * (bins[1:] ** 2 - bins[:-1] ** 2) * height
|
volume = np.pi * (bins[1:] ** 2 - bins[:-1] ** 2) * length
|
||||||
res = hist / volume
|
return hist / volume
|
||||||
if returnx:
|
|
||||||
return np.vstack((runningmean(bins, 2), res))
|
|
||||||
else:
|
|
||||||
return res
|
|
||||||
|
|
||||||
|
|
||||||
def shell_density(
|
def shell_density(
|
||||||
atoms,
|
atoms: CoordinateFrame,
|
||||||
shell_radius,
|
shell_radius: float,
|
||||||
bins,
|
bins: ArrayLike,
|
||||||
shell_thickness=0.5,
|
shell_thickness: float = 0.5,
|
||||||
symmetry_axis=(0, 0, 1),
|
symmetry_axis: ArrayLike = (0, 0, 1),
|
||||||
origin=(0, 0, 0),
|
origin: Optional[ArrayLike] = None,
|
||||||
):
|
) -> np.ndarray:
|
||||||
"""
|
"""
|
||||||
Compute the density distribution on a cylindrical shell.
|
Compute the density distribution on a cylindrical shell.
|
||||||
|
|
||||||
@ -240,6 +240,8 @@ def shell_density(
|
|||||||
Returns:
|
Returns:
|
||||||
Two-dimensional density distribution of the atoms in the defined shell.
|
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)
|
cartesian = rotate_axis(atoms - origin, symmetry_axis)
|
||||||
radius, theta = polar_coordinates(cartesian[:, 0], cartesian[:, 1])
|
radius, theta = polar_coordinates(cartesian[:, 0], cartesian[:, 1])
|
||||||
shell_indices = (shell_radius <= radius) & (
|
shell_indices = (shell_radius <= radius) & (
|
||||||
@ -250,40 +252,13 @@ def shell_density(
|
|||||||
return hist
|
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(
|
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.
|
Compute the distribution of next neighbors with the same residue name.
|
||||||
"""
|
"""
|
||||||
@ -299,15 +274,14 @@ def next_neighbor_distribution(
|
|||||||
|
|
||||||
|
|
||||||
def hbonds(
|
def hbonds(
|
||||||
D,
|
D: CoordinateFrame,
|
||||||
H,
|
H: CoordinateFrame,
|
||||||
A,
|
A: CoordinateFrame,
|
||||||
box,
|
DA_lim: float = 0.35,
|
||||||
DA_lim=0.35,
|
HA_lim: float = 0.35,
|
||||||
HA_lim=0.35,
|
min_cos: float = np.cos(30 * np.pi / 180),
|
||||||
min_cos=np.cos(30 * np.pi / 180),
|
full_output: bool = False,
|
||||||
full_output=False,
|
) -> Union[np.ndarray, tuple[np.ndarray, np.ndarray, np.ndarray]]:
|
||||||
):
|
|
||||||
"""
|
"""
|
||||||
Compute h-bond pairs
|
Compute h-bond pairs
|
||||||
|
|
||||||
@ -327,8 +301,10 @@ def hbonds(
|
|||||||
List of (D,A)-pairs in hbonds.
|
List of (D,A)-pairs in hbonds.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
def dist_DltA(D, H, A, box, max_dist=0.35):
|
def dist_DltA(
|
||||||
ppoints, pind = pbc_points(D, box, thickness=max_dist + 0.1, index=True)
|
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)
|
Dtree = spatial.cKDTree(ppoints)
|
||||||
Atree = spatial.cKDTree(A)
|
Atree = spatial.cKDTree(A)
|
||||||
pairs = Dtree.sparse_distance_matrix(Atree, max_dist, output_type="ndarray")
|
pairs = Dtree.sparse_distance_matrix(Atree, max_dist, output_type="ndarray")
|
||||||
@ -337,8 +313,10 @@ def hbonds(
|
|||||||
pairs[:, 0] = pind[pairs[:, 0]]
|
pairs[:, 0] = pind[pairs[:, 0]]
|
||||||
return pairs
|
return pairs
|
||||||
|
|
||||||
def dist_AltD(D, H, A, box, max_dist=0.35):
|
def dist_AltD(
|
||||||
ppoints, pind = pbc_points(A, box, thickness=max_dist + 0.1, index=True)
|
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)
|
Atree = spatial.cKDTree(ppoints)
|
||||||
Dtree = spatial.cKDTree(D)
|
Dtree = spatial.cKDTree(D)
|
||||||
pairs = Atree.sparse_distance_matrix(Dtree, max_dist, output_type="ndarray")
|
pairs = Atree.sparse_distance_matrix(Dtree, max_dist, output_type="ndarray")
|
||||||
@ -348,10 +326,11 @@ def hbonds(
|
|||||||
pairs[:, 1] = pind[pairs[:, 1]]
|
pairs[:, 1] = pind[pairs[:, 1]]
|
||||||
return pairs
|
return pairs
|
||||||
|
|
||||||
|
box = D.box
|
||||||
if len(D) <= len(A):
|
if len(D) <= len(A):
|
||||||
pairs = dist_DltA(D, H, A, box, DA_lim)
|
pairs = dist_DltA(D, A, DA_lim)
|
||||||
else:
|
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)
|
vDH = pbc_diff(D[pairs[:, 0]], H[pairs[:, 0]], box)
|
||||||
vDA = pbc_diff(D[pairs[:, 0]], A[pairs[:, 1]], box)
|
vDA = pbc_diff(D[pairs[:, 0]], A[pairs[:, 1]], box)
|
||||||
@ -378,13 +357,11 @@ def hbonds(
|
|||||||
return pairs[is_bond]
|
return pairs[is_bond]
|
||||||
|
|
||||||
|
|
||||||
def calc_cluster_sizes(frame, r_max=0.35):
|
def calc_cluster_sizes(atoms: CoordinateFrame, r_max: float = 0.35) -> np.ndarray:
|
||||||
frame_PBC, indices_PBC = pbc_points(
|
frame_PBC, indices_PBC = pbc_points(atoms, thickness=r_max + 0.1, index=True)
|
||||||
frame, frame.box, thickness=r_max + 0.1, index=True
|
|
||||||
)
|
|
||||||
tree = KDTree(frame_PBC)
|
tree = KDTree(frame_PBC)
|
||||||
matrix = tree.sparse_distance_matrix(tree, r_max, output_type="ndarray")
|
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:
|
for entry in matrix:
|
||||||
if entry[2] > 0:
|
if entry[2] > 0:
|
||||||
new_matrix[indices_PBC[entry[0]], indices_PBC[entry[1]]] = 1
|
new_matrix[indices_PBC[entry[0]], indices_PBC[entry[1]]] = 1
|
||||||
|
Loading…
Reference in New Issue
Block a user