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

Author: harryswift01

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}")

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):

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()