diff --git a/src/mdevaluate/atoms.py b/src/mdevaluate/atoms.py index 5f0813e..22582a4 100644 --- a/src/mdevaluate/atoms.py +++ b/src/mdevaluate/atoms.py @@ -10,7 +10,7 @@ from .checksum import checksum from .coordinates import CoordinateFrame -def compare_regex(str_list: list[str], exp: str): +def compare_regex(str_list: list[str], exp: str) -> np.ndarray: """ Compare a list of strings with a regular expression. """ @@ -79,10 +79,10 @@ class AtomSubset: the keyworss below. Names are matched as a regular expression with `re.match`. Args: - atom_name: Specification of the atom name - residue_name: Specification of the resiude name - residue_id: Residue ID or list of IDs - indices: List of atom indices + atom_name: Specification of the atom name + residue_name: Specification of the resiude name + residue_id: Residue ID or list of IDs + indices: List of atom indices """ new_subset = self if atom_name is not None: @@ -176,7 +176,7 @@ class AtomSubset: return checksum(self.description) -def gyration_radius(position: CoordinateFrame): +def gyration_radius(position: CoordinateFrame) -> np.ndarray: r""" Calculates a list of all radii of gyration of all molecules given in the coordinate frame, weighted with the masses of the individual atoms. @@ -185,7 +185,8 @@ def gyration_radius(position: CoordinateFrame): position: Coordinate frame object ..math:: - R_G = \left(\frac{\sum_{i=1}^{n} m_i |\vec{r_i} - \vec{r_{COM}}|^2 }{\sum_{i=1}^{n} m_i } + R_G = \left(\frac{\sum_{i=1}^{n} m_i |\vec{r_i} + - \vec{r_{COM}}|^2 }{\sum_{i=1}^{n} m_i } \rigth)^{\frac{1}{2}} """ gyration_radii = np.array([]) @@ -207,7 +208,7 @@ def layer_of_atoms( thickness: float, plane_normal: npt.ArrayLike, plane_offset: Optional[npt.ArrayLike] = np.array([0, 0, 0]), -): +) -> np.array: if plane_offset is None: np.array([0, 0, 0]) atoms = atoms - plane_offset @@ -215,45 +216,14 @@ def layer_of_atoms( return np.abs(distance) <= thickness -def distance_to_atoms(ref, atoms, box=None): - """Get the minimal distance from atoms to ref. - The result is an array of with length == len(atoms) - """ - out = np.empty(atoms.shape[0]) - for i, atom in enumerate(atoms): - diff = (pbc_diff(atom, ref, box) ** 2).sum(axis=1).min() - out[i] = np.sqrt(diff) - return out - - -def distance_to_atoms_KDtree(ref, atoms, box=None, thickness=None): - """ - Get the minimal distance from atoms to ref. - The result is an array of with length == len(atoms) - Can be faster than distance_to_atoms. - Thickness defaults to box/5. If this is too small results may be wrong. - If box is not given then periodic boundary conditions are not applied! - """ - if thickness == None: - thickness = box / 5 - if box is not None: - start_coords = np.copy(atoms) % box - all_frame_coords = pbc_points(ref, box, thickness=thickness) - else: - start_coords = atoms - all_frame_coords = ref - - tree = KDTree(all_frame_coords) - return tree.query(start_coords)[0] - - def next_neighbors( - atoms, - query_atoms=None, - number_of_neighbors=1, - distance_upper_bound=np.inf, - distinct=False, -): + atoms: CoordinateFrame, + query_atoms: Optional[CoordinateFrame] = None, + number_of_neighbors: int = 1, + distance_upper_bound: float = np.inf, + distinct: bool = False, + **kwargs +) -> (np.ndarray, np.ndarray): """ Find the N next neighbors of a set of atoms. @@ -269,15 +239,29 @@ def next_neighbors( If this is true, the atoms and query atoms are taken as distinct sets of atoms """ - tree = KDTree(atoms) dnn = 0 if query_atoms is None: query_atoms = atoms + dnn = 1 elif not distinct: dnn = 1 - dist, indices = tree.query( - query_atoms, - number_of_neighbors + dnn, - distance_upper_bound=distance_upper_bound, - ) - return indices[:, dnn:] + + box = atoms.box + if np.all(np.diag(np.diag(box)) == box): + atoms = atoms % np.diag(box) + tree = KDTree(atoms, boxsize=np.diag(box)) + distances, indices = tree.query( + query_atoms, number_of_neighbors, distance_upper_bound=distance_upper_bound + ) + else: + atoms_pbc, atoms_pbc_index = pbc_points( + query_atoms, box, thickness=distance_upper_bound+0.1, index=True, **kwargs + ) + tree = KDTree(atoms_pbc) + distances, indices = tree.query( + query_atoms, number_of_neighbors, distance_upper_bound=distance_upper_bound + ) + indices = atoms_pbc_index[indices] + + return distances[:, dnn:], indices[:, dnn:] +