Source code for pdb2pqr.cells

"""Cell list to facilitate neighbor searching."""
import logging

_LOGGER = logging.getLogger(__name__)

[docs]class Cells: """Accelerate the search for nearby atoms. A pure all versus all search is O(n^2) - for every atom, every other atom must be searched. This is rather inefficient, especially for large biomolecules where cells may be tens of angstroms apart. The cell class breaks down the xyz biomolecule space into several 3-D cells of desired size - then by simply examining atoms that fall into the adjacent cells one can quickly find nearby cells. """
[docs] def __init__(self, cellsize): """Initialize the cells. :param cellsize: the size of each cell (in Angstroms) :type cellsize: int """ self.cellmap = {} self.cellsize = cellsize
[docs] def assign_cells(self, biomolecule): """Place each atom in a virtual cell for easy neighbor comparison. :param biomolecule: biomolecule with atoms to assign to cells :type biomolecule: Biomolecule """ for atom in biomolecule.atoms: atom.cell = None self.add_cell(atom)
[docs] def add_cell(self, atom): """Add an atom to the cell. :param atom: the atom to add :type atom: Atom """ size = self.cellsize x = atom.x x = (int(x) - 1) // size * size if x < 0 else int(x) // size * size y = atom.y y = (int(y) - 1) // size * size if y < 0 else int(y) // size * size z = atom.z z = (int(z) - 1) // size * size if z < 0 else int(z) // size * size key = (x, y, z) try: self.cellmap[key].append(atom) except KeyError: self.cellmap[key] = [atom] atom.cell = key
[docs] def remove_cell(self, atom): """Remove an atom from a cell. :param atom: the atom to remove :type atom: Atom """ oldcell = atom.cell if oldcell is None: return atom.cell = None self.cellmap[oldcell].remove(atom)
[docs] def get_near_cells(self, atom): """Find all atoms in cells bordering an atom. :param atom: the atom to test :type atom: Atom :return: a list of nearby atoms :rtype: [Atom] """ closeatoms = [] cell = atom.cell if cell is not None: x = cell[0] y = cell[1] z = cell[2] size = self.cellsize for i in range(-1 * size, 2 * size, size): for j in range(-1 * size, 2 * size, size): for k in range(-1 * size, 2 * size, size): newkey = (x + i, y + j, z + k) try: newatoms = self.cellmap[newkey] for atom2 in newatoms: if atom == atom2: continue closeatoms.append(atom2) except KeyError: pass return closeatoms