Source code for MDAnalysis.lib.NeighborSearch

# -*- Mode: python; tab-width: 4; indent-tabs-mode:nil; coding:utf-8 -*-
# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4
#
# MDAnalysis --- http://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 GNU Public Licence, v2 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.
#
# 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`.
"""
from __future__ import absolute_import

import numpy as np
from Bio.KDTree import KDTree

from MDAnalysis.core.groups import AtomGroup, Atom

[docs]class AtomNeighborSearch(object): """This class can be used to find all atoms/residues/segements within the radius of a given query position. This class is using the BioPython KDTree for the neighborsearch. This class also does not apply PBC to the distance calculattions. So you have to ensure yourself that the trajectory has been corrected for PBC artifacts. """ def __init__(self, atom_group, bucket_size=10): """ Parameters ---------- atom_list : AtomGroup list of atoms bucket_size : int Number of entries in leafs of the KDTree. If you suffer poor performance you can play around with this number. Increasing the `bucket_size` will speed up the construction of the KDTree but slow down the search. """ self.atom_group = atom_group self._u = atom_group.universe self.kdtree = KDTree(dim=3, bucket_size=bucket_size) self.kdtree.set_coords(atom_group.positions)
[docs] def search(self, atoms, radius, level='A'): """ Return all atoms/residues/segments that are within *radius* of the atoms in *atoms*. Parameters ---------- atoms : AtomGroup, MDAnalysis.core.groups.Atom list of atoms 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*. """ if isinstance(atoms, Atom): positions = atoms.position.reshape(1, 3) else: positions = atoms.positions indices = [] for pos in positions: self.kdtree.search(pos, radius) indices.append(self.kdtree.get_indices()) unique_idx = np.unique([i for l in indices for i in l]).astype(np.int64) return self._index2level(unique_idx, level)
def _index2level(self, indices, level): """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*. """ n_atom_list = self.atom_group[indices] if level == 'A': if not n_atom_list: return [] else: return n_atom_list elif level == 'R': return list({a.residue for a in n_atom_list}) elif level == 'S': return list(set([a.segment for a in n_atom_list])) else: raise NotImplementedError('{0}: level not implemented'.format(level))