From aba9dd3021204dfd57f34e1569449dea9e6d7aa8 Mon Sep 17 00:00:00 2001 From: harryswift01 Date: Mon, 2 Mar 2026 13:47:43 +0000 Subject: [PATCH 1/4] Fix water entropy crash for pure solvent systems and ensure consistent universe usage --- CodeEntropy/entropy/workflow.py | 47 ++++++++++++++++++++++++--------- 1 file changed, 35 insertions(+), 12 deletions(-) diff --git a/CodeEntropy/entropy/workflow.py b/CodeEntropy/entropy/workflow.py index d79b46e..56e2326 100644 --- a/CodeEntropy/entropy/workflow.py +++ b/CodeEntropy/entropy/workflow.py @@ -93,9 +93,15 @@ def __init__( def execute(self) -> None: """Run the full entropy workflow and emit results. - This method orchestrates the complete pipeline, populates shared data, - and triggers the DAG/graph executions. Final results are logged and saved - via `ResultsReporter`. + This orchestrates the complete entropy pipeline: + 1. Build trajectory slice. + 2. Apply atom selection to create a reduced universe. + 3. Detect hierarchy levels. + 4. Group molecules. + 5. Split groups into water and non-water. + 6. Optionally compute water entropy (only if solute exists). + 7. Run level DAG and entropy graph. + 8. Finalize and persist results. """ traj = self._build_trajectory_slice() console.print( @@ -109,9 +115,11 @@ def execute(self) -> None: reduced_universe, self._args.grouping ) - nonwater_groups, water_groups = self._split_water_groups(groups) + nonwater_groups, water_groups = self._split_water_groups( + reduced_universe, groups + ) - if self._args.water_entropy and water_groups: + if self._args.water_entropy and water_groups and nonwater_groups: self._compute_water_entropy(traj, water_groups) else: nonwater_groups.update(water_groups) @@ -254,17 +262,32 @@ def _detect_levels(self, reduced_universe: Any) -> Any: return levels def _split_water_groups( - self, groups: Mapping[int, Any] + self, + universe: Any, + groups: Mapping[int, Any], ) -> Tuple[Dict[int, Any], Dict[int, Any]]: """Partition molecule groups into water and non-water groups. + This method identifies which molecule groups correspond to water + molecules based on residue membership. + Args: - groups: Mapping of group id -> molecule ids. + universe (Any): + The MDAnalysis Universe used to build the molecule groups + (typically the reduced_universe). + groups (Mapping[int, Any]): + Mapping of group_id -> list of molecule fragment indices. Returns: - Tuple of (nonwater_groups, water_groups). + Tuple[Dict[int, Any], Dict[int, Any]]: + A tuple containing: + + - nonwater_groups: + Mapping of group_id -> molecule ids that are NOT water. + - water_groups: + Mapping of group_id -> molecule ids that contain water. """ - water_atoms = self._universe.select_atoms("water") + water_atoms = universe.select_atoms("water") water_resids = {res.resid for res in water_atoms.residues} water_groups = { @@ -272,7 +295,7 @@ def _split_water_groups( for gid, mol_ids in groups.items() if any( res.resid in water_resids - for mol in [self._universe.atoms.fragments[i] for i in mol_ids] + for mol in [universe.atoms.fragments[i] for i in mol_ids] for res in mol.residues ) } @@ -293,10 +316,10 @@ def _compute_water_entropy( if not water_groups or not self._args.water_entropy: return - water_entropy = WaterEntropy(self._args) + water_entropy = WaterEntropy(self._args, self._reporter) for group_id in water_groups.keys(): - water_entropy._calculate_water_entropy( + water_entropy.calculate_and_log( universe=self._universe, start=traj.start, end=traj.end, From a0a0932b2465c29fd1caea57186e6c0d0ea58e41 Mon Sep 17 00:00:00 2001 From: harryswift01 Date: Mon, 2 Mar 2026 15:05:56 +0000 Subject: [PATCH 2/4] fix unit tests in relation to PR #297 --- .../unit/CodeEntropy/entropy/test_workflow.py | 34 +++++++++++++------ 1 file changed, 23 insertions(+), 11 deletions(-) diff --git a/tests/unit/CodeEntropy/entropy/test_workflow.py b/tests/unit/CodeEntropy/entropy/test_workflow.py index 29b93cc..eb87e39 100644 --- a/tests/unit/CodeEntropy/entropy/test_workflow.py +++ b/tests/unit/CodeEntropy/entropy/test_workflow.py @@ -102,15 +102,15 @@ def test_execute_water_entropy_branch_calls_water_entropy_solver(): patch("CodeEntropy.entropy.workflow.EntropyGraph") as GraphCls, ): water_instance = WaterCls.return_value - water_instance._calculate_water_entropy = MagicMock() + water_instance.calculate_and_log = MagicMock() LevelDAGCls.return_value.build.return_value.execute.return_value = None GraphCls.return_value.build.return_value.execute.return_value = {} wf.execute() - water_instance._calculate_water_entropy.assert_called_once() - _, kwargs = water_instance._calculate_water_entropy.call_args + water_instance.calculate_and_log.assert_called_once() + _, kwargs = water_instance.calculate_and_log.call_args assert kwargs["universe"] is universe assert kwargs["start"] == 0 assert kwargs["end"] == 5 @@ -190,7 +190,7 @@ def test_split_water_groups_returns_empty_when_none(): universe_operations=MagicMock(), ) - groups, water = wf._split_water_groups({0: [1, 2]}) + groups, water = wf._split_water_groups(wf._universe, {0: [1, 2]}) assert water == {} @@ -253,11 +253,17 @@ def test_compute_water_entropy_updates_selection_string_and_calls_internal_metho with patch("CodeEntropy.entropy.workflow.WaterEntropy") as WaterCls: inst = WaterCls.return_value - inst._calculate_water_entropy = MagicMock() + inst.calculate_and_log = MagicMock() wf._compute_water_entropy(traj, water_groups) - inst._calculate_water_entropy.assert_called_once() + inst.calculate_and_log.assert_called_once_with( + universe=wf._universe, + start=traj.start, + end=traj.end, + step=traj.step, + group_id=9, + ) assert wf._args.selection_string == "not water" @@ -345,7 +351,7 @@ def test_split_water_groups_partitions_correctly(): ) groups = {0: [0], 1: [1]} - nonwater, water = wf._split_water_groups(groups) + nonwater, water = wf._split_water_groups(universe, groups) assert 0 in water assert 1 in nonwater @@ -366,13 +372,19 @@ def test_compute_water_entropy_instantiates_waterentropy_and_updates_selection_s with patch("CodeEntropy.entropy.workflow.WaterEntropy") as WaterCls: inst = WaterCls.return_value - inst._calculate_water_entropy = MagicMock() + inst.calculate_and_log = MagicMock() wf._compute_water_entropy(traj, water_groups) - WaterCls.assert_called_once_with(args) - inst._calculate_water_entropy.assert_called_once() - assert wf._args.selection_string == "not water" + WaterCls.assert_called_once_with(args, reporter) + inst.calculate_and_log.assert_called_once_with( + universe=universe, + start=traj.start, + end=traj.end, + step=traj.step, + group_id=9, + ) + assert args.selection_string == "not water" def test_detect_levels_calls_hierarchy_builder(): From a0b8dcd6413b9b9071f217d18f026e9cd4890ce3 Mon Sep 17 00:00:00 2001 From: harryswift01 Date: Tue, 3 Mar 2026 09:18:58 +0000 Subject: [PATCH 3/4] update regression tests to include a test for water --- tests/regression/baselines/water.json | 40 ++++++++++++++++++++++ tests/regression/configs/water/config.yaml | 12 +++++++ tests/regression/test_regression.py | 1 + 3 files changed, 53 insertions(+) create mode 100644 tests/regression/baselines/water.json create mode 100644 tests/regression/configs/water/config.yaml diff --git a/tests/regression/baselines/water.json b/tests/regression/baselines/water.json new file mode 100644 index 0000000..0199ecc --- /dev/null +++ b/tests/regression/baselines/water.json @@ -0,0 +1,40 @@ +{ + "args": { + "top_traj_file": [ + "../../test_data/Liquids_simulation_data/GAFF/water/molecules.top", + "../../test_data/Liquids_simulation_data/GAFF/water/trajectory.crd" + ], + "force_file": "../../test_data/Liquids_simulation_data/GAFF/water/forces.frc", + "file_format": "MDCRD", + "kcal_force_units": false, + "selection_string": "all", + "start": 0, + "end": 1, + "step": 1, + "bin_width": 30, + "temperature": 298.0, + "verbose": false, + "output_file": "/home/tdo96567/BioSim/temp/water/job008/output_file.json", + "force_partitioning": 0.5, + "water_entropy": true, + "grouping": "molecules", + "combined_forcetorque": true, + "customised_axes": true + }, + "provenance": { + "python": "3.14.0", + "platform": "Linux-6.6.87.2-microsoft-standard-WSL2-x86_64-with-glibc2.39", + "codeentropy_version": "2.0.0", + "git_sha": "cba3d8ea4118e00b25ee5a58d7ba951e4894b5c0" + }, + "groups": { + "0": { + "components": { + "united_atom:Transvibrational": 79.20298312418278, + "united_atom:Rovibrational": 50.90260688502127, + "united_atom:Conformational": 0.0 + }, + "total": 130.10559000920404 + } + } +} diff --git a/tests/regression/configs/water/config.yaml b/tests/regression/configs/water/config.yaml new file mode 100644 index 0000000..2197e42 --- /dev/null +++ b/tests/regression/configs/water/config.yaml @@ -0,0 +1,12 @@ +--- + +run1: + force_file: ".testdata/water/forces.frc" + top_traj_file: + - ".testdata/water/molecules.top" + - ".testdata/water/trajectory.crd" + selection_string: "all" + start: 0 + step: 1 + end: 1 + file_format: "MDCRD" diff --git a/tests/regression/test_regression.py b/tests/regression/test_regression.py index d5ca7f1..56c9126 100644 --- a/tests/regression/test_regression.py +++ b/tests/regression/test_regression.py @@ -111,6 +111,7 @@ def _compare_grouped( "methane", "methanol", pytest.param("octonol", marks=pytest.mark.slow), + "water", ], ) def test_regression_matches_baseline( From 8fd566c5ae070f7c50afddea4243647a9cdc3ddf Mon Sep 17 00:00:00 2001 From: harryswift01 Date: Wed, 4 Mar 2026 09:07:30 +0000 Subject: [PATCH 4/4] test(regression): set absolute tolerance to 0.5 kmol - Update regression comparison to use an absolute tolerance of 0.5 kmol - Replace previous numerical tolerance (1e-8) with a physically meaningful threshold - Keep relative tolerance (rtol=1e-9) unchanged for floating-point stability --- tests/regression/test_regression.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/regression/test_regression.py b/tests/regression/test_regression.py index 56c9126..f877899 100644 --- a/tests/regression/test_regression.py +++ b/tests/regression/test_regression.py @@ -151,5 +151,5 @@ def test_regression_matches_baseline( got_payload=run.payload, baseline_payload=baseline_payload, rtol=1e-9, - atol=1e-8, + atol=0.5, )