diff --git a/src/mdevaluate/distribution.py b/src/mdevaluate/distribution.py index bebe26d..cb20d21 100644 --- a/src/mdevaluate/distribution.py +++ b/src/mdevaluate/distribution.py @@ -12,6 +12,7 @@ from .coordinates import ( Coordinates, CoordinateFrame, next_neighbors, + number_of_neighbors, ) from .autosave import autosave_data from .pbc import pbc_diff, pbc_points @@ -54,11 +55,11 @@ def time_average( @autosave_data(nargs=2, kwargs_keys=("coordinates_b",)) def time_distribution( - function: Callable, - coordinates: Coordinates, - coordinates_b: Optional[Coordinates] = None, - skip: float = 0.1, - segments: int = 100, + 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. @@ -88,7 +89,7 @@ def time_distribution( return times, result -def gr( +def rdf( atoms_a: CoordinateFrame, atoms_b: Optional[CoordinateFrame] = None, bins: Optional[ArrayLike] = None, @@ -122,10 +123,9 @@ def gr( if bins is None: bins = np.arange(0, 1, 0.01) - box = atoms_b.box - 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) + 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, @@ -144,7 +144,7 @@ def gr( ) distances = np.concatenate(new_distances) else: - distances = distances.flatten() + 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_a) @@ -310,26 +310,27 @@ def next_neighbor_distribution( def hbonds( - D: CoordinateFrame, - H: CoordinateFrame, - A: CoordinateFrame, + atoms: CoordinateFrame, + donator_indices: ArrayLike, + hydrogen_indices: ArrayLike, + acceptor_indices: ArrayLike, DA_lim: float = 0.35, HA_lim: float = 0.35, - min_cos: float = np.cos(30 * np.pi / 180), + max_angle_deg: float = 30, full_output: bool = False, ) -> Union[NDArray, tuple[NDArray, NDArray, NDArray]]: """ Compute h-bond pairs Args: - D: Set of coordinates for donators. - H: Set of coordinates for hydrogen atoms. Should have the same + 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. - A: Set of coordinates for acceptors. - DA_lim (opt.): Minimum distance beteen donator and acceptor. - HA_lim (opt.): Minimum distance beteen hydrogen and acceptor. - min_cos (opt.): Minimum cosine for the HDA angle. Default is - equivalent to a maximum angle of 30 degree. + 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 @@ -362,6 +363,10 @@ def hbonds( 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)