Fix Averaging and Conformational Analysis for Protein Multimers

CodeEntropy/levels.py

@@ -205,6 +205,7 @@ def get_dihedrals(self, data_container, level):
         # if residue level, looking for dihedrals involving residues
         if level == "residue":
             num_residues = len(data_container.residues)
+            logger.debug(f"Number Residues: {num_residues}")
             if num_residues < 4:
                 logger.debug("no residue level dihedrals")
 
@@ -249,7 +250,7 @@ def get_dihedrals(self, data_container, level):
                     atom_group = atom1 + atom2 + atom3 + atom4
                     dihedrals.append(atom_group.dihedral)
 
-        logger.debug(f"Dihedrals: {dihedrals}")
+        logger.debug(f"Level: {level}, Dihedrals: {dihedrals}")
 
         return dihedrals
 
@@ -309,6 +310,7 @@ def compute_dihedral_conformations(
                 if state
             ]
 
+        logger.debug(f"level: {level}, states: {states}")
         return states
 
     def get_beads(self, data_container, level):
@@ -1138,11 +1140,11 @@ def build_conformational_states(
                                 )
 
                                 if key in states_ua:
-                                    states_ua[key].append(states)
+                                    states_ua[key].extend(states)
                                 else:
                                     states_ua[key] = states
 
-                        elif level == "res":
+                        elif level == "residue":
                             states = self.compute_dihedral_conformations(
                                 mol,
                                 level,
@@ -1157,7 +1159,7 @@ def build_conformational_states(
                             if states_res[group_id] is None:
                                 states_res[group_id] = states
                             else:
-                                states_res[group_id] += states
+                                states_res[group_id].extend(states)
 
                         progress.advance(task)
 

docs/getting_started.rst

@@ -130,6 +130,12 @@ The top_traj_file argument is necessary to identify your simulation data, the ot
      - ``molecules``
      - ``str``
 
+Averaging
+^^^^^^^^^
+The code is able to average over molecules of the same type.
+The grouping arguement is used to control how the averaging is done.
+The default is "molecules" which defines molecules by the number and names of the atoms and groups molecules that are the same.
+You can also use "each" which makes each molecule its own group, effectively not averaging over molecules.
 
 Example #1
 ^^^^^^^^^^

tests/test_CodeEntropy/test_entropy.py

@@ -777,6 +777,46 @@ def test_process_conformational_residue_level(self):
         results = [entry[3] for entry in df]
         self.assertIn(3.33, results)
 
+    def test_process_conformational_entropy_no_states_entry(self):
+        """
+        Tests that `_process_conformational_entropy` logs zero entropy when
+        the group_id is not present in the states dictionary.
+        """
+        # Setup minimal mock universe
+        u = MagicMock()
+
+        # Setup managers and arguments
+        args = MagicMock()
+        run_manager = MagicMock()
+        level_manager = MagicMock()
+        data_logger = DataLogger()
+        group_molecules = MagicMock()
+        manager = EntropyManager(
+            run_manager, args, u, data_logger, level_manager, group_molecules
+        )
+
+        # States dict does NOT contain group_id=1
+        states = {0: np.ones((10, 3))}
+
+        # Mock entropy calculator
+        ce = MagicMock()
+
+        # Run method with group_id=1 (not in states)
+        manager._process_conformational_entropy(
+            group_id=1,
+            mol_container=MagicMock(),
+            ce=ce,
+            level="residue",
+            states=states,
+            number_frames=10,
+        )
+
+        # Assert entropy is zero
+        self.assertEqual(data_logger.molecule_data[0][3], 0)
+
+        # Assert calculator was not called
+        ce.conformational_entropy_calculation.assert_not_called()
+
     def test_compute_entropies_united_atom(self):
         """
         Test that _process_united_atom_entropy is called correctly for 'united_atom'

tests/test_CodeEntropy/test_levels.py

@@ -1253,7 +1253,7 @@ def test_build_conformational_states_united_atom_accumulates_states(self):
             ce,
         )
 
-        assert states_ua[(0, 0)] == ["ua_state_1", ["ua_state_2"]]
+        assert states_ua[(0, 0)] == ["ua_state_1", "ua_state_2"]
 
         # Confirm compute_dihedral_conformations was called twice (once per molecule)
         assert level_manager.compute_dihedral_conformations.call_count == 2
@@ -1291,7 +1291,7 @@ def test_build_conformational_states_residue_level_accumulates_states(self):
 
         # Setup inputs with 2 molecules in same group
         groups = {0: [0, 1]}  # Both mol 0 and mol 1 are in group 0
-        levels = [["res"], ["res"]]
+        levels = [["residue"], ["residue"]]
         start, end, step = 0, 10, 1
         number_frames = 10
         bin_width = 0.1