extended AtomGroup with: bond orders, the ability to extend an atomgr…

prody/atomic/atomgroup.py

@@ -127,7 +127,8 @@ class AtomGroup(Atomic):
                  '_timestamps', '_kdtrees',
                  '_bmap', '_angmap', '_dmap', '_imap',
                  '_domap', '_acmap', '_nbemap', '_cmap',
-                 '_bonds', '_angles', '_dihedrals', '_impropers',
+                 '_bonds', '_bondOrders', '_bondIndex', '_angles',
+                 '_dihedrals', '_impropers',
                  '_donors', '_acceptors', '_nbexclusions', '_crossterms',
                  '_cslabels', '_acsi', '_n_csets', '_data',
                  '_fragments', '_flags', '_flagsts', '_subsets',
@@ -144,6 +145,8 @@ def __init__(self, title='Unnamed'):
         self._kdtrees = None
         self._bmap = None
         self._bonds = None
+        self._bondOrders = None
+        self._bondIndex = None
         self._angmap = None
         self._angles = None
         self._dmap = None
@@ -229,17 +232,35 @@ def __len__(self):
 
         return self._n_atoms
 
-    def __add__(self, other):
-
+    def __add__(self, other, newAG=True):
+        """This method adds two atom groups. By default this creates a new atom group.
+
+        e.g.:
+           newAG = ag1 + ag2
+           len(newAG) == (len(ag1) + len(ag2))
+
+        if newAG is set to false, ag1 is extended with the atoms of ag2
+           oldLength = len(ag1)
+           oldID = id(ag1)
+           ag1.__add__(ag2, newAG=False)
+           len(ag1) == (oldLength + len(ag2))
+           id(ag1) === oldID
+        """
         if not isinstance(other, AtomGroup):
             raise TypeError('unsupported operand type(s) for +: {0} and '
                             '{1}'.format(repr(type(self).__name__),
                                          repr(type(other).__name__)))
 
-        new = AtomGroup(self._title + ' + ' + other._title)
+        oldSize = len(self)
+        if newAG:
+            new = AtomGroup(self._title + ' + ' + other._title)
+        else:
+            new = self
+        self._n_atoms += other._n_atoms
+
         if self._n_csets:
             if self._n_csets == other._n_csets:
-                new.setCoords(np.concatenate((self._coords, other._coords), 1))
+                new.setCoords(np.concatenate((self._coords, other._coords), 1), overwrite=True)
                 this = self._anisous
                 that = other._anisous
                 if this is not None and that is not None:
@@ -268,7 +289,7 @@ def __add__(self, other):
             if this is not None or that is not None:
                 if this is None:
                     shape = list(that.shape)
-                    shape[0] = len(self)
+                    shape[0] = oldSize
                     this = np.zeros(shape, that.dtype)
                 if that is None:
                     shape = list(this.shape)
@@ -285,7 +306,17 @@ def __add__(self, other):
             for flag in other._flags:
                 if flag not in keys: keys.append(flag)
 
-        for key in keys:
+        # remove aliases
+        skip = []
+        uniqueKeys = []
+        for k in keys:
+            aliases = FLAG_ALIASES.get(k, [k])
+            if aliases[0] not in uniqueKeys and aliases[0] not in skip:
+                uniqueKeys.append(aliases[0])
+                if len(aliases) > 1:
+                    skip.extend(list(aliases[1:]))
+
+        for key in uniqueKeys:
             this = None
             that = None
             if self._flags:
@@ -295,23 +326,34 @@ def __add__(self, other):
             if this is not None or that is not None:
                 if this is None:
                     shape = list(that.shape)
-                    shape[0] = len(self)
+                    shape[0] = oldSize
                     this = np.zeros(shape, that.dtype)
                 if that is None:
                     shape = list(this.shape)
                     shape[0] = len(other)
                     that = np.zeros(shape, this.dtype)
                 new._setFlags(key, np.concatenate((this, that)))
-
-        new._n_atoms = self._n_atoms + other._n_atoms
+
+        if self._bondOrders is not None:
+            if other._bondOrders is not None:
+                bo = np.concatenate([self._bondsOrders, other._bondOrders])
+            else:
+                bo = np.concatenate([self._bondsOrders, np.ones(len(other), np.int8)])
+        else:
+            if other._bondOrders is not None:
+                bo = np.concatenate([np.ones(len(self), np.int8), self._bondsOrders])
+            else:
+                bo = None
 
         if self._bonds is not None and other._bonds is not None:
             new.setBonds(np.concatenate([self._bonds,
-                                         other._bonds + self._n_atoms]))
+                                         other._bonds + oldSize]), bo)
         elif self._bonds is not None:
-            new.setBonds(self._bonds.copy())
+            new.setBonds(self._bonds.copy(), bo)
+
         elif other._bonds is not None:
-            new.setBonds(other._bonds + self._n_atoms)
+            bonds = other._bonds + self._n_atoms
+            new.setBonds(bonds, bo)
 
         if self._angles is not None and other._angles is not None:
             new.setAngles(np.concatenate([self._angles,
@@ -395,6 +437,10 @@ def __eq__(self, other):
                 all(np.all(self._data[key] == other._data[key])
                     for key in self._data))
 
+    def __hash__(self):
+        return object.__hash__(self)
+
+
     def __iter__(self):
         """Yield atom instances."""
 
@@ -493,7 +539,7 @@ def _getCoords(self):
         if self._coords is not None:
             return self._coords[self._acsi]
 
-    def setCoords(self, coords, label=''):
+    def setCoords(self, coords, label='', overwrite=False):
         """Set coordinates of atoms.  *coords* may be any array like object
         or an object instance with :meth:`getCoords` method.  If the shape of
         coordinate array is ``(n_csets > 1, n_atoms, 3)``, it will replace all
@@ -502,7 +548,9 @@ def setCoords(self, coords, label=''):
         If shape of *coords* is ``(n_atoms, 3)`` or ``(1, n_atoms, 3)``, it
         will replace the active coordinate set.  *label* argument may be used
         to label coordinate set(s).  *label* may be a string or a list of
-        strings length equal to the number of coordinate sets."""
+        strings length equal to the number of coordinate sets. The optional
+        argument *overwrite* can be set to True to force resizing the
+        coordinates array when the number of atoms in the AtomGroup changed."""
 
         atoms = coords
         try:
@@ -524,15 +572,17 @@ def setCoords(self, coords, label=''):
             raise TypeError('coords must be a numpy array or an '
                             'object with `getCoords` method')
 
-        self._setCoords(coords, label=label)
+        self._setCoords(coords, label=label, overwrite=overwrite)
 
     def _setCoords(self, coords, label='', overwrite=False):
         """Set coordinates without data type checking.  *coords* must
         be a :class:`~numpy.ndarray`, but may have data type other than
         :class:`~numpy.float64`, e.g. :class:`~numpy.float32`.  *label*
         argument may be used to label coordinate sets.  *label* may be
         a string or a list of strings length equal to the number of
-        coordinate sets."""
+        coordinate sets. The optional argument *overwrite* can be set
+        to True to force resizing the coordinates array when the number
+        of atoms in the AtomGroup changed."""
 
         n_atoms = self._n_atoms
         if n_atoms:
@@ -1204,15 +1254,35 @@ def setCSLabels(self, labels):
         else:
             raise TypeError('labels must be a list')
 
-    def setBonds(self, bonds):
+    def setBondOrders(self, bondOrders):
+        """Set covalent bond order. *bondOrders* must be a list or an array
+        of integers and provide a value for each bond. Possible values are
+        1:single, 2:double, 3:triple, 4:aromatic, 5:amide. The bond order of
+        all bonds must be set at once. This method must be called after
+        the setBonds() has been called. The bond order is stored in the
+        *_bondOrders* array."""
+        if len(bondOrders)!=len(self._bonds):
+            raise ValueError('invalid bond order list, bond and bond order length mismatch')
+        if min(bondOrders)<1 or max(bondOrders)>5:
+            raise ValueError('invalid bond order value, values must range from 1 to 5')
+
+        if bondOrders is None:
+            self._bondOrders = bondOrders
+        else:
+            self._bondOrders = np.array(bondOrders, np.int8)
+
+    def setBonds(self, bonds, bondOrders=None):
         """Set covalent bonds between atoms.  *bonds* must be a list or an
-        array of pairs of indices.  All bonds must be set at once.  Bonding
-        information can be used to make atom selections, e.g. ``"bonded to
-        index 1"``.  See :mod:`.select` module documentation for details.
-        Also, a data array with number of bonds will be generated and stored
-        with label *numbonds*.  This can be used in atom selections, e.g.
-        ``'numbonds 0'`` can be used to select ions in a system. If *bonds* 
-        is empty or **None**, then all bonds will be removed for this 
+        array of pairs of indices. Optionally *bondOrders* can be specified
+        (see :meth:`setBondOrder` method).  All bonds must be
+        set at once.  Bonding information can be used to make atom selections,
+        e.g. ``"bonded to index 1"``.  See :mod:`.select` module documentation
+        for details. Also, a data array with number of bonds will be generated
+        and stored with label *numbonds*.  This can be used in atom selections,
+        e.g. ``'numbonds 0'`` can be used to select ions in a system. The keys
+        in the *_bondIndex* dictionary is a string representation of the bond's
+        atom indices i.e. '%d %d'%(i, j) with i<j. If  *bonds*  is empty or
+        **None**, then all bonds will be removed for this
         :class:`.AtomGroup`. """
 
         if bonds is None or len(bonds) == 0:
@@ -1229,18 +1299,28 @@ def setBonds(self, bonds):
             raise ValueError('bonds.shape must be (n_bonds, 2)')
         if bonds.min() < 0:
             raise ValueError('negative atom indices are not valid')
+
         n_atoms = self._n_atoms
         if bonds.max() >= n_atoms:
             raise ValueError('atom indices are out of range')
+
         bonds.sort(1)
         bonds = bonds[bonds[:, 1].argsort(), ]
         bonds = bonds[bonds[:, 0].argsort(), ]
         bonds = np.unique(bonds, axis=0)
 
+        d = {}
+        for n, b in enumerate(bonds):
+            key = '%d %d'%(b[0], b[1])
+            d[key] = n
+        self._bondIndex = d
+
         self._bmap, self._data['numbonds'] = evalBonds(bonds, n_atoms)
         self._bonds = bonds
         self._fragments = None
 
+        self.setBondOrders(bondOrders)
+
     def numBonds(self):
         """Returns number of bonds.  Use :meth:`setBonds` or 
         :meth:`inferBonds` for setting bonds."""

prody/atomic/bond.py

@@ -15,23 +15,32 @@ class Bond(object):
     * :func:`len` returns bond length, i.e. :meth:`getLength`
     * :func:`iter` yields :class:`~.Atom` instances"""
 
-    __slots__ = ['_ag', '_acsi', '_indices']
+    __slots__ = ['_ag', '_acsi', '_indices', '_bondOrder']
+    _bondType = {1:'single', 2:'double', 3:'triple', 4:'aromatic', 5:'amide', }
 
     def __init__(self, ag, indices, acsi=None):
 
         self._ag = ag
-        self._indices = np.array(indices)
+        i, j  = self._indices = np.array(indices)
+        if self._ag._bondOrders is not None:
+            if i>j:
+                a=i
+                i=j
+                j=a
+            self._bondOrder = self._ag._bondOrders[self._ag._bondIndex['%d %d'%(i,j)]]
+        else:
+            self._bondOrder = 1 # single bond by default
         if acsi is None:
             self._acsi = ag.getACSIndex()
         else:
             self._acsi = acsi
 
     def __repr__(self):
-
         one, two = self._indices
         names = self._ag._getNames()
-        return '<Bond: {0}({1})--{2}({3}) from {4}>'.format(
-                            names[one], one, names[two], two, str(self._ag))
+        return '<Bond: {0}({1})--{2}({3}) from {4}, bond order: {5}>'.format(
+            names[one], one, names[two], two, str(self._ag),
+            self._bondType[self._bondOrder])
 
     def __str__(self):