Add utility function to get representative model

Add a new method to System that returns a single
"representative" model. If the user hasn't provided
one explicitly, return the model containing the most
chains instead. Closes #174.

Author: benmwebb

ihm/__init__.py

@@ -268,6 +268,27 @@ def _all_models(self):
                 seen_models[model] = None
                 yield group, model
 
+    def get_representative_model(self):
+        """Get a single :class:`~model.ModelRepresentative` that best
+           represents the entire System. If no representatives have been
+           provided, the :class:`~model.Model` containing the most chains
+           is returned instead."""
+        repmodel = None
+        repmodel_chains = -1
+        for mg in self._all_model_groups():
+            if mg.representatives:
+                return mg.representatives[0]
+            for m in mg:
+                asyms = frozenset(id(rep.asym_unit)
+                                  for rep in m.representation)
+                if len(asyms) > repmodel_chains:
+                    repmodel_chains = len(asyms)
+                    repmodel = m
+        if repmodel is None:
+            raise ValueError("No models")
+        else:
+            return repmodel
+
     def update_locations_in_repositories(self, repos):
         """Update all :class:`~ihm.location.Location` objects in the system
            that lie within a checked-out :class:`~ihm.location.Repository`

test/test_main.py

@@ -1410,6 +1410,39 @@ def test_data_usage(self):
         d = ihm.Disclaimer("foo", name="fooname", url="foourl")
         self.assertEqual(d.type, "disclaimer")
 
+    def test_get_representative_model(self):
+        """Test get_representative_model()"""
+        s = ihm.System()
+        # No models should raise an error
+        self.assertRaises(ValueError, s.get_representative_model)
+
+        e = ihm.Entity('AHC')
+        a1 = ihm.AsymUnit(e, 'foo')
+        a2 = ihm.AsymUnit(e, 'bar')
+        s.entities.append(e)
+        s.asym_units.extend((a1, a2))
+        s1 = ihm.representation.AtomicSegment(
+            a1, starting_model=None, rigid=True)
+        s2 = ihm.representation.AtomicSegment(
+            a2, starting_model=None, rigid=True)
+        m1 = ihm.model.Model(assembly=None, protocol=None,
+                             representation=[s1])
+        m2 = ihm.model.Model(assembly=None, protocol=None,
+                             representation=[s1, s2])
+        mg = ihm.model.ModelGroup((m1, m2))
+        state = ihm.model.State([mg])
+        s.state_groups.append(ihm.model.StateGroup([state]))
+
+        # m2 contains two chains to m1's one, so it should be selected
+        # by default
+        self.assertIs(s.get_representative_model(), m2)
+
+        # If we explicitly denote m1 as the representative, it should
+        # be returned instead
+        mr = ihm.model.ModelRepresentative(m1, "medoid")
+        mg.representatives.append(mr)
+        self.assertIs(s.get_representative_model(), mr)
+
 
 if __name__ == '__main__':
     unittest.main()