Support atoms in showOccupancies

prody/ensemble/functions.py

@@ -353,23 +353,38 @@ def calcOccupancies(pdb_ensemble, normed=False):
 
 def showOccupancies(pdbensemble, *args, **kwargs):
     """Show occupancies for the PDB ensemble using :func:`~matplotlib.pyplot.
-    plot`.  Occupancies are calculated using :meth:`calcOccupancies`."""
+    plot`.  Occupancies are calculated using :meth:`calcOccupancies`.
+
+    :arg atoms: atoms for showing residue numbers along the x-axis.
+        Default option is to use ``pdbensemble.getAtoms()``.
+    :type atoms: :class:`.Atomic`
+    """
 
     import matplotlib.pyplot as plt
 
     normed = kwargs.pop('normed', False)
+    atoms = kwargs.pop('atoms', None)
 
     if not isinstance(pdbensemble, PDBEnsemble):
         raise TypeError('pdbensemble must be a PDBEnsemble instance')
+    if atoms is None:
+        atoms = pdbensemble.getAtoms()
     weights = calcOccupancies(pdbensemble, normed)
     if weights is None:
         return None
-    show = plt.plot(weights, *args, **kwargs)
+    if atoms is not None:
+        if weights.shape[0] != atoms.numAtoms():
+            raise ValueError('size mismatch between occupancies ({0}) and atoms ({1})'
+                             .format(weights.shape[0], atoms.numAtoms()))
+        show = plt.plot(atoms.getResnums(), weights, *args, **kwargs)
+        plt.xlabel('Residue number')
+    else:
+        show = plt.plot(weights, *args, **kwargs)
+        plt.xlabel('Atom index')
     axis = list(plt.axis())
     axis[2] = 0
     axis[3] += 1
     plt.axis(axis)
-    plt.xlabel('Atom index')
     plt.ylabel('Sum of weights')
     if SETTINGS['auto_show']:
         showFigure()

prody/tests/ensemble/test_functions.py

@@ -1,10 +1,13 @@
 """This module contains unit tests for :mod:`~prody.ensemble`."""
 
-from prody.tests import TestCase
+from prody.tests import TestCase, skipUnless, MATPLOTLIB
+
+if MATPLOTLIB:
+    import matplotlib.pyplot as plt
 
 from numpy.testing import assert_equal
 
-from prody import (calcOccupancies, trimPDBEnsemble, PDBEnsemble, 
+from prody import (calcOccupancies, trimPDBEnsemble, PDBEnsemble, showOccupancies,
                    parsePDB, buildPDBEnsemble, bestMatch, sameChid, 
                    sameChainPos)
 from prody.tests.datafiles import *
@@ -89,4 +92,24 @@ def testResults(self):
         ens6 = buildPDBEnsemble(biomols, match_func=bestMatch)
         assert_equal(ens6.numConfs(), 5, 
             'buildPDBEnsemble with sameChainPos on biomols did not include all AMPAR dimers')        
-
+
+
+@skipUnless(MATPLOTLIB, 'matplotlib not found')
+class TestShowOccupancies(TestCase):
+
+    def tearDown(self):
+        plt.close('all')
+
+    def testUsesExplicitAtoms(self):
+        show = showOccupancies(PDBENSEMBLEA, normed=True, atoms=ATOMS)
+        line = show[0]
+        assert_equal(line.get_xdata(), ATOMS.getResnums(),
+                     'showOccupancies failed to use explicit atoms for x-axis')
+        assert_equal(plt.gca().get_xlabel(), 'Residue number',
+                     'showOccupancies failed to label the residue-number x-axis')
+
+    def testUsesEnsembleAtomsByDefault(self):
+        show = showOccupancies(PDBENSEMBLEA, normed=True)
+        line = show[0]
+        assert_equal(line.get_xdata(), PDBENSEMBLEA.getAtoms().getResnums(),
+                     'showOccupancies failed to use ensemble atoms for x-axis')