Skip to content

Commit

Permalink
adding initial tests that check correct workflow for angle tool
Browse files Browse the repository at this point in the history
  • Loading branch information
@Jgmedina95
Jgmedina95 committed Jan 24, 2025
1 parent 59071e3 commit 15b7ccc
Show file tree
Hide file tree
Showing 2 changed files with 58 additions and 186 deletions.
8 changes: 5 additions & 3 deletions mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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=(
Expand All @@ -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. """

Expand All @@ -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,
Expand Down Expand Up @@ -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",
]:
Expand Down
236 changes: 53 additions & 183 deletions tests/test_analysis/test_bond_angles_dihedrals.py
View file
Original file line number Diff line number Diff line change
@@ -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

0 comments on commit 15b7ccc

Please sign in to comment.