Concurrence

contact_map/concurrence.py

@@ -2,6 +2,8 @@
 import mdtraj as md
 import numpy as np
 
+import contact_map
+
 try:
     import matplotlib.pyplot as plt
 except ImportError:
@@ -22,7 +24,7 @@ class Concurrence(object):
     ----------
         values : list of list of bool
             the whether a contact is present for each contact pair at each
-            point in time; outer list is length number of frames, inner list
+            point in time; inner list is length number of frames, outer list
             in length number of (included) contacts
         labels : list of string
             labels for each contact pair
@@ -63,6 +65,41 @@ def __getitem__(self, label):
 
         # return sum(coincidence_list) / np.sqrt(norm_sq)
 
+def _regularize_contact_input(contact_input, atom_or_res):
+    """Clean input for concurrence objects.
+
+    The allowed inputs are the :class:`.ContactFrequency`, or the
+    :class:`.ContactObject` coming from the ``.atom_contacts`` or
+    ``.residue_contacts`` attribute of the contact frequency, or the list
+    coming from the ``.most_common()`` method for the
+    :class:`.ContactObject`.
+
+    Parameters
+    ----------
+    contact_input : many possible types; see method description
+        input to the contact concurrences
+    atom_or_res : string
+        whether to treat this as an atom-based or residue-based contact;
+        allowed values are "atom", "res", and "residue"
+
+    Returns
+    -------
+    list :
+        list in the format of ``ContactCount.most_common()``
+    """
+    if isinstance(contact_input, contact_map.ContactFrequency):
+        if atom_or_res in ["atom", "atoms"]:
+            contact_input = contact_input.atom_contacts.most_common()
+        elif atom_or_res in ["residue", "residues", "res"]:
+            contact_input = contact_input.residue_contacts.most_common()
+        else:
+            raise RuntimeError("Bad value for atom_or_res: " +
+                               str(atom_or_res))
+    elif isinstance(contact_input, contact_map.ContactCount):
+        contact_input = contact_input.most_common()
+
+    return contact_input
+
 
 class AtomContactConcurrence(Concurrence):
     """Contact concurrences for atom contacts.
@@ -77,13 +114,15 @@ class AtomContactConcurrence(Concurrence):
         cutoff, in nm. Should be the same as used in the contact map.
     """
     def __init__(self, trajectory, atom_contacts, cutoff=0.45):
-        # TODO: the use of atom_contacts as input from most_common is weird
+        atom_contacts = _regularize_contact_input(atom_contacts, "atom")
         atom_pairs = [[contact[0][0].index, contact[0][1].index]
                       for contact in atom_contacts]
         labels = [str(contact[0]) for contact in atom_contacts]
         distances = md.compute_distances(trajectory, atom_pairs=atom_pairs)
         vector_f = np.vectorize(lambda d: d < cutoff)
-        values = list(map(list, zip(*vector_f(distances))))
+        # transpose because distances is ndarray shape (n_frames,
+        # n_contacts); values should be list shape (n_contacts, n_frames)
+        values = vector_f(distances).T.tolist()
         super(AtomContactConcurrence, self).__init__(values=values,
                                                      labels=labels)
 
@@ -105,7 +144,8 @@ class ResidueContactConcurrence(Concurrence):
     """
     def __init__(self, trajectory, residue_contacts, cutoff=0.45,
                  select="and symbol != 'H'"):
-        # TODO: the use of residue_contacts as input from most_common is weird
+        residue_contacts = _regularize_contact_input(residue_contacts,
+                                                     "residue")
         residue_pairs = [[contact[0][0], contact[0][1]]
                          for contact in residue_contacts]
         labels = [str(contact[0]) for contact in residue_contacts]
@@ -120,7 +160,7 @@ def __init__(self, trajectory, residue_contacts, cutoff=0.45,
             distances = md.compute_distances(trajectory,
                                              atom_pairs=atom_pairs)
             min_dists = [min(dists) for dists in distances]
-            values.append(list(map(lambda d: d < cutoff, min_dists)))
+            values.append([d < cutoff for d in min_dists])
 
         super(ResidueContactConcurrence, self).__init__(values=values,
                                                         labels=labels)
@@ -144,7 +184,7 @@ class ConcurrencePlotter(object):
     def __init__(self, concurrence=None, labels=None, x_values=None):
         self.concurrence = concurrence
         self.labels = self.get_concurrence_labels(concurrence, labels)
-        self.x_values = x_values
+        self._x_values = x_values
 
     @staticmethod
     def get_concurrence_labels(concurrence, labels=None):

contact_map/tests/test_concurrence.py

@@ -20,6 +20,29 @@ def setup_module():
                                 n_neighbors_ignored=0)
 
 
+@pytest.mark.parametrize("contact_type", ('atoms', 'residues'))
+def test_regularize_contact_input(contact_type):
+    from contact_map.concurrence import _regularize_contact_input \
+            as regularize
+    most_common = {
+        'atoms': contacts.atom_contacts.most_common(),
+        'residues': contacts.residue_contacts.most_common()
+    }[contact_type]
+    contact_count = {
+        'atoms': contacts.atom_contacts,
+        'residues': contacts.residue_contacts
+    }[contact_type]
+    assert regularize(most_common, contact_type) == most_common
+    assert regularize(contact_count, contact_type) == most_common
+    assert regularize(contacts, contact_type) == most_common
+
+def test_regularize_contact_input_bad_type():
+    from contact_map.concurrence import _regularize_contact_input \
+            as regularize
+    with pytest.raises(RuntimeError):
+        regularize(contacts, "foo")
+
+
 class ContactConcurrenceTester(object):
     def _test_default_labels(self, concurrence):
         assert len(concurrence.labels) == len(self.labels) / 2