Refactor water residue handling and add test for dynamic residue group:

- Replace manual filtering of water residues with MDAnalysis selection language (`select_atoms("resname WAT")`) for more robust residue handling
- Dynamically determine `residue_group` from Sorient_dict keys, including only residues present in the water selection
- Update group label logging to use the new selection method
- Add unit test `test_calculate_water_entropy_adds_resname` to ensure the `residue_names.add(resname)` path is exercised

Author: harryswift01

CodeEntropy/entropy.py

@@ -156,18 +156,6 @@ def _handle_water_entropy(self, start, end, step, water_groups):
             return
 
         for group_id, atom_indices in water_groups.items():
-            water_selection = self._universe.atoms[atom_indices]
-
-            water_residues = [
-                res for res in water_selection.residues if res.resname == "WAT"
-            ]
-
-            residue_group = "_".join(sorted({res.resname for res in water_residues}))
-            residue_count = len(water_residues)
-            atom_count = len(water_selection.atoms)
-            self._data_logger.add_group_label(
-                group_id, residue_group, residue_count, atom_count
-            )
 
             self._calculate_water_entropy(
                 universe=self._universe,
@@ -657,17 +645,31 @@ def _calculate_water_entropy(self, universe, start, end, step, group_id=None):
                     "Transvibrational": 0.0,
                     "Rovibrational": 0.0,
                 }
-
             results[mol_id][entropy_type] += value
 
         for mol_id, components in results.items():
-            total = 0.0
             for entropy_type in ["Orientational", "Transvibrational", "Rovibrational"]:
                 S_component = components[entropy_type]
                 self._data_logger.add_results_data(
                     group_id, "water", entropy_type, S_component
                 )
-                total += S_component
+
+        water_selection = universe.select_atoms("resname WAT")
+        actual_water_residues = len(water_selection.residues)
+
+        residue_names = set()
+        for res_dict in Sorient_dict.values():
+            for resname in res_dict.keys():
+                if resname.upper() in water_selection.residues.resnames:
+                    residue_names.add(resname)
+
+        residue_group = "_".join(sorted(residue_names)) if residue_names else "WAT"
+        residue_count = actual_water_residues
+        atom_count = len(water_selection.atoms)
+
+        self._data_logger.add_group_label(
+            group_id, residue_group, residue_count, atom_count
+        )
 
     def _calculate_water_orientational_entropy(
         self, Sorient_dict, group_id, water_count

tests/test_CodeEntropy/test_entropy.py

@@ -1480,6 +1480,54 @@ def test_calculate_water_entropy_minimal(self, mock_get_entropy):
             ]
         )
 
+    @patch(
+        "waterEntropy.recipes.interfacial_solvent.get_interfacial_water_orient_entropy"
+    )
+    def test_calculate_water_entropy_adds_resname(self, mock_get_entropy):
+        """
+        Test _calculate_water_entropy with Sorient_dict containing a water residue
+        so that residue_names.add(resname) is executed.
+        """
+        # Mock vibrations object
+        mock_vibrations = MagicMock()
+        mock_vibrations.translational_S = {("res1", "WAT"): 2.0}
+        mock_vibrations.rotational_S = {("res1", "WAT"): 3.0}
+
+        # Sorient_dict contains a water residue key "WAT"
+        mock_get_entropy.return_value = (
+            {1: {"WAT": [1.0, 5]}},  # orientational
+            None,
+            mock_vibrations,
+            None,
+            1,  # water_count
+        )
+
+        # Mock universe.select_atoms to return a selection containing "WAT"
+        mock_water_selection = MagicMock()
+        mock_residues_group = MagicMock()
+        mock_residues_group.resnames = ["WAT"]  # this is key
+        mock_water_selection.residues = mock_residues_group
+        mock_water_selection.atoms = [1, 2, 3]  # mock atom count
+        mock_universe = MagicMock()
+        mock_universe.select_atoms.return_value = mock_water_selection
+        mock_universe.trajectory = [1, 2]  # 2 frames
+
+        group_id = 0
+
+        # Call the function
+        self.entropy_manager._data_logger = MagicMock()  # mock logger
+        self.entropy_manager._calculate_water_entropy(
+            mock_universe, start=0, end=1, step=1, group_id=group_id
+        )
+
+        # Check that residue_group is "WAT" and residue_names.add was triggered
+        self.entropy_manager._data_logger.add_group_label.assert_called_with(
+            group_id,
+            "WAT",
+            len(mock_water_selection.residues),
+            len(mock_water_selection.atoms),
+        )
+
     # TODO test for error handling on invalid inputs