diff --git a/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py b/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py index 3b9faee7..02fefcf6 100644 --- a/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py +++ b/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py @@ -26,6 +26,8 @@ class ComputingAnglesSchema(BaseModel): "all (makes all of the previous analysis)" ), ) + # This arg is here, but is not used in the code. As of now it will get the analysis + # of all the residues in the simulation selection: Optional[str] = Field( "backbone and sidechain", description=( @@ -36,7 +38,7 @@ class ComputingAnglesSchema(BaseModel): class ComputeAngles(BaseTool): - name = "compute_angles" + name = "ComputeAngles" description = """Analyze dihedral angles from a trajectory file. The tool allows for analysis of the phi-psi angles, chi1-chi2 angles, or both. """ @@ -53,7 +55,7 @@ def _run(self, input): input = self.validate_input(**input) except ValueError as e: - return f"Failed. Error using the PCA Tool: {str(e)}" + return f"Failed. Error using the ComputeAngle Tool: {str(e)}" ( traj_id, @@ -323,7 +325,7 @@ def validate_input(self, **input): for key in keys: if key not in [ "trajectory_fileid", - "pc_percentage", + "topology_fileid", "analysis", "selection", ]: diff --git a/tests/test_analysis/test_bond_angles_dihedrals.py b/tests/test_analysis/test_bond_angles_dihedrals.py index 9d343c27..0e04c81b 100644 --- a/tests/test_analysis/test_bond_angles_dihedrals.py +++ b/tests/test_analysis/test_bond_angles_dihedrals.py @@ -1,190 +1,60 @@ -from unittest.mock import MagicMock, patch - -import numpy as np -import pytest +from unittest.mock import patch from mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool import ( ComputeAngles, - ComputeChi1, - ComputeChi2, - ComputeChi3, - ComputeChi4, - ComputeDihedrals, - ComputeOmega, - ComputePhi, - ComputePsi, - RamachandranPlot, ) -# Fixture to patch 'load_single_traj' -@pytest.fixture -def patched_load_single_traj(): - with patch( - "mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool.load_single_traj" - ) as mock_load_single_traj: - yield mock_load_single_traj - - -@pytest.fixture -def compute_angles_tool(get_registry): - path_registry = get_registry("raw", True) - return ComputeAngles(path_registry) - - -@pytest.fixture -def compute_dihedrals_tool(get_registry): - path_registry = get_registry("raw", True) - return ComputeDihedrals(path_registry) - - -@pytest.fixture -def compute_phi_tool(get_registry): - path_registry = get_registry("raw", True) - return ComputePhi(path_registry) - - -@pytest.fixture -def compute_psi_tool(get_registry): - path_registry = get_registry("raw", True) - return ComputePsi(path_registry) - - -@pytest.fixture -def compute_chi1_tool(get_registry): - path_registry = get_registry("raw", True) - return ComputeChi1(path_registry) - - -@pytest.fixture -def compute_chi2_tool(get_registry): - path_registry = get_registry("raw", True) - return ComputeChi2(path_registry) - - -@pytest.fixture -def compute_chi3_tool(get_registry): - path_registry = get_registry("raw", True) - return ComputeChi3(path_registry) - - -@pytest.fixture -def compute_chi4_tool(get_registry): - path_registry = get_registry("raw", True) - return ComputeChi4(path_registry) - - -@pytest.fixture -def compute_omega_tool(get_registry): - path_registry = get_registry("raw", True) - return ComputeOmega(path_registry) - +def test_compute_angles_tool_bad_inputs(get_registry): + reg = get_registry("raw", True, map_path=True, include_peptide_trajectory=True) + angles_tool = ComputeAngles(path_registry=reg) + bad_input_files = { + "trajectory_fileid": "pep_traj_987654_3", + "topology_fileid": "pep_traj_987654_3", + "analysis": "both", + } + + error_catching = angles_tool._run(bad_input_files) + assert "Trajectory File ID not in path registry" in error_catching + assert "Topology File ID not in path registry" in error_catching + + +# patch and or moch save_results_to_file +# @patch("mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool.save_results_to_file") +def test_compute_angles_ram_values(get_registry): + reg = get_registry("raw", True, dynamic=True, include_hydrogens=True) + angles_tool = ComputeAngles(path_registry=reg) + phi_psi_input_files = { + "trajectory_fileid": "pep_traj_987654", + "topology_fileid": "pep_traj_987654", + "analysis": "phi-psi", + } + chi_innput_files = { + "trajectory_fileid": "pep_traj_987654", + "topology_fileid": "pep_traj_987654", + "analysis": "chi1-chi2", + } + # traj = md.load(reg.get_mapped_path("pep_traj_987654")) -@pytest.fixture -def ramachandran_plot(get_registry): - path_registry = get_registry("raw", True) - return RamachandranPlot(path_registry) - - -@patch("mdtraj.compute_angles") -@patch("matplotlib.pyplot.savefig") -def test_run_success_compute_angles( - mock_savefig, mock_compute_angles, patched_load_single_traj, compute_angles_tool -): - # Create a mock trajectory - mock_traj = MagicMock() - patched_load_single_traj.return_value = mock_traj - - # Define the expected output from compute_angles - expected_angles = np.array([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]) - mock_compute_angles.return_value = expected_angles - - # Mock the path registry get_mapped_path method - compute_angles_tool.path_registry.get_mapped_path = MagicMock( - return_value="angles_plot.png" - ) - - # Call the _run method - traj_file = "rec0_butane_123456" - top_file = "top_sim0_butane_123456" - angle_indices = [(0, 1, 2), (1, 2, 3)] - result = compute_angles_tool._run(traj_file, angle_indices, top_file) - - # Assertions - patched_load_single_traj.assert_called_once_with( - compute_angles_tool.path_registry, traj_file, top_file - ) - mock_compute_angles.assert_called_once_with( - mock_traj, angle_indices, periodic=True, opt=True - ) - compute_angles_tool.path_registry.get_mapped_path.assert_called_once_with( - "angles_plot.png" - ) - mock_savefig.assert_called_once_with("angles_plot.png") - assert result == "Succeeded. Bond angles computed, saved to file and plot saved." - - -def test_run_fail_compute_angles(patched_load_single_traj, compute_angles_tool): - # Simulate the trajectory loading failure - patched_load_single_traj.return_value = None - - # Call the _run method - traj_file = "rec0_butane_123456" - top_file = "top_sim0_butane_123456" - angle_indices = [(0, 1, 2), (1, 2, 3)] - result = compute_angles_tool._run(traj_file, angle_indices, top_file) - - # Assertions - patched_load_single_traj.assert_called_once_with( - compute_angles_tool.path_registry, traj_file, top_file - ) - assert result == "Failed. Trajectory could not be loaded." - - -# Similar tests for other classes (ComputeChi1, ComputeChi2, etc.) - - -@patch("matplotlib.pyplot.savefig") -@patch("mdtraj.compute_phi") -@patch("mdtraj.compute_psi") -def test_run_success_ramachandran_plot( - mock_compute_psi, - mock_compute_phi, - mock_savefig, - patched_load_single_traj, - ramachandran_plot, -): - # Create a mock trajectory - mock_traj = MagicMock() - patched_load_single_traj.return_value = mock_traj - - # Define the expected output from compute_phi and compute_psi - expected_phi = ([(0, 1, 2, 3)], [[0.7, 0.8, 0.9]]) - expected_psi = ([(0, 1, 2, 3)], [[1.0, 1.1, 1.2]]) - mock_compute_phi.return_value = expected_phi - mock_compute_psi.return_value = expected_psi - - # Mock the path registry get_mapped_path method - ramachandran_plot.path_registry.get_mapped_path = MagicMock( - return_value="ramachandran_plot.png" - ) - - # Call the _run method - traj_file = "rec0_butane_123456" - top_file = "top_sim0_butane_123456" - result = ramachandran_plot._run(traj_file, top_file) - - # Assertions - patched_load_single_traj.assert_called_once_with( - ramachandran_plot.path_registry, traj_file, top_file - ) - mock_compute_phi.assert_called_once_with(mock_traj, periodic=True, opt=True) - mock_compute_psi.assert_called_once_with(mock_traj, periodic=True, opt=True) - ramachandran_plot.path_registry.get_mapped_path.assert_called_once_with( - "ramachandran_plot.png" - ) - # Ensure savefig is called - print(mock_savefig.call_args_list) - mock_savefig.assert_called_once_with("ramachandran_plot.png") - - assert result == "Succeeded. Ramachandran plot generated and saved to file." + with patch( + "mdagent.tools.base_tools.analysis_tools.ComputeAngles.compute_and_plot_phi_psi" + ) as mock_compute_and_plot_phi_psi: + with patch( + "mdagent.tools.base_tools.analysis_tools.ComputeAngles.compute_and_plot_chi1_chi2" + ) as mock_compute_and_plot_chi1_chi2: + mock_compute_and_plot_phi_psi.return_value = ("mockid", "mockresult") + # instance.return_value = ("mockid", "mockresult") + angles_tool._run(phi_psi_input_files) + # print(result) + assert mock_compute_and_plot_phi_psi.called + # assert compute_and_plot_chi1_chi2 is not called + assert not mock_compute_and_plot_chi1_chi2.called + + # =========================================================================# + mock_compute_and_plot_chi1_chi2.return_value = ("mockid", "mockresult") + angles_tool._run(chi_innput_files) + assert mock_compute_and_plot_chi1_chi2.called + # assert compute_and_plot_phi_psi is not called + assert ( + mock_compute_and_plot_phi_psi.assert_called_once + ) # already called once