Refinements of the integration of waterEntropy into CodeEntropy:

- More robust way of calculating the total entropy of each molecule if it is water
- Refactored how `_finalize_molecule_results` function operates, rather than a call per molecule it is run after all molecules have been processed
- Additional test cases to fully capture the refactored `_calculate_water_entropy` function

Author: harryswift01

CodeEntropy/entropy.py

@@ -125,7 +125,7 @@ def execute(self):
                         number_frames,
                     )
 
-            self._finalize_molecule_results(molecule_id, level)
+        self._finalize_molecule_results()
 
         self._data_logger.log_tables()
 
@@ -312,21 +312,24 @@ def _process_conformational_residue_level(
         )
         self._log_result(mol_id, level, "Conformational", S_conf)
 
-    def _finalize_molecule_results(self, mol_id, level):
+    def _finalize_molecule_results(self):
         """
         Summarizes entropy for a molecule and saves results to file.
 
         Args:
             mol_id (int): ID of the molecule.
-            level (str): Current level name (used for tagging final results).
-        """
-        S_total = self._results_df[self._results_df["Molecule ID"] == mol_id][
-            "Result"
-        ].sum()
-        self._log_result(mol_id, "Molecule Total", "Molecule Total Entropy", S_total)
-        self._data_logger.save_dataframes_as_json(
-            self._results_df, self._residue_results_df, self._args.output_file
-        )
+        """
+        logger.info(f"len(self._results_df) {len(self._results_df)}")
+        for mol_id in range(len(self._results_df)):
+            S_total = self._results_df[self._results_df["Molecule ID"] == mol_id][
+                "Result"
+            ].sum()
+            self._log_result(
+                mol_id, "Molecule Total", "Molecule Total Entropy", S_total
+            )
+            self._data_logger.save_dataframes_as_json(
+                self._results_df, self._residue_results_df, self._args.output_file
+            )
 
     def _log_result(self, mol_id, level, entropy_type, value):
         """
@@ -392,13 +395,36 @@ def _calculate_water_entropy(self, universe, start, end, step):
         self._calculate_water_vibrational_translational_entropy(vibrations)
         self._calculate_water_vibrational_rotational_entropy(vibrations)
 
-        # Compute and log per-molecule totals
-        unique_mol_ids = set(self._residue_results_df["Molecule ID"])
-        for mol_id in unique_mol_ids:
-            S_total = self._residue_results_df[
-                self._residue_results_df["Molecule ID"] == mol_id
-            ]["Result"].sum()
-            self._log_result(mol_id, "water", "Molecule Total Entropy", S_total)
+        # Aggregate entropy components per molecule
+        results = {}
+
+        for _, row in self._residue_results_df.iterrows():
+            entropy_type = row["Type"].split()[0]
+            value = row["Result"]
+
+            if entropy_type == "Orientational":
+                mol_id = row["Molecule ID"]
+            else:
+                mol_id = row["Residue"].split("_")[0]
+
+            if mol_id not in results:
+                results[mol_id] = {
+                    "Orientational": 0.0,
+                    "Transvibrational": 0.0,
+                    "Rovibrational": 0.0,
+                }
+
+            results[mol_id][entropy_type] += value
+
+        # Log per-molecule entropy components and total
+        for mol_id, components in results.items():
+            total = 0.0
+            for entropy_type in ["Orientational", "Transvibrational", "Rovibrational"]:
+                S_component = components[entropy_type]
+                self._log_result(mol_id, "water", entropy_type, S_component)
+                total += S_component
+
+            self._log_result(mol_id, "Molecule Total", "Molecule Total Entropy", total)
 
     def _calculate_water_orientational_entropy(self, Sorient_dict):
         """

tests/test_CodeEntropy/test_entropy.py

@@ -5,6 +5,7 @@
 from unittest.mock import MagicMock, call, patch
 
 import numpy as np
+import pandas as pd
 import pytest
 
 from CodeEntropy.entropy import EntropyManager, VibrationalEntropy
@@ -258,6 +259,67 @@ def test_calculate_water_entropy(self, mock_get_entropy):
         )
         self.entropy_manager._log_result.assert_not_called()
 
+    @patch(
+        "waterEntropy.recipes.interfacial_solvent.get_interfacial_water_orient_entropy"
+    )
+    def test_calculate_water_entropy_minimal(self, mock_get_entropy):
+        """
+        This twst verifies that _calculate_water_entropy correctly logs
+        entropy components and total for a single molecule with minimal data.
+        """
+        self.entropy_manager._log_residue_data = MagicMock()
+        self.entropy_manager._log_result = MagicMock()
+
+        # Minimal mocked return from get_interfacial_water_orient_entropy
+        mock_get_entropy.return_value = (
+            {},  # Sorient_dict (not used here)
+            None,
+            MagicMock(
+                translational_S={("ACE_1", "WAT"): 10.0},
+                rotational_S={("ACE_1", "WAT"): 2.0},
+            ),
+            None,
+        )
+
+        # Minimal internal state
+        self.entropy_manager._residue_results_df = pd.DataFrame(
+            [
+                {
+                    "Molecule ID": "ACE",
+                    "Residue": "1",
+                    "Type": "Orientational (J/mol/K)",
+                    "Result": 5.0,
+                },
+                {
+                    "Molecule ID": "WAT",
+                    "Residue": "ACE_1",
+                    "Type": "Transvibrational (J/mol/K)",
+                    "Result": 10.0,
+                },
+                {
+                    "Molecule ID": "WAT",
+                    "Residue": "ACE_1",
+                    "Type": "Rovibrational (J/mol/K)",
+                    "Result": 2.0,
+                },
+            ]
+        )
+
+        # Call the real method
+        mock_universe = MagicMock()
+        self.entropy_manager._calculate_water_entropy(mock_universe, 0, 10, 1)
+
+        # Assert that only ACE is logged with correct values
+        self.entropy_manager._log_result.assert_has_calls(
+            [
+                call("ACE", "water", "Orientational", 5.0),
+                call("ACE", "water", "Transvibrational", 10.0),
+                call("ACE", "water", "Rovibrational", 2.0),
+                call("ACE", "Molecule Total", "Molecule Total Entropy", 17.0),
+            ],
+            any_order=False,
+        )
+
     # TODO test for error handling on invalid inputs