Skip to content

Commit c247cbb

Browse files
harryswift01harryswift01
committed
add atomic unit tests for force merge logic in run_entropy_workflow:
- `test_merges_forcefile_correctly`: verifies that coordinate and force universes are merged and loaded correctly - `test_converts_kcal_to_kj`: checks that forces are scaled by 4.184 when `kcal_force_units=True` - `test_logs_debug_messages`: ensures expected debug logs are emitted during force merge
1 parent 5c03f80 commit c247cbb

File tree

1 file changed

+177
-0
lines changed

1 file changed

+177
-0
lines changed

‎tests/test_CodeEntropy/test_run.py‎

Lines changed: 177 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -513,6 +513,183 @@ def test_run_entropy_workflow_missing_selection_string(self):
513513
with self.assertRaisesRegex(ValueError, "Missing 'selection_string' argument."):
514514
run_manager.run_entropy_workflow()
515515

516+
@patch("CodeEntropy.run.EntropyManager")
517+
@patch("CodeEntropy.run.GroupMolecules")
518+
@patch("CodeEntropy.run.LevelManager")
519+
@patch("CodeEntropy.run.AnalysisFromFunction")
520+
@patch("CodeEntropy.run.mda.Merge")
521+
@patch("CodeEntropy.run.mda.Universe")
522+
def test_merges_forcefile_correctly(
523+
self, MockUniverse, MockMerge, MockAnalysis, *_mocks
524+
):
525+
"""
526+
Ensure that coordinates and forces are merged correctly
527+
when a force file is provided.
528+
"""
529+
run_manager = RunManager("mock/job001")
530+
mock_logger = MagicMock()
531+
run_manager._logging_config = MagicMock(
532+
setup_logging=MagicMock(return_value=mock_logger)
533+
)
534+
run_manager._config_manager = MagicMock()
535+
run_manager._data_logger = MagicMock()
536+
run_manager.show_splash = MagicMock()
537+
538+
args = MagicMock(
539+
top_traj_file=["topol.tpr", "traj.xtc"],
540+
force_file="forces.xtc",
541+
file_format="xtc",
542+
selection_string="all",
543+
verbose=False,
544+
output_file="output.json",
545+
kcal_force_units=False,
546+
)
547+
run_manager._config_manager.load_config.return_value = {"run1": {}}
548+
parser = run_manager._config_manager.setup_argparse.return_value
549+
parser.parse_known_args.return_value = (args, [])
550+
run_manager._config_manager.merge_configs.return_value = args
551+
run_manager._config_manager.input_parameters_validation = MagicMock()
552+
553+
mock_u, mock_u_force = MagicMock(), MagicMock()
554+
MockUniverse.side_effect = [mock_u, mock_u_force]
555+
mock_atoms, mock_atoms_force = MagicMock(), MagicMock()
556+
mock_u.select_atoms.return_value = mock_atoms
557+
mock_u_force.select_atoms.return_value = mock_atoms_force
558+
559+
coords, forces = np.random.rand(2, 3, 3), np.random.rand(2, 3, 3)
560+
MockAnalysis.side_effect = [
561+
MagicMock(
562+
run=MagicMock(return_value=MagicMock(results={"timeseries": coords}))
563+
),
564+
MagicMock(
565+
run=MagicMock(return_value=MagicMock(results={"timeseries": forces}))
566+
),
567+
]
568+
mock_merge = MagicMock()
569+
MockMerge.return_value = mock_merge
570+
571+
run_manager.run_entropy_workflow()
572+
573+
MockUniverse.assert_any_call("topol.tpr", ["traj.xtc"], format="xtc")
574+
MockUniverse.assert_any_call("topol.tpr", "forces.xtc", format="xtc")
575+
MockMerge.assert_called_once_with(mock_atoms)
576+
mock_merge.load_new.assert_called_once_with(coords, forces=forces)
577+
self.assertEqual(mock_logger.debug.call_count, 3)
578+
579+
@patch("CodeEntropy.run.EntropyManager")
580+
@patch("CodeEntropy.run.GroupMolecules")
581+
@patch("CodeEntropy.run.LevelManager")
582+
@patch("CodeEntropy.run.AnalysisFromFunction")
583+
@patch("CodeEntropy.run.mda.Merge")
584+
@patch("CodeEntropy.run.mda.Universe")
585+
def test_converts_kcal_to_kj(self, MockUniverse, MockMerge, MockAnalysis, *_mocks):
586+
"""
587+
Ensure that forces are scaled by 4.184 when kcal_force_units=True.
588+
"""
589+
run_manager = RunManager("mock/job001")
590+
mock_logger = MagicMock()
591+
run_manager._logging_config = MagicMock(
592+
setup_logging=MagicMock(return_value=mock_logger)
593+
)
594+
run_manager._config_manager = MagicMock()
595+
run_manager._data_logger = MagicMock()
596+
run_manager.show_splash = MagicMock()
597+
598+
args = MagicMock(
599+
top_traj_file=["topol.tpr", "traj.xtc"],
600+
force_file="forces.xtc",
601+
file_format="xtc",
602+
selection_string="all",
603+
verbose=False,
604+
output_file="output.json",
605+
kcal_force_units=True,
606+
)
607+
run_manager._config_manager.load_config.return_value = {"run1": {}}
608+
parser = run_manager._config_manager.setup_argparse.return_value
609+
parser.parse_known_args.return_value = (args, [])
610+
run_manager._config_manager.merge_configs.return_value = args
611+
run_manager._config_manager.input_parameters_validation = MagicMock()
612+
613+
mock_u, mock_u_force = MagicMock(), MagicMock()
614+
MockUniverse.side_effect = [mock_u, mock_u_force]
615+
mock_atoms, mock_atoms_force = MagicMock(), MagicMock()
616+
mock_u.select_atoms.return_value = mock_atoms
617+
mock_u_force.select_atoms.return_value = mock_atoms_force
618+
619+
coords = np.random.rand(2, 3, 3)
620+
forces = np.random.rand(2, 3, 3)
621+
forces_orig = forces.copy()
622+
MockAnalysis.side_effect = [
623+
MagicMock(
624+
run=MagicMock(return_value=MagicMock(results={"timeseries": coords}))
625+
),
626+
MagicMock(
627+
run=MagicMock(return_value=MagicMock(results={"timeseries": forces}))
628+
),
629+
]
630+
mock_merge = MagicMock()
631+
MockMerge.return_value = mock_merge
632+
633+
run_manager.run_entropy_workflow()
634+
_, kwargs = mock_merge.load_new.call_args
635+
np.testing.assert_allclose(kwargs["forces"], forces_orig * 4.184)
636+
637+
@patch("CodeEntropy.run.EntropyManager")
638+
@patch("CodeEntropy.run.GroupMolecules")
639+
@patch("CodeEntropy.run.LevelManager")
640+
@patch("CodeEntropy.run.AnalysisFromFunction")
641+
@patch("CodeEntropy.run.mda.Merge")
642+
@patch("CodeEntropy.run.mda.Universe")
643+
def test_logs_debug_messages(self, MockUniverse, MockMerge, MockAnalysis, *_mocks):
644+
"""
645+
Ensure that loading and merging steps produce debug logs.
646+
"""
647+
run_manager = RunManager("mock/job001")
648+
mock_logger = MagicMock()
649+
run_manager._logging_config = MagicMock(
650+
setup_logging=MagicMock(return_value=mock_logger)
651+
)
652+
run_manager._config_manager = MagicMock()
653+
run_manager._data_logger = MagicMock()
654+
run_manager.show_splash = MagicMock()
655+
656+
args = MagicMock(
657+
top_traj_file=["topol.tpr", "traj.xtc"],
658+
force_file="forces.xtc",
659+
file_format="xtc",
660+
selection_string="all",
661+
verbose=False,
662+
output_file="output.json",
663+
kcal_force_units=False,
664+
)
665+
run_manager._config_manager.load_config.return_value = {"run1": {}}
666+
parser = run_manager._config_manager.setup_argparse.return_value
667+
parser.parse_known_args.return_value = (args, [])
668+
run_manager._config_manager.merge_configs.return_value = args
669+
run_manager._config_manager.input_parameters_validation = MagicMock()
670+
671+
mock_u, mock_u_force = MagicMock(), MagicMock()
672+
MockUniverse.side_effect = [mock_u, mock_u_force]
673+
mock_atoms, mock_atoms_force = MagicMock(), MagicMock()
674+
mock_u.select_atoms.return_value = mock_atoms
675+
mock_u_force.select_atoms.return_value = mock_atoms_force
676+
677+
arr = np.random.rand(2, 3, 3)
678+
MockAnalysis.side_effect = [
679+
MagicMock(
680+
run=MagicMock(return_value=MagicMock(results={"timeseries": arr}))
681+
),
682+
MagicMock(
683+
run=MagicMock(return_value=MagicMock(results={"timeseries": arr}))
684+
),
685+
]
686+
MockMerge.return_value = MagicMock()
687+
688+
run_manager.run_entropy_workflow()
689+
debug_msgs = [c[0][0] for c in mock_logger.debug.call_args_list]
690+
self.assertTrue(any("forces.xtc" in m for m in debug_msgs))
691+
self.assertTrue(any("Merging forces" in m for m in debug_msgs))
692+
516693
@patch("CodeEntropy.run.AnalysisFromFunction")
517694
@patch("CodeEntropy.run.mda.Merge")
518695
def test_new_U_select_frame(self, MockMerge, MockAnalysisFromFunction):

0 commit comments

Comments
 (0)