Merge pull request #110 from CAHanke/main

Add support for multi-state schemes from IHM dictionary

docs/flr.rst

@@ -106,5 +106,11 @@ The :mod:`ihm.flr` Python module
 .. autoclass:: FPSMPPAtomPosition
    :members:
 
+.. autoclass:: KineticRateFretAnalysisConnection
+   :members:
+
+.. autoclass:: RelaxationTimeFretAnalysisConnection
+   :members:
+
 .. autoclass:: FLRData
    :members:

docs/index.rst

@@ -48,7 +48,8 @@ API Reference:
    reader
    dictionary
    flr
-
+   multi_state_scheme
+
 
 Indices and tables
 ==================

docs/multi_state_scheme.rst

@@ -0,0 +1,29 @@
+.. highlight:: rest
+
+.. _multi_state_scheme_module:
+
+The :mod:`ihm.multi_state_scheme` Python module
+===============================================
+
+.. automodule:: ihm.multi_state_scheme
+
+.. autoclass:: MultiStateScheme
+   :members:
+
+.. autoclass:: Connectivity
+   :members:
+
+.. autoclass:: EquilibriumConstant
+   :members:
+
+.. autoclass:: PopulationEquilibriumConstant
+   :members:
+
+.. autoclass:: KineticRateEquilibriumConstant
+   :members:
+
+.. autoclass:: KineticRate
+   :members:
+
+.. autoclass:: RelaxationTime
+   :members:

ihm/__init__.py

@@ -222,6 +222,10 @@ def __init__(self, title=None, id='model', model_details=None):
         #: See :class:`~ihm.flr.FLRData`.
         self.flr_data = []
 
+        #: All multi-state schemes
+        #: See :class:`~ihm.multi_state_scheme.MultiStateScheme`.
+        self.multi_state_schemes = []
+
     def _make_complete_assembly(self):
         """Fill in the complete assembly with all asym units"""
         # Clear out any existing components
@@ -311,10 +315,22 @@ def _all_model_groups(self, only_in_states=True):
            by a State object; otherwise, also include ModelGroups referenced
            by an OrderedProcess or Ensemble."""
         # todo: raise an error if a modelgroup is present in multiple states
+        seen_model_groups = []
         for state_group in self.state_groups:
             for state in state_group:
                 for model_group in state:
+                    seen_model_groups.append(model_group)
+                    yield model_group
+        for mssc in self._all_multi_state_scheme_connectivities():
+            for model_group in mssc.begin_state:
+                if model_group not in seen_model_groups:
+                    seen_model_groups.append(model_group)
                     yield model_group
+            if mssc.end_state:
+                for model_group in mssc.end_state:
+                    if model_group not in seen_model_groups:
+                        seen_model_groups.append(model_group)
+                        yield model_group
         if not only_in_states:
             for ensemble in self.ensembles:
                 if ensemble.model_group:
@@ -397,7 +413,14 @@ def _all_dataset_groups(self):
             (step.dataset_group for step in self._all_protocol_steps()
              if step.dataset_group),
             (step.dataset_group for step in self._all_analysis_steps()
-             if step.dataset_group))
+             if step.dataset_group),
+            (rt.dataset_group for rt in self._all_relaxation_times()
+             if rt.dataset_group),
+            (kr.dataset_group for kr in self._all_kinetic_rates()
+             if kr.dataset_group),
+            (mssc.dataset_group for mssc in
+             self._all_multi_state_scheme_connectivities()
+             if mssc.dataset_group))
 
     def _all_templates(self):
         """Iterate over all Templates in the system."""
@@ -471,7 +494,11 @@ def _all_ensemble_locations():
             (step.script_file for step in self._all_protocol_steps()
                 if step.script_file),
             (step.script_file for step in self._all_analysis_steps()
-                if step.script_file))
+                if step.script_file),
+            (rt.external_file for rt in self._all_relaxation_times()
+                if rt.external_file),
+            (kr.external_file for kr in self._all_kinetic_rates()
+                if kr.external_file))
 
     def _all_geometric_objects(self):
         """Iterate over all GeometricObjects in the system.
@@ -563,6 +590,57 @@ def _all_entity_ranges(self):
                 for comp in f._all_entities_or_asyms()),
             (d.asym_unit for d in self._all_densities())))
 
+    def _all_multi_state_schemes(self):
+        for mss in self.multi_state_schemes:
+            yield mss
+
+    def _all_multi_state_scheme_connectivities(self):
+        """Iterate over all multi-state scheme connectivities"""
+        for mss in self.multi_state_schemes:
+            for mssc in mss.get_connectivities():
+                yield mssc
+
+    def _all_kinetic_rates(self):
+        """Iterate over all kinetic rates within multi-state schemes"""
+        return _remove_identical(itertools.chain(
+            (mssc.kinetic_rate for mssc in
+             self._all_multi_state_scheme_connectivities()
+             if mssc.kinetic_rate),
+            (c.kinetic_rate for f in
+             self.flr_data for c in f.kinetic_rate_fret_analysis_connections
+             if self.flr_data)))
+
+    def _all_relaxation_times(self):
+        """Iterate over all relaxation times.
+        This includes relaxation times from
+        :class:`ihm.multi_state_scheme.MultiStateScheme`
+        and those assigned to connectivities in
+        :class:`ihm.multi_state_scheme.Connectivity`"""
+        seen_relaxation_times = []
+        for mss in self._all_multi_state_schemes():
+            for rt in mss.get_relaxation_times():
+                if rt in seen_relaxation_times:
+                    continue
+                seen_relaxation_times.append(rt)
+                yield rt
+        for mssc in self._all_multi_state_scheme_connectivities():
+            if mssc.relaxation_time:
+                rt = mssc.relaxation_time
+                if rt in seen_relaxation_times:
+                    continue
+                seen_relaxation_times.append(rt)
+                yield rt
+        # Get the relaxation times from the
+        # flr.RelaxationTimeFRETAnalysisConnection objects
+        if self.flr_data:
+            for f in self.flr_data:
+                for c in f.relaxation_time_fret_analysis_connections:
+                    rt = c.relaxation_time
+                    if rt in seen_relaxation_times:
+                        continue
+                    seen_relaxation_times.append(rt)
+                    yield rt
+
     def _before_write(self):
         """Do any setup necessary before writing out to a file"""
         # Here, we initialize all RestraintGroups by removing any assigned ID