Changed input of hbonds to atom_indices

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)