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`.
+    When atoms are available, :func:`~prody.dynamics.plotting.showAtomicLines`
+    is used to label the x-axis.
+
+    :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
+    from prody.dynamics.plotting import showAtomicLines
 
     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:
+        show = showAtomicLines(weights, *args, atoms=atoms, **kwargs)[0]
+        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,15 @@
 """This module contains unit tests for :mod:`~prody.ensemble`."""
 
-from prody.tests import TestCase
+import numpy as np
+
+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 +94,54 @@ 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 _draw(self):
+        plt.gcf().canvas.draw()
+        return plt.gca().xaxis.get_major_formatter()
+
+    def testUsesExplicitAtoms(self):
+        show = showOccupancies(PDBENSEMBLEA, normed=True, atoms=ATOMS)
+        line = show[0]
+        assert_equal(line.get_xdata(), np.arange(ATOMS.numAtoms(), dtype=float),
+                     'showOccupancies failed to use atomic plotting positions')
+        formatter = self._draw()
+        assert_equal(str(formatter(0, None)), str(ATOMS.getResnums()[0]),
+                     'showOccupancies failed to label the x-axis using residue numbers')
+        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(),
+                     np.arange(PDBENSEMBLEA.getAtoms().numAtoms(), dtype=float),
+                     'showOccupancies failed to use ensemble atoms for atomic plotting positions')
+
+    def testUsesAtomicLabelsForMultipleChains(self):
+        atoms = parsePDB(DATA_FILES['3hsy']['path'], subset='ca')
+        ensemble = PDBEnsemble('multi-chain')
+        ensemble.setAtoms(atoms)
+        ensemble.setCoords(atoms.getCoords())
+        ensemble.addCoordset(atoms)
+
+        show = showOccupancies(ensemble, normed=True)
+        line = show[0]
+        assert_equal(line.get_xdata(), np.arange(atoms.numAtoms(), dtype=float),
+                     'showOccupancies failed to use sequential positions for multiple chains')
+
+        formatter = self._draw()
+        chids = atoms.getChids()
+        first_chain = chids[0]
+        second_index = np.where(chids != first_chain)[0][0]
+        assert_equal(str(formatter(0, None)), '{0}:{1}'.format(chids[0], atoms.getResnums()[0]),
+                     'showOccupancies failed to label the first chain correctly')
+        assert_equal(str(formatter(second_index, None)),
+                     '{0}:{1}'.format(chids[second_index], atoms.getResnums()[second_index]),
+                     'showOccupancies failed to label later chains correctly')