From bc4e40022d39f4cb1077ea9ac8f781d86f3c48c6 Mon Sep 17 00:00:00 2001
From: harryswift01 <harry.swift@stfc.ac.uk>
Date: Thu, 12 Jun 2025 08:50:52 +0100
Subject: [PATCH 1/9] Integration of `waterEntropy` into `CodeEntropy`: -
 Automatically detecting any water molecules in the universe and calling
 `_calculate_water_entropy` function and changing the `selection_string` to
 `not water` - `waterEntropy` function is called within
 `_calculate_water_entropy` and results are logged into `data_logger` -
 Removal of the `water_entropy` keyword in the `arg_config_manager` as
 automatic detection is now used - Additional test cases to check the
 functions are being called in the correct order and sequence - Updated
 `pyproject.toml` to ensure dependencies are up to date

---
 CodeEntropy/config/arg_config_manager.py |   5 -
 CodeEntropy/entropy.py                   |  76 +++++++++++-
 config.yaml                              |   1 -
 pyproject.toml                           |   3 +-
 tests/test_CodeEntropy/test_entropy.py   | 140 ++++++++++++++++++++++-
 5 files changed, 212 insertions(+), 13 deletions(-)

diff --git a/CodeEntropy/config/arg_config_manager.py b/CodeEntropy/config/arg_config_manager.py
index 89987b4..dda4cc2 100644
--- a/CodeEntropy/config/arg_config_manager.py
+++ b/CodeEntropy/config/arg_config_manager.py
@@ -59,11 +59,6 @@
         "default": "output_file.json",
     },
     "force_partitioning": {"type": float, "help": "Force partitioning", "default": 0.5},
-    "water_entropy": {
-        "type": bool,
-        "help": "Calculate water entropy",
-        "default": False,
-    },
 }
 
 
diff --git a/CodeEntropy/entropy.py b/CodeEntropy/entropy.py
index d4975cc..4bff031 100644
--- a/CodeEntropy/entropy.py
+++ b/CodeEntropy/entropy.py
@@ -3,6 +3,7 @@
 
 import numpy as np
 import pandas as pd
+import waterEntropy.recipes.interfacial_solvent as GetSolvent
 from numpy import linalg as la
 
 logger = logging.getLogger(__name__)
@@ -59,6 +60,11 @@ def execute(self):
         """
         start, end, step = self._get_trajectory_bounds()
         number_frames = self._get_number_frames(start, end, step)
+
+        if self._universe.select_atoms("water").n_atoms > 0:
+            self._calculate_water_entropy(self._universe, start, end, step)
+            self._args.selection_string = "not water"
+
         reduced_atom = self._get_reduced_universe()
         number_molecules, levels = self._level_manager.select_levels(reduced_atom)
 
@@ -240,9 +246,15 @@ def _process_united_atom_level(
             S_rot += S_rot_res
             S_conf += S_conf_res
 
-            self._log_residue_data(mol_id, residue_id, "Transvibrational", S_trans_res)
-            self._log_residue_data(mol_id, residue_id, "Rovibrational", S_rot_res)
-            self._log_residue_data(mol_id, residue_id, "Conformational", S_conf_res)
+            self._log_residue_data(
+                residue.resid, residue.resname, "Transvibrational", S_trans_res
+            )
+            self._log_residue_data(
+                residue.resid, residue.resname, "Rovibrational", S_rot_res
+            )
+            self._log_residue_data(
+                residue.resid, residue.resname, "Conformational", S_conf_res
+            )
 
         self._log_result(mol_id, level, "Transvibrational", S_trans)
         self._log_result(mol_id, level, "Rovibrational", S_rot)
@@ -334,7 +346,8 @@ def _log_result(self, mol_id, level, entropy_type, value):
                 "Result": [value],
             }
         )
-        self._results_df = pd.concat([self._results_df, row], ignore_index=True)
+        if self.results_df.empty:
+            self._results_df = pd.concat([self._results_df, row], ignore_index=True)
         self._data_logger.add_results_data(mol_id, level, entropy_type, value)
 
     def _log_residue_data(self, mol_id, residue_id, entropy_type, value):
@@ -360,6 +373,61 @@ def _log_residue_data(self, mol_id, residue_id, entropy_type, value):
         )
         self._data_logger.add_residue_data(mol_id, residue_id, entropy_type, value)
 
+    def _calculate_water_entropy(self, universe, start, end, step):
+        """
+        Calculates orientational and vibrational entropy for water molecules.
+
+        Args:
+            universe: MDAnalysis Universe object.
+            start (int): Start frame.
+            end (int): End frame.
+            step (int): Step size.
+        """
+        Sorient_dict, _, vibrations, _ = (
+            GetSolvent.get_interfacial_water_orient_entropy(universe, start, end, step)
+        )
+
+        # Log per-residue entropy using helper functions
+        self._calculate_water_orientational_entropy(Sorient_dict)
+        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)
+
+    def _calculate_water_orientational_entropy(self, Sorient_dict):
+        """
+        Logs orientational entropy values directly from Sorient_dict.
+        """
+        for resid, resname_dict in Sorient_dict.items():
+            for resname, values in resname_dict.items():
+                if isinstance(values, list) and len(values) == 2:
+                    Sor, count = values
+                    self._log_residue_data(resname, resid, "Orientational", Sor)
+
+    def _calculate_water_vibrational_translational_entropy(self, vibrations):
+        """
+        Logs summed translational entropy values per residue-solvent pair.
+        """
+        for (resid, mol_id), entropy in vibrations.translational_S.items():
+            if isinstance(entropy, (list, np.ndarray)):
+                entropy = float(np.sum(entropy))
+            self._log_residue_data(mol_id, f"{resid}", "Transvibrational", entropy)
+
+    def _calculate_water_vibrational_rotational_entropy(self, vibrations):
+        """
+        Logs summed rotational entropy values per residue-solvent pair.
+        """
+        for (resid, mol_id), entropy in vibrations.rotational_S.items():
+            if isinstance(entropy, (list, np.ndarray)):
+                entropy = float(np.sum(entropy))
+            self._log_residue_data(mol_id, f"{resid}", "Rovibrational", entropy)
+
 
 class VibrationalEntropy(EntropyManager):
     """
diff --git a/config.yaml b/config.yaml
index 30c0133..4053a95 100644
--- a/config.yaml
+++ b/config.yaml
@@ -12,4 +12,3 @@ run1:
   thread: 
   output_file: 
   force_partitioning:
-  water_entropy:
diff --git a/pyproject.toml b/pyproject.toml
index e714acd..3bcd90a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -38,7 +38,8 @@ dependencies = [
     "psutil==5.9.5",
     "PyYAML==6.0.2",
     "python-json-logger==3.3.0",
-    "tabulate==0.9.0"
+    "tabulate==0.9.0",
+    "waterEntropy==1.0.2"
 ]
 
 [project.urls]
diff --git a/tests/test_CodeEntropy/test_entropy.py b/tests/test_CodeEntropy/test_entropy.py
index 7522cfe..88798cb 100644
--- a/tests/test_CodeEntropy/test_entropy.py
+++ b/tests/test_CodeEntropy/test_entropy.py
@@ -2,12 +2,12 @@
 import shutil
 import tempfile
 import unittest
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, call, patch
 
 import numpy as np
 import pytest
 
-from CodeEntropy.entropy import VibrationalEntropy
+from CodeEntropy.entropy import EntropyManager, VibrationalEntropy
 from CodeEntropy.main import main
 from CodeEntropy.run import RunManager
 
@@ -28,6 +28,10 @@ def setUp(self):
         self._orig_dir = os.getcwd()
         os.chdir(self.test_dir)
 
+        self.entropy_manager = EntropyManager(
+            MagicMock(), MagicMock(), MagicMock(), MagicMock(), MagicMock()
+        )
+
     def tearDown(self):
         """
         Clean up after each test.
@@ -122,6 +126,138 @@ def test_vibrational_entropy_polymer_torque(self):
 
         assert S_vib == pytest.approx(48.45003266069881)
 
+    def test_calculate_water_orientational_entropy(self):
+        """
+        Test that orientational entropy values are correctly extracted from Sorient_dict
+        and logged using _log_residue_data. Verifies that the entropy values are passed
+        correctly and that no molecule-level aggregation is performed unless
+        implemented.
+        """
+        self.entropy_manager._log_residue_data = MagicMock()
+        self.entropy_manager._log_result = MagicMock()
+
+        Sorient_dict = {1: {"mol1": [1.0, 2]}, 2: {"mol1": [3.0, 4]}}
+
+        self.entropy_manager._calculate_water_orientational_entropy(Sorient_dict)
+
+        self.entropy_manager._log_residue_data.assert_has_calls(
+            [
+                call("mol1", 1, "Orientational", 1.0),
+                call("mol1", 2, "Orientational", 3.0),
+            ]
+        )
+
+    def test_calculate_water_vibrational_translational_entropy(self):
+        """
+        Test that translational vibrational entropy values are correctly summed
+        and logged per residue using _log_residue_data. Also verifies that the
+        molecule-level average is computed and logged using _log_result.
+        """
+        self.entropy_manager._log_residue_data = MagicMock()
+        self.entropy_manager._log_result = MagicMock()
+
+        mock_vibrations = MagicMock()
+        mock_vibrations.translational_S = {
+            ("res1", "mol1"): [1.0, 2.0],
+            ("res2", "mol1"): 3.0,
+        }
+
+        self.entropy_manager._calculate_water_vibrational_translational_entropy(
+            mock_vibrations
+        )
+
+        self.entropy_manager._log_residue_data.assert_has_calls(
+            [
+                call("mol1", "res1", "Transvibrational", 3.0),
+                call("mol1", "res2", "Transvibrational", 3.0),
+            ]
+        )
+
+    def test_calculate_water_vibrational_rotational_entropy(self):
+        """
+        Test that rotational vibrational entropy values are correctly summed
+        and logged per residue using _log_residue_data. Also verifies that the
+        molecule-level average is computed and logged using _log_result.
+        """
+        self.entropy_manager._log_residue_data = MagicMock()
+        self.entropy_manager._log_result = MagicMock()
+
+        mock_vibrations = MagicMock()
+        mock_vibrations.rotational_S = {
+            ("res1", "mol1"): [2.0, 4.0],
+            ("res2", "mol1"): 6.0,
+        }
+
+        self.entropy_manager._calculate_water_vibrational_rotational_entropy(
+            mock_vibrations
+        )
+
+        self.entropy_manager._log_residue_data.assert_has_calls(
+            [
+                call("mol1", "res1", "Rovibrational", 6.0),
+                call("mol1", "res2", "Rovibrational", 6.0),
+            ]
+        )
+
+    def test_empty_vibrational_entropy_dicts(self):
+        """
+        Test that no logging occurs when both translational and rotational
+        entropy dictionaries are empty. Ensures that the methods handle empty
+        input gracefully without errors or unnecessary logging.
+        """
+        self.entropy_manager._log_residue_data = MagicMock()
+        self.entropy_manager._log_result = MagicMock()
+
+        mock_vibrations = MagicMock()
+        mock_vibrations.translational_S = {}
+        mock_vibrations.rotational_S = {}
+
+        self.entropy_manager._calculate_water_vibrational_translational_entropy(
+            mock_vibrations
+        )
+        self.entropy_manager._calculate_water_vibrational_rotational_entropy(
+            mock_vibrations
+        )
+
+        self.entropy_manager._log_residue_data.assert_not_called()
+        self.entropy_manager._log_result.assert_not_called()
+
+    @patch(
+        "waterEntropy.recipes.interfacial_solvent.get_interfacial_water_orient_entropy"
+    )
+    def test_calculate_water_entropy(self, mock_get_entropy):
+        """
+        Integration-style test that verifies _calculate_water_entropy correctly
+        delegates to the orientational and vibrational entropy methods and logs
+        the expected values. Uses a mocked return from
+        get_interfacial_water_orient_entropy.
+        """
+        self.entropy_manager._log_residue_data = MagicMock()
+        self.entropy_manager._log_result = MagicMock()
+
+        mock_vibrations = MagicMock()
+        mock_vibrations.translational_S = {("res1", "mol1"): 2.0}
+        mock_vibrations.rotational_S = {("res1", "mol1"): 3.0}
+
+        mock_get_entropy.return_value = (
+            {1: {"mol1": [1.0, 5]}},
+            None,
+            mock_vibrations,
+            None,
+        )
+
+        mock_universe = MagicMock()
+        self.entropy_manager._calculate_water_entropy(mock_universe, 0, 10, 1)
+
+        self.entropy_manager._log_residue_data.assert_has_calls(
+            [
+                call("mol1", 1, "Orientational", 1.0),
+                call("mol1", "res1", "Transvibrational", 2.0),
+                call("mol1", "res1", "Rovibrational", 3.0),
+            ]
+        )
+        self.entropy_manager._log_result.assert_not_called()
+
     # TODO test for error handling on invalid inputs
 
 

From e744dc00c47de5a259282756da23ec137231c229 Mon Sep 17 00:00:00 2001
From: harryswift01 <harry.swift@stfc.ac.uk>
Date: Fri, 13 Jun 2025 16:27:36 +0100
Subject: [PATCH 2/9] 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

---
 CodeEntropy/entropy.py                 | 62 ++++++++++++++++++--------
 tests/test_CodeEntropy/test_entropy.py | 62 ++++++++++++++++++++++++++
 2 files changed, 106 insertions(+), 18 deletions(-)

diff --git a/CodeEntropy/entropy.py b/CodeEntropy/entropy.py
index 4bff031..b334ffb 100644
--- a/CodeEntropy/entropy.py
+++ b/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):
         """
diff --git a/tests/test_CodeEntropy/test_entropy.py b/tests/test_CodeEntropy/test_entropy.py
index 88798cb..f39d789 100644
--- a/tests/test_CodeEntropy/test_entropy.py
+++ b/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
 
 

From f5c201b972b263321f24ec7ef521ec49dd5cafce Mon Sep 17 00:00:00 2001
From: harryswift01 <harry.swift@stfc.ac.uk>
Date: Fri, 13 Jun 2025 16:34:04 +0100
Subject: [PATCH 3/9] small refinements to the comments within
 `test_entropy.py` to tidy it up

---
 tests/test_CodeEntropy/test_entropy.py | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/tests/test_CodeEntropy/test_entropy.py b/tests/test_CodeEntropy/test_entropy.py
index f39d789..a30fd3a 100644
--- a/tests/test_CodeEntropy/test_entropy.py
+++ b/tests/test_CodeEntropy/test_entropy.py
@@ -270,9 +270,8 @@ def test_calculate_water_entropy_minimal(self, mock_get_entropy):
         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},
@@ -281,7 +280,6 @@ def test_calculate_water_entropy_minimal(self, mock_get_entropy):
             None,
         )
 
-        # Minimal internal state
         self.entropy_manager._residue_results_df = pd.DataFrame(
             [
                 {
@@ -305,11 +303,9 @@ def test_calculate_water_entropy_minimal(self, mock_get_entropy):
             ]
         )
 
-        # 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),

From a4cad6742808e230f52274cd1274eb131dab5363 Mon Sep 17 00:00:00 2001
From: harryswift01 <harry.swift@stfc.ac.uk>
Date: Mon, 16 Jun 2025 08:27:41 +0100
Subject: [PATCH 4/9] Additional refinments to ensure consistency across
 `CodeEntropy` - Remove duplicated adding and logging results, using
 `DataLogger` class to exclusivly handle all of the data handling -
 Restructured how results are diplayed, users will now see the `resname`
 rather than an arbitary internal counter to differentiate the molecules -
 Changing of `_args.selection_string` to ensure user's input is not
 overwritten if selection is not `all` - Updated test cases to reflect the
 changes made within this commit

---
 CodeEntropy/config/data_logger.py          |  29 ++-
 CodeEntropy/entropy.py                     | 207 ++++++++++-----------
 tests/test_CodeEntropy/test_data_logger.py |  17 +-
 tests/test_CodeEntropy/test_entropy.py     | 141 ++++++--------
 4 files changed, 191 insertions(+), 203 deletions(-)

diff --git a/CodeEntropy/config/data_logger.py b/CodeEntropy/config/data_logger.py
index 483db7d..223a66c 100644
--- a/CodeEntropy/config/data_logger.py
+++ b/CodeEntropy/config/data_logger.py
@@ -1,5 +1,6 @@
 import json
 import logging
+import re
 
 from tabulate import tabulate
 
@@ -23,13 +24,19 @@ def save_dataframes_as_json(self, molecule_df, residue_df, output_file):
         with open(output_file, "w") as out:
             json.dump(data, out, indent=4)
 
-    def add_results_data(self, molecule, level, type, S_molecule):
+    def clean_residue_name(self, resname):
+        """Ensures residue names are stripped and cleaned before being stored"""
+        return re.sub(r"[-–—]", "", str(resname))
+
+    def add_results_data(self, resname, level, entropy_type, value):
         """Add data for molecule-level entries"""
-        self.molecule_data.append([molecule, level, type, f"{S_molecule}"])
+        resname = self.clean_residue_name(resname)
+        self.molecule_data.append((resname, level, entropy_type, value))
 
-    def add_residue_data(self, molecule, residue, type, S_trans_residue):
+    def add_residue_data(self, resid, resname, level, entropy_type, value):
         """Add data for residue-level entries"""
-        self.residue_data.append([molecule, residue, type, f"{S_trans_residue}"])
+        resname = self.clean_residue_name(resname)
+        self.residue_data.append([resid, resname, level, entropy_type, value])
 
     def log_tables(self):
         """Log both tables at once"""
@@ -38,8 +45,10 @@ def log_tables(self):
             logger.info("Molecule Data Table:")
             table_str = tabulate(
                 self.molecule_data,
-                headers=["Molecule ID", "Level", "Type", "Result (J/mol/K)"],
+                headers=["Residue Name", "Level", "Type", "Result (J/mol/K)"],
                 tablefmt="grid",
+                numalign="center",
+                stralign="center",
             )
             logger.info(f"\n{table_str}")
 
@@ -48,7 +57,15 @@ def log_tables(self):
             logger.info("Residue Data Table:")
             table_str = tabulate(
                 self.residue_data,
-                headers=["Molecule ID", "Residue", "Type", "Result (J/mol/K)"],
+                headers=[
+                    "Residue ID",
+                    "Residue Name",
+                    "Level",
+                    "Type",
+                    "Result (J/mol/K)",
+                ],
                 tablefmt="grid",
+                numalign="center",
+                stralign="center",
             )
             logger.info(f"\n{table_str}")
diff --git a/CodeEntropy/entropy.py b/CodeEntropy/entropy.py
index b334ffb..a9dcb72 100644
--- a/CodeEntropy/entropy.py
+++ b/CodeEntropy/entropy.py
@@ -1,5 +1,6 @@
 import logging
 import math
+from collections import defaultdict
 
 import numpy as np
 import pandas as pd
@@ -33,25 +34,6 @@ def __init__(self, run_manager, args, universe, data_logger, level_manager):
         self._level_manager = level_manager
         self._GAS_CONST = 8.3144598484848
 
-        self._results_df = pd.DataFrame(
-            columns=["Molecule ID", "Level", "Type", "Result"]
-        )
-        self._residue_results_df = pd.DataFrame(
-            columns=["Molecule ID", "Residue", "Type", "Result"]
-        )
-
-    @property
-    def results_df(self):
-        """Returns the dataframe containing entropy results at all levels."""
-        return self._results_df
-
-    @property
-    def residue_results_df(self):
-        """
-        Returns the dataframe containing united-atom level results for each residue.
-        """
-        return self._residue_results_df
-
     def execute(self):
         """
         Executes the full entropy computation workflow over selected molecules and
@@ -63,7 +45,11 @@ def execute(self):
 
         if self._universe.select_atoms("water").n_atoms > 0:
             self._calculate_water_entropy(self._universe, start, end, step)
-            self._args.selection_string = "not water"
+
+            if self._args.selection_string != "all":
+                self._args.selection_string += " and not water"
+            else:
+                self._args.selection_string = "not water"
 
         reduced_atom = self._get_reduced_universe()
         number_molecules, levels = self._level_manager.select_levels(reduced_atom)
@@ -246,19 +232,25 @@ def _process_united_atom_level(
             S_rot += S_rot_res
             S_conf += S_conf_res
 
-            self._log_residue_data(
-                residue.resid, residue.resname, "Transvibrational", S_trans_res
+            self._data_logger.add_residue_data(
+                mol_id, residue.resname, level, "Transvibrational", S_trans_res
             )
-            self._log_residue_data(
-                residue.resid, residue.resname, "Rovibrational", S_rot_res
+            self._data_logger.add_residue_data(
+                mol_id, residue.resname, level, "Rovibrational", S_rot_res
             )
-            self._log_residue_data(
-                residue.resid, residue.resname, "Conformational", S_conf_res
+            self._data_logger.add_residue_data(
+                mol_id, residue.resname, level, "Conformational", S_conf_res
             )
 
-        self._log_result(mol_id, level, "Transvibrational", S_trans)
-        self._log_result(mol_id, level, "Rovibrational", S_rot)
-        self._log_result(mol_id, level, "Conformational", S_conf)
+        self._data_logger.add_results_data(
+            residue.resname, level, "Transvibrational", S_trans
+        )
+        self._data_logger.add_results_data(
+            residue.resname, level, "Rovibrational", S_rot
+        )
+        self._data_logger.add_results_data(
+            residue.resname, level, "Conformational", S_conf
+        )
 
     def _process_vibrational_only_levels(
         self, mol_id, mol_container, ve, level, start, end, step, n_frames, highest
@@ -286,9 +278,13 @@ def _process_vibrational_only_levels(
         S_rot = ve.vibrational_entropy_calculation(
             torque_matrix, "torque", self._args.temperature, highest
         )
-
-        self._log_result(mol_id, level, "Transvibrational", S_trans)
-        self._log_result(mol_id, level, "Rovibrational", S_rot)
+        residue = mol_container.residues[mol_id]
+        self._data_logger.add_results_data(
+            residue.resname, level, "Transvibrational", S_trans
+        )
+        self._data_logger.add_results_data(
+            residue.resname, level, "Rovibrational", S_rot
+        )
 
     def _process_conformational_residue_level(
         self, mol_id, mol_container, ce, level, start, end, step, n_frames
@@ -310,71 +306,48 @@ def _process_conformational_residue_level(
         S_conf = ce.conformational_entropy_calculation(
             mol_container, dihedrals, bin_width, start, end, step, n_frames
         )
-        self._log_result(mol_id, level, "Conformational", S_conf)
+        residue = mol_container.residues[mol_id]
+        self._data_logger.add_results_data(
+            residue.resname, level, "Conformational", S_conf
+        )
 
     def _finalize_molecule_results(self):
         """
-        Summarizes entropy for a molecule and saves results to file.
-
-        Args:
-            mol_id (int): ID of the molecule.
-        """
-        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):
+        Aggregates and logs total entropy per molecule using residue_data grouped by
+        resid.
         """
-        Logs and stores a single entropy value in the global results dataframe.
+        entropy_by_molecule = defaultdict(float)
 
-        Args:
-            mol_id (int): Molecule ID.
-            level (str): Entropy level or type.
-            entropy_type (str): Type of entropy (e.g., 'Transvibrational').
-            value (float): Entropy value.
-        """
-        row = pd.DataFrame(
-            {
-                "Molecule ID": [mol_id],
-                "Level": [level],
-                "Type": [f"{entropy_type} (J/mol/K)"],
-                "Result": [value],
-            }
-        )
-        if self.results_df.empty:
-            self._results_df = pd.concat([self._results_df, row], ignore_index=True)
-        self._data_logger.add_results_data(mol_id, level, entropy_type, value)
+        for mol_id, level, entropy_type, result in self._data_logger.molecule_data:
+            if level != "Molecule Total":
+                try:
+                    entropy_by_molecule[mol_id] += float(result)
+                except ValueError:
+                    logger.warning(f"Skipping invalid entry: {mol_id}, {result}")
 
-    def _log_residue_data(self, mol_id, residue_id, entropy_type, value):
-        """
-        Logs and stores per-residue entropy data.
+        for mol_id, total_entropy in entropy_by_molecule.items():
+            self._data_logger.molecule_data.append(
+                (mol_id, "Molecule Total", "Molecule Total Entropy", total_entropy)
+            )
 
-        Args:
-            mol_id (int): Molecule ID.
-            residue_id (int): Residue index within the molecule.
-            entropy_type (str): Entropy category.
-            value (float): Entropy value.
-        """
-        row = pd.DataFrame(
-            {
-                "Molecule ID": [mol_id],
-                "Residue": [residue_id],
-                "Type": [f"{entropy_type} (J/mol/K)"],
-                "Result": [value],
-            }
-        )
-        self._residue_results_df = pd.concat(
-            [self._residue_results_df, row], ignore_index=True
+        # Save to file
+        self._data_logger.save_dataframes_as_json(
+            pd.DataFrame(
+                self._data_logger.molecule_data,
+                columns=["Molecule ID", "Level", "Type", "Result (J/mol/K)"],
+            ),
+            pd.DataFrame(
+                self._data_logger.residue_data,
+                columns=[
+                    "Residue ID",
+                    "Residue Name",
+                    "Level",
+                    "Type",
+                    "Result (J/mol/K)",
+                ],
+            ),
+            self._args.output_file,
         )
-        self._data_logger.add_residue_data(mol_id, residue_id, entropy_type, value)
 
     def _calculate_water_entropy(self, universe, start, end, step):
         """
@@ -398,14 +371,10 @@ def _calculate_water_entropy(self, universe, start, end, step):
         # 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]
+        for row in self._data_logger.residue_data:
+            mol_id = row[1]
+            entropy_type = row[3].split()[0]
+            value = float(row[4])
 
             if mol_id not in results:
                 results[mol_id] = {
@@ -421,11 +390,11 @@ def _calculate_water_entropy(self, universe, start, end, step):
             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)
+                self._data_logger.add_results_data(
+                    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):
         """
         Logs orientational entropy values directly from Sorient_dict.
@@ -434,25 +403,53 @@ def _calculate_water_orientational_entropy(self, Sorient_dict):
             for resname, values in resname_dict.items():
                 if isinstance(values, list) and len(values) == 2:
                     Sor, count = values
-                    self._log_residue_data(resname, resid, "Orientational", Sor)
+                    self._data_logger.add_residue_data(
+                        resid, resname, "Water", "Orientational", Sor
+                    )
 
     def _calculate_water_vibrational_translational_entropy(self, vibrations):
         """
         Logs summed translational entropy values per residue-solvent pair.
         """
-        for (resid, mol_id), entropy in vibrations.translational_S.items():
+        for (solute_id, _), entropy in vibrations.translational_S.items():
             if isinstance(entropy, (list, np.ndarray)):
                 entropy = float(np.sum(entropy))
-            self._log_residue_data(mol_id, f"{resid}", "Transvibrational", entropy)
+
+            if "_" in solute_id:
+                resname, resid_str = solute_id.rsplit("_", 1)
+                try:
+                    resid = int(resid_str)
+                except ValueError:
+                    resid = -1
+            else:
+                resname = solute_id
+                resid = -1
+
+            self._data_logger.add_residue_data(
+                resid, resname, "Water", "Transvibrational", entropy
+            )
 
     def _calculate_water_vibrational_rotational_entropy(self, vibrations):
         """
         Logs summed rotational entropy values per residue-solvent pair.
         """
-        for (resid, mol_id), entropy in vibrations.rotational_S.items():
+        for (solute_id, _), entropy in vibrations.rotational_S.items():
             if isinstance(entropy, (list, np.ndarray)):
                 entropy = float(np.sum(entropy))
-            self._log_residue_data(mol_id, f"{resid}", "Rovibrational", entropy)
+
+            if "_" in solute_id:
+                resname, resid_str = solute_id.rsplit("_", 1)
+                try:
+                    resid = int(resid_str)
+                except ValueError:
+                    resid = -1
+            else:
+                resname = solute_id
+                resid = -1
+
+            self._data_logger.add_residue_data(
+                resid, resname, "Water", "Rovibrational", entropy
+            )
 
 
 class VibrationalEntropy(EntropyManager):
diff --git a/tests/test_CodeEntropy/test_data_logger.py b/tests/test_CodeEntropy/test_data_logger.py
index 803b95f..782b0b6 100644
--- a/tests/test_CodeEntropy/test_data_logger.py
+++ b/tests/test_CodeEntropy/test_data_logger.py
@@ -53,20 +53,23 @@ def test_add_results_data(self):
         Test that add_results_data correctly appends a molecule-level entry.
         """
         self.logger.add_results_data(
-            0, "united_atom", "Transvibrational (J/mol/K)", 653.404
+            0, "united_atom", "Transvibrational", 653.4041220313459
         )
         self.assertEqual(
             self.logger.molecule_data,
-            [[0, "united_atom", "Transvibrational (J/mol/K)", "653.404"]],
+            [("0", "united_atom", "Transvibrational", 653.4041220313459)],
         )
 
     def test_add_residue_data(self):
         """
         Test that add_residue_data correctly appends a residue-level entry.
         """
-        self.logger.add_residue_data(0, 0, "Transvibrational (J/mol/K)", 122.612)
+        self.logger.add_residue_data(
+            0, "DA", "united_atom", "Transvibrational", 122.61216935211893
+        )
         self.assertEqual(
-            self.logger.residue_data, [[0, 0, "Transvibrational (J/mol/K)", "122.612"]]
+            self.logger.residue_data,
+            [[0, "DA", "united_atom", "Transvibrational", 122.61216935211893]],
         )
 
     def test_save_dataframes_as_json(self):
@@ -124,9 +127,11 @@ def test_log_tables(self, mock_logger):
         logger.
         """
         self.logger.add_results_data(
-            0, "united_atom", "Transvibrational (J/mol/K)", 653.404
+            0, "united_atom", "Transvibrational", 653.4041220313459
+        )
+        self.logger.add_residue_data(
+            0, "DA", "united_atom", "Transvibrational", 122.61216935211893
         )
-        self.logger.add_residue_data(0, 0, "Transvibrational (J/mol/K)", 122.612)
 
         self.logger.log_tables()
 
diff --git a/tests/test_CodeEntropy/test_entropy.py b/tests/test_CodeEntropy/test_entropy.py
index a30fd3a..47a66b2 100644
--- a/tests/test_CodeEntropy/test_entropy.py
+++ b/tests/test_CodeEntropy/test_entropy.py
@@ -5,7 +5,6 @@
 from unittest.mock import MagicMock, call, patch
 
 import numpy as np
-import pandas as pd
 import pytest
 
 from CodeEntropy.entropy import EntropyManager, VibrationalEntropy
@@ -130,36 +129,29 @@ def test_vibrational_entropy_polymer_torque(self):
     def test_calculate_water_orientational_entropy(self):
         """
         Test that orientational entropy values are correctly extracted from Sorient_dict
-        and logged using _log_residue_data. Verifies that the entropy values are passed
-        correctly and that no molecule-level aggregation is performed unless
-        implemented.
+        and logged using add_residue_data.
         """
-        self.entropy_manager._log_residue_data = MagicMock()
-        self.entropy_manager._log_result = MagicMock()
-
         Sorient_dict = {1: {"mol1": [1.0, 2]}, 2: {"mol1": [3.0, 4]}}
 
         self.entropy_manager._calculate_water_orientational_entropy(Sorient_dict)
 
-        self.entropy_manager._log_residue_data.assert_has_calls(
+        self.entropy_manager._data_logger.add_residue_data.assert_has_calls(
             [
-                call("mol1", 1, "Orientational", 1.0),
-                call("mol1", 2, "Orientational", 3.0),
+                call(1, "mol1", "Water", "Orientational", 1.0),
+                call(2, "mol1", "Water", "Orientational", 3.0),
             ]
         )
 
     def test_calculate_water_vibrational_translational_entropy(self):
         """
         Test that translational vibrational entropy values are correctly summed
-        and logged per residue using _log_residue_data. Also verifies that the
+        and logged per residue using add_residue_data. Also verifies that the
         molecule-level average is computed and logged using _log_result.
         """
-        self.entropy_manager._log_residue_data = MagicMock()
-        self.entropy_manager._log_result = MagicMock()
-
         mock_vibrations = MagicMock()
         mock_vibrations.translational_S = {
             ("res1", "mol1"): [1.0, 2.0],
+            ("resB_invalid", "mol1"): 4.0,
             ("res2", "mol1"): 3.0,
         }
 
@@ -167,36 +159,11 @@ def test_calculate_water_vibrational_translational_entropy(self):
             mock_vibrations
         )
 
-        self.entropy_manager._log_residue_data.assert_has_calls(
+        self.entropy_manager._data_logger.add_residue_data.assert_has_calls(
             [
-                call("mol1", "res1", "Transvibrational", 3.0),
-                call("mol1", "res2", "Transvibrational", 3.0),
-            ]
-        )
-
-    def test_calculate_water_vibrational_rotational_entropy(self):
-        """
-        Test that rotational vibrational entropy values are correctly summed
-        and logged per residue using _log_residue_data. Also verifies that the
-        molecule-level average is computed and logged using _log_result.
-        """
-        self.entropy_manager._log_residue_data = MagicMock()
-        self.entropy_manager._log_result = MagicMock()
-
-        mock_vibrations = MagicMock()
-        mock_vibrations.rotational_S = {
-            ("res1", "mol1"): [2.0, 4.0],
-            ("res2", "mol1"): 6.0,
-        }
-
-        self.entropy_manager._calculate_water_vibrational_rotational_entropy(
-            mock_vibrations
-        )
-
-        self.entropy_manager._log_residue_data.assert_has_calls(
-            [
-                call("mol1", "res1", "Rovibrational", 6.0),
-                call("mol1", "res2", "Rovibrational", 6.0),
+                call(-1, "res1", "Water", "Transvibrational", 3.0),
+                call(-1, "resB", "Water", "Transvibrational", 4.0),
+                call(-1, "res2", "Water", "Transvibrational", 3.0),
             ]
         )
 
@@ -223,6 +190,31 @@ def test_empty_vibrational_entropy_dicts(self):
         self.entropy_manager._log_residue_data.assert_not_called()
         self.entropy_manager._log_result.assert_not_called()
 
+    def test_calculate_water_vibrational_rotational_entropy(self):
+        """
+        Test that rotational vibrational entropy values are correctly summed
+        and logged per residue using add_residue_data. Also verifies that the
+        residue ID parsing handles both valid and invalid formats.
+        """
+        mock_vibrations = MagicMock()
+        mock_vibrations.rotational_S = {
+            ("resA_101", "mol1"): [2.0, 3.0],
+            ("resB_invalid", "mol1"): 4.0,
+            ("resC", "mol1"): 5.0,
+        }
+
+        self.entropy_manager._calculate_water_vibrational_rotational_entropy(
+            mock_vibrations
+        )
+
+        self.entropy_manager._data_logger.add_residue_data.assert_has_calls(
+            [
+                call(101, "resA", "Water", "Rovibrational", 5.0),
+                call(-1, "resB", "Water", "Rovibrational", 4.0),
+                call(-1, "resC", "Water", "Rovibrational", 5.0),
+            ]
+        )
+
     @patch(
         "waterEntropy.recipes.interfacial_solvent.get_interfacial_water_orient_entropy"
     )
@@ -230,12 +222,8 @@ def test_calculate_water_entropy(self, mock_get_entropy):
         """
         Integration-style test that verifies _calculate_water_entropy correctly
         delegates to the orientational and vibrational entropy methods and logs
-        the expected values. Uses a mocked return from
-        get_interfacial_water_orient_entropy.
+        the expected values.
         """
-        self.entropy_manager._log_residue_data = MagicMock()
-        self.entropy_manager._log_result = MagicMock()
-
         mock_vibrations = MagicMock()
         mock_vibrations.translational_S = {("res1", "mol1"): 2.0}
         mock_vibrations.rotational_S = {("res1", "mol1"): 3.0}
@@ -250,26 +238,22 @@ def test_calculate_water_entropy(self, mock_get_entropy):
         mock_universe = MagicMock()
         self.entropy_manager._calculate_water_entropy(mock_universe, 0, 10, 1)
 
-        self.entropy_manager._log_residue_data.assert_has_calls(
+        self.entropy_manager._data_logger.add_residue_data.assert_has_calls(
             [
-                call("mol1", 1, "Orientational", 1.0),
-                call("mol1", "res1", "Transvibrational", 2.0),
-                call("mol1", "res1", "Rovibrational", 3.0),
+                call(1, "mol1", "Water", "Orientational", 1.0),
+                call(-1, "res1", "Water", "Transvibrational", 2.0),
+                call(-1, "res1", "Water", "Rovibrational", 3.0),
             ]
         )
-        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.
+        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()
-
         mock_get_entropy.return_value = (
             {},
             None,
@@ -280,40 +264,25 @@ def test_calculate_water_entropy_minimal(self, mock_get_entropy):
             None,
         )
 
-        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,
-                },
-            ]
-        )
+        # Simulate residue-level results already collected
+        self.entropy_manager._data_logger.residue_data = [
+            [1, "ACE", "Water", "Orientational", 5.0],
+            [1, "ACE_1", "Water", "Transvibrational", 10.0],
+            [1, "ACE_1", "Water", "Rovibrational", 2.0],
+        ]
 
         mock_universe = MagicMock()
         self.entropy_manager._calculate_water_entropy(mock_universe, 0, 10, 1)
 
-        self.entropy_manager._log_result.assert_has_calls(
+        self.entropy_manager._data_logger.add_results_data.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,
+                call("ACE", "water", "Transvibrational", 0.0),
+                call("ACE", "water", "Rovibrational", 0.0),
+                call("ACE_1", "water", "Orientational", 0.0),
+                call("ACE_1", "water", "Transvibrational", 10.0),
+                call("ACE_1", "water", "Rovibrational", 2.0),
+            ]
         )
 
     # TODO test for error handling on invalid inputs

From eee173a219769f847ed86f954fa861b993eafb98 Mon Sep 17 00:00:00 2001
From: harryswift01 <harry.swift@stfc.ac.uk>
Date: Mon, 16 Jun 2025 13:17:49 +0100
Subject: [PATCH 5/9] updating the version of `waterEntropy` in
 `pyproject.toml` to the current version `v.1.0.3`

---
 pyproject.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index 3bcd90a..63e307f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -39,7 +39,7 @@ dependencies = [
     "PyYAML==6.0.2",
     "python-json-logger==3.3.0",
     "tabulate==0.9.0",
-    "waterEntropy==1.0.2"
+    "waterEntropy==1.0.3"
 ]
 
 [project.urls]

From adca1bf9d6ce420f8c45b434282c696582a5be9d Mon Sep 17 00:00:00 2001
From: harryswift01 <harry.swift@stfc.ac.uk>
Date: Tue, 17 Jun 2025 16:23:49 +0100
Subject: [PATCH 6/9] additioanl `logger.debug` statements to output in debug
 mode the `self._data_logger.molecule_data` and
 `self._data_logger.residue_data` after each level

---
 CodeEntropy/entropy.py | 38 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 38 insertions(+)

diff --git a/CodeEntropy/entropy.py b/CodeEntropy/entropy.py
index a9dcb72..fbbd371 100644
--- a/CodeEntropy/entropy.py
+++ b/CodeEntropy/entropy.py
@@ -51,6 +51,15 @@ def execute(self):
             else:
                 self._args.selection_string = "not water"
 
+            logger.debug(
+                "WaterEntropy: molecule_data: %s",
+                self._data_logger.molecule_data,
+            )
+            logger.debug(
+                "WaterEntropy: residue_data: %s",
+                self._data_logger.residue_data,
+            )
+
         reduced_atom = self._get_reduced_universe()
         number_molecules, levels = self._level_manager.select_levels(reduced_atom)
 
@@ -87,6 +96,16 @@ def execute(self):
                         number_frames,
                         highest_level,
                     )
+
+                    logger.debug(
+                        "United Atom Level: molecule_data: %s",
+                        self._data_logger.molecule_data,
+                    )
+                    logger.debug(
+                        "United Atom Level: residue_data: %s",
+                        self._data_logger.residue_data,
+                    )
+
                 elif level in ("polymer", "residue"):
                     self._process_vibrational_only_levels(
                         molecule_id,
@@ -99,6 +118,16 @@ def execute(self):
                         number_frames,
                         highest_level,
                     )
+
+                    logger.debug(
+                        "Vibrational Level: molecule_data: %s",
+                        self._data_logger.molecule_data,
+                    )
+                    logger.debug(
+                        "Vibrational Level: residue_data: %s",
+                        self._data_logger.residue_data,
+                    )
+
                 if level == "residue":
                     self._process_conformational_residue_level(
                         molecule_id,
@@ -111,6 +140,15 @@ def execute(self):
                         number_frames,
                     )
 
+                    logger.debug(
+                        "Confirmational Level: molecule_data: %s",
+                        self._data_logger.molecule_data,
+                    )
+                    logger.debug(
+                        "Confirmational Level: residue_data: %s",
+                        self._data_logger.residue_data,
+                    )
+
         self._finalize_molecule_results()
 
         self._data_logger.log_tables()

From dd01d9e0dfc9e21187704bb4d8a18241ce28187c Mon Sep 17 00:00:00 2001
From: harryswift01 <harry.swift@stfc.ac.uk>
Date: Tue, 17 Jun 2025 16:45:04 +0100
Subject: [PATCH 7/9] add back in the option to disable water entropy
 calculations

---
 CodeEntropy/config/arg_config_manager.py | 5 +++++
 CodeEntropy/entropy.py                   | 3 ++-
 config.yaml                              | 9 +++++----
 3 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/CodeEntropy/config/arg_config_manager.py b/CodeEntropy/config/arg_config_manager.py
index dda4cc2..8b6de77 100644
--- a/CodeEntropy/config/arg_config_manager.py
+++ b/CodeEntropy/config/arg_config_manager.py
@@ -59,6 +59,11 @@
         "default": "output_file.json",
     },
     "force_partitioning": {"type": float, "help": "Force partitioning", "default": 0.5},
+    "disable_water_entropy": {
+        "type": bool,
+        "help": "If set to True, disables the calculation of water entropy",
+        "default": False,
+    },
 }
 
 
diff --git a/CodeEntropy/entropy.py b/CodeEntropy/entropy.py
index fbbd371..dfdfc2c 100644
--- a/CodeEntropy/entropy.py
+++ b/CodeEntropy/entropy.py
@@ -43,7 +43,8 @@ def execute(self):
         start, end, step = self._get_trajectory_bounds()
         number_frames = self._get_number_frames(start, end, step)
 
-        if self._universe.select_atoms("water").n_atoms > 0:
+        has_water = self._universe.select_atoms("water").n_atoms > 0
+        if has_water and not self._args.disable_water_entropy:
             self._calculate_water_entropy(self._universe, start, end, step)
 
             if self._args.selection_string != "all":
diff --git a/config.yaml b/config.yaml
index 4053a95..2a7ba86 100644
--- a/config.yaml
+++ b/config.yaml
@@ -2,13 +2,14 @@
 
 run1:
   top_traj_file:
-  selection_string: 
+  selection_string:
   start:
-  end: 
+  end:
   step:
   bin_width:
   temperature:
   verbose:
-  thread: 
-  output_file: 
+  thread:
+  output_file:
   force_partitioning:
+  disable_water_entropy:

From 26f3d00ec430868a0d2f07acaba64e4970b6a047 Mon Sep 17 00:00:00 2001
From: harryswift01 <harry.swift@stfc.ac.uk>
Date: Wed, 18 Jun 2025 09:30:03 +0100
Subject: [PATCH 8/9] rename `disable_water_entropy` to `water_entropy` for
 improved clarity

---
 CodeEntropy/config/arg_config_manager.py | 6 +++---
 CodeEntropy/entropy.py                   | 2 +-
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/CodeEntropy/config/arg_config_manager.py b/CodeEntropy/config/arg_config_manager.py
index 8b6de77..1bd9fbf 100644
--- a/CodeEntropy/config/arg_config_manager.py
+++ b/CodeEntropy/config/arg_config_manager.py
@@ -59,10 +59,10 @@
         "default": "output_file.json",
     },
     "force_partitioning": {"type": float, "help": "Force partitioning", "default": 0.5},
-    "disable_water_entropy": {
+    "water_entropy": {
         "type": bool,
-        "help": "If set to True, disables the calculation of water entropy",
-        "default": False,
+        "help": "If set to False, disables the calculation of water entropy",
+        "default": True,
     },
 }
 
diff --git a/CodeEntropy/entropy.py b/CodeEntropy/entropy.py
index dfdfc2c..814ca1b 100644
--- a/CodeEntropy/entropy.py
+++ b/CodeEntropy/entropy.py
@@ -44,7 +44,7 @@ def execute(self):
         number_frames = self._get_number_frames(start, end, step)
 
         has_water = self._universe.select_atoms("water").n_atoms > 0
-        if has_water and not self._args.disable_water_entropy:
+        if has_water and self._args.water_entropy:
             self._calculate_water_entropy(self._universe, start, end, step)
 
             if self._args.selection_string != "all":

From e235cf3845cc2209f291c48a3b3b9fe6447ff114 Mon Sep 17 00:00:00 2001
From: harryswift01 <harry.swift@stfc.ac.uk>
Date: Wed, 18 Jun 2025 09:37:51 +0100
Subject: [PATCH 9/9] refined `logger.debug` statements for outputting the
 `molecule_data` and `residue_data` after each level

---
 CodeEntropy/entropy.py | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/CodeEntropy/entropy.py b/CodeEntropy/entropy.py
index 814ca1b..e06c72e 100644
--- a/CodeEntropy/entropy.py
+++ b/CodeEntropy/entropy.py
@@ -99,11 +99,13 @@ def execute(self):
                     )
 
                     logger.debug(
-                        "United Atom Level: molecule_data: %s",
+                        "%s level: molecule_data: %s",
+                        level,
                         self._data_logger.molecule_data,
                     )
                     logger.debug(
-                        "United Atom Level: residue_data: %s",
+                        "%s level: residue_data: %s",
+                        level,
                         self._data_logger.residue_data,
                     )
 
@@ -121,11 +123,13 @@ def execute(self):
                     )
 
                     logger.debug(
-                        "Vibrational Level: molecule_data: %s",
+                        "%s level: molecule_data: %s",
+                        level,
                         self._data_logger.molecule_data,
                     )
                     logger.debug(
-                        "Vibrational Level: residue_data: %s",
+                        "%s level: residue_data: %s",
+                        level,
                         self._data_logger.residue_data,
                     )
 
@@ -142,11 +146,13 @@ def execute(self):
                     )
 
                     logger.debug(
-                        "Confirmational Level: molecule_data: %s",
+                        "%s level: molecule_data: %s",
+                        level,
                         self._data_logger.molecule_data,
                     )
                     logger.debug(
-                        "Confirmational Level: residue_data: %s",
+                        "%s level: residue_data: %s",
+                        level,
                         self._data_logger.residue_data,
                     )