diff --git a/CodeEntropy/config/arg_config_manager.py b/CodeEntropy/config/arg_config_manager.py index 89987b4..1bd9fbf 100644 --- a/CodeEntropy/config/arg_config_manager.py +++ b/CodeEntropy/config/arg_config_manager.py @@ -61,8 +61,8 @@ "force_partitioning": {"type": float, "help": "Force partitioning", "default": 0.5}, "water_entropy": { "type": bool, - "help": "Calculate water entropy", - "default": False, + "help": "If set to False, disables the calculation of water entropy", + "default": True, }, } 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 d4975cc..e06c72e 100644 --- a/CodeEntropy/entropy.py +++ b/CodeEntropy/entropy.py @@ -1,8 +1,10 @@ import logging import math +from collections import defaultdict 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__) @@ -32,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 @@ -59,6 +42,25 @@ def execute(self): """ start, end, step = self._get_trajectory_bounds() number_frames = self._get_number_frames(start, end, step) + + has_water = self._universe.select_atoms("water").n_atoms > 0 + if has_water and self._args.water_entropy: + self._calculate_water_entropy(self._universe, start, end, step) + + if self._args.selection_string != "all": + self._args.selection_string += " and not water" + 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) @@ -95,6 +97,18 @@ def execute(self): number_frames, highest_level, ) + + logger.debug( + "%s level: molecule_data: %s", + level, + self._data_logger.molecule_data, + ) + logger.debug( + "%s level: residue_data: %s", + level, + self._data_logger.residue_data, + ) + elif level in ("polymer", "residue"): self._process_vibrational_only_levels( molecule_id, @@ -107,6 +121,18 @@ def execute(self): number_frames, highest_level, ) + + logger.debug( + "%s level: molecule_data: %s", + level, + self._data_logger.molecule_data, + ) + logger.debug( + "%s level: residue_data: %s", + level, + self._data_logger.residue_data, + ) + if level == "residue": self._process_conformational_residue_level( molecule_id, @@ -119,7 +145,18 @@ def execute(self): number_frames, ) - self._finalize_molecule_results(molecule_id, level) + logger.debug( + "%s level: molecule_data: %s", + level, + self._data_logger.molecule_data, + ) + logger.debug( + "%s level: residue_data: %s", + level, + self._data_logger.residue_data, + ) + + self._finalize_molecule_results() self._data_logger.log_tables() @@ -240,13 +277,25 @@ 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._data_logger.add_residue_data( + mol_id, residue.resname, level, "Transvibrational", S_trans_res + ) + self._data_logger.add_residue_data( + mol_id, residue.resname, level, "Rovibrational", S_rot_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 @@ -274,9 +323,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 @@ -298,67 +351,150 @@ 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, 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). + Aggregates and logs total entropy per molecule using residue_data grouped by + resid. """ - 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) + entropy_by_molecule = defaultdict(float) + + 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}") + + 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) + ) + + # Save to file self._data_logger.save_dataframes_as_json( - self._results_df, self._residue_results_df, self._args.output_file + 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, ) - def _log_result(self, mol_id, level, entropy_type, value): + def _calculate_water_entropy(self, universe, start, end, step): """ - Logs and stores a single entropy value in the global results dataframe. + Calculates orientational and vibrational entropy for water molecules. 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], - } + 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) ) - 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): + # 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) + + # Aggregate entropy components per molecule + results = {} + + 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] = { + "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._data_logger.add_results_data( + mol_id, "water", entropy_type, S_component + ) + total += S_component + + 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._data_logger.add_residue_data( + resid, resname, "Water", "Orientational", Sor + ) + + def _calculate_water_vibrational_translational_entropy(self, vibrations): """ - Logs and stores per-residue entropy data. + Logs summed translational entropy values per residue-solvent pair. + """ + for (solute_id, _), entropy in vibrations.translational_S.items(): + if isinstance(entropy, (list, np.ndarray)): + entropy = float(np.sum(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 - ) - self._data_logger.add_residue_data(mol_id, residue_id, entropy_type, value) + 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 (solute_id, _), entropy in vibrations.rotational_S.items(): + if isinstance(entropy, (list, np.ndarray)): + entropy = float(np.sum(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/config.yaml b/config.yaml index 30c0133..2a7ba86 100644 --- a/config.yaml +++ b/config.yaml @@ -2,14 +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: - water_entropy: + disable_water_entropy: diff --git a/pyproject.toml b/pyproject.toml index e714acd..63e307f 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.3" ] [project.urls] 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 7522cfe..47a66b2 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,165 @@ 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 add_residue_data. + """ + Sorient_dict = {1: {"mol1": [1.0, 2]}, 2: {"mol1": [3.0, 4]}} + + self.entropy_manager._calculate_water_orientational_entropy(Sorient_dict) + + self.entropy_manager._data_logger.add_residue_data.assert_has_calls( + [ + 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 add_residue_data. Also verifies that the + molecule-level average is computed and logged using _log_result. + """ + mock_vibrations = MagicMock() + mock_vibrations.translational_S = { + ("res1", "mol1"): [1.0, 2.0], + ("resB_invalid", "mol1"): 4.0, + ("res2", "mol1"): 3.0, + } + + self.entropy_manager._calculate_water_vibrational_translational_entropy( + mock_vibrations + ) + + self.entropy_manager._data_logger.add_residue_data.assert_has_calls( + [ + call(-1, "res1", "Water", "Transvibrational", 3.0), + call(-1, "resB", "Water", "Transvibrational", 4.0), + call(-1, "res2", "Water", "Transvibrational", 3.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() + + 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" + ) + 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. + """ + 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._data_logger.add_residue_data.assert_has_calls( + [ + call(1, "mol1", "Water", "Orientational", 1.0), + call(-1, "res1", "Water", "Transvibrational", 2.0), + call(-1, "res1", "Water", "Rovibrational", 3.0), + ] + ) + + @patch( + "waterEntropy.recipes.interfacial_solvent.get_interfacial_water_orient_entropy" + ) + def test_calculate_water_entropy_minimal(self, mock_get_entropy): + """ + Verifies that _calculate_water_entropy correctly logs entropy components + and total for a single molecule with minimal data. + """ + mock_get_entropy.return_value = ( + {}, + None, + MagicMock( + translational_S={("ACE_1", "WAT"): 10.0}, + rotational_S={("ACE_1", "WAT"): 2.0}, + ), + None, + ) + + # 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._data_logger.add_results_data.assert_has_calls( + [ + call("ACE", "water", "Orientational", 5.0), + 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