1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
|
# -*- Mode: python; tab-width: 4; indent-tabs-mode:nil; coding:utf-8 -*-
# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4
#
# MDAnalysis --- https://www.mdanalysis.org
# Copyright (c) 2006-2017 The MDAnalysis Development Team and contributors
# (see the file AUTHORS for the full list of names)
#
# Released under the Lesser GNU Public Licence, v2.1 or any higher version
#
# Please cite your use of MDAnalysis in published work:
#
# R. J. Gowers, M. Linke, J. Barnoud, T. J. E. Reddy, M. N. Melo, S. L. Seyler,
# D. L. Dotson, J. Domanski, S. Buchoux, I. M. Kenney, and O. Beckstein.
# MDAnalysis: A Python package for the rapid analysis of molecular dynamics
# simulations. In S. Benthall and S. Rostrup editors, Proceedings of the 15th
# Python in Science Conference, pages 102-109, Austin, TX, 2016. SciPy.
# doi: 10.25080/majora-629e541a-00e
#
# N. Michaud-Agrawal, E. J. Denning, T. B. Woolf, and O. Beckstein.
# MDAnalysis: A Toolkit for the Analysis of Molecular Dynamics Simulations.
# J. Comput. Chem. 32 (2011), 2319--2327, doi:10.1002/jcc.21787
#
"""
Neighbor Search wrapper for MDAnalysis --- :mod:`MDAnalysis.lib.NeighborSearch`
===============================================================================
This module contains classes that allow neighbor searches directly with
`AtomGroup` objects from `MDAnalysis`.
"""
import numpy as np
from MDAnalysis.lib.distances import capped_distance
from MDAnalysis.lib.util import unique_int_1d
from MDAnalysis.core.groups import AtomGroup, SegmentGroup, ResidueGroup
import numpy.typing as npt
from typing import Optional, Union, List
class AtomNeighborSearch(object):
"""This class can be used to find all atoms/residues/segments within the
radius of a given query position.
For the neighbor search, this class is a wrapper around
:class:`~MDAnalysis.lib.distances.capped_distance`.
"""
def __init__(
self, atom_group: AtomGroup, box: Optional[npt.ArrayLike] = None
) -> None:
"""
Parameters
----------
atom_list : AtomGroup
list of atoms
box : array-like or ``None``, optional, default ``None``
Simulation cell dimensions in the form of
:attr:`MDAnalysis.trajectory.timestep.Timestep.dimensions` when
periodic boundary conditions should be taken into account for
the calculation of contacts.
"""
self.atom_group = atom_group
self._u = atom_group.universe
self._box = box
def search(
self, atoms: AtomGroup, radius: float, level: str = "A"
) -> Optional[Union[AtomGroup, ResidueGroup, SegmentGroup]]:
"""
Return all atoms/residues/segments that are within *radius* of the
atoms in *atoms*.
Parameters
----------
atoms : AtomGroup, MDAnalysis.core.groups.AtomGroup
AtomGroup object
radius : float
Radius for search in Angstrom.
level : str
char (A, R, S). Return atoms(A), residues(R) or segments(S) within
*radius* of *atoms*.
Returns
-------
AtomGroup : :class:`~MDAnalysis.core.groups.AtomGroup`
When ``level='A'``, AtomGroup is being returned.
ResidueGroup : :class:`~MDAnalysis.core.groups.ResidueGroup`
When ``level='R'``, ResidueGroup is being returned.
SegmentGroup : :class:`~MDAnalysis.core.groups.SegmentGroup`
When ``level='S'``, SegmentGroup is being returned.
.. versionchanged:: 2.0.0
Now returns :class:`AtomGroup` (when empty this is now an empty
:class:`AtomGroup` instead of an empty list), :class:`ResidueGroup`,
or a :class:`SegmentGroup`
"""
unique_idx = []
try:
# For atom groups, take the positions attribute
position = atoms.positions
except AttributeError:
# For atom, take the position attribute
position = atoms.position
pairs = capped_distance(
position,
self.atom_group.positions,
radius,
box=self._box,
return_distances=False,
)
if pairs.size > 0:
unique_idx = unique_int_1d(np.asarray(pairs[:, 1], dtype=np.intp))
return self._index2level(unique_idx, level)
def _index2level(
self, indices: List[int], level: str
) -> Union[AtomGroup, ResidueGroup, SegmentGroup]:
"""Convert list of atom_indices in a AtomGroup to either the
Atoms or segments/residues containing these atoms.
Parameters
----------
indices
list of atom indices
level : str
char (A, R, S). Return atoms(A), residues(R) or segments(S) within
*radius* of *atoms*.
"""
atomgroup = self.atom_group[indices]
if level == "A":
return atomgroup
elif level == "R":
return atomgroup.residues
elif level == "S":
return atomgroup.segments
else:
raise NotImplementedError(
"{0}: level not implemented".format(level)
)
|
