code and unit test for salt bridge tool

mdagent/tools/base_tools/__init__.py

@@ -10,6 +10,7 @@
     RadiusofGyrationPlot,
 )
 from .analysis_tools.rmsd_tools import ComputeLPRMSD, ComputeRMSD, ComputeRMSF
+from .analysis_tools.salt_bridge_tool import SaltBridgeTool
 from .analysis_tools.sasa import SolventAccessibleSurfaceArea
 from .analysis_tools.vis_tools import VisFunctions, VisualizeProtein
 from .preprocess_tools.clean_tools import CleaningToolFunction
@@ -44,7 +45,7 @@
     "RadiusofGyrationPerFrame",
     "RadiusofGyrationPlot",
     "RDFTool",
-    "RMSDCalculator",
+    "SaltBridgeTool",
     "Scholar2ResultLLM",
     "SerpGitTool",
     "SetUpandRunFunction",

mdagent/tools/base_tools/analysis_tools/__init__.py

@@ -5,6 +5,7 @@
 from .ppi_tools import PPIDistance
 from .rgy import RadiusofGyrationAverage, RadiusofGyrationPerFrame, RadiusofGyrationPlot
 from .rmsd_tools import ComputeLPRMSD, ComputeRMSD, ComputeRMSF
+from .salt_bridge_tool import SaltBridgeTool
 from .sasa import SolventAccessibleSurfaceArea
 from .vis_tools import VisFunctions, VisualizeProtein
 
@@ -20,7 +21,7 @@
     "RadiusofGyrationAverage",
     "RadiusofGyrationPerFrame",
     "RadiusofGyrationPlot",
-    "RMSDCalculator",
+    "SaltBridgeTool",
     "SimulationOutputFigures",
     "SolventAccessibleSurfaceArea",
     "VisFunctions",

mdagent/tools/base_tools/analysis_tools/salt_bridge_tool.py

@@ -0,0 +1,106 @@
+from typing import Optional
+
+import mdtraj as md
+from langchain.tools import BaseTool
+from pydantic import BaseModel, Field
+
+from mdagent.utils import PathRegistry
+
+
+class SaltBridgeFunction:  # this class defines a method called find_salt_bridge
+    # using MD traj and top files and threshold distance default, residue pair list
+    # used to account for salt bridge analysis
+    def __init__(self, path_registry):
+        self.path_registry = path_registry
+        self.salt_bridges = []  # stores paired salt bridges
+        self.traj = None
+
+    def find_salt_bridges(
+        self, traj_file, top_file=None, threshold_distance=0.4, residue_pairs=None
+    ):
+        # add two files here in similar format as line 14 above
+        traj_file_path = self.path_registry.get_mapped_path(traj_file)
+        ending = traj_file_path.split(".")[-1]
+        if ending in ["dcd", "xtc", "xyz"] and top_file is not None:
+            top_file_path = self.path_registry.get_mapped_path(top_file)
+            self.traj = md.load(traj_file_path, top=top_file_path)
+        else:
+            self.traj = md.load(traj_file_path)
+
+        if residue_pairs is None:
+            residue_pairs = [
+                ("ARG", "ASP"),
+                ("ARG", "GLU"),
+                ("LYS", "ASP"),
+                ("LYS", "GLU"),
+            ]
+
+        for pair in residue_pairs:
+            donor_residues = self.traj.topology.select(f'residue_name == "{pair[0]}"')
+            acceptor_residues = self.traj.topology.select(
+                f'residue_name == "{pair[1]}"'
+            )
+
+            for donor_idx in donor_residues:
+                for acceptor_idx in acceptor_residues:
+                    distances = md.compute_distances(
+                        self.traj, [[donor_idx, acceptor_idx]]
+                    )
+                    if any(d <= threshold_distance for d in distances):
+                        self.salt_bridges.append((donor_idx, acceptor_idx))
+        return self.salt_bridges
+
+    def get_results_string(self):
+        msg = "Salt bridges found: "
+        for bridge in self.salt_bridges:
+            msg += (
+                f"Residue {self.traj.topology.atom(bridge[0]).residue.index + 1} "
+                f"({self.traj.topology.atom(bridge[0]).residue.name}) - "
+                f"Residue {self.traj.topology.atom(bridge[1]).residue.index + 1} "
+                f"({self.traj.topology.atom(bridge[1]).residue.name})"
+            )
+        return msg
+
+
+class SaltBridgeToolInput(BaseModel):
+    trajectory_fileid: str = Field(
+        None, description="Trajectory file. Either dcd, hdf5, xtc, or xyz"
+    )
+
+    topology_fileid: Optional[str] = Field(None, description="Topology file")
+
+    threshold_distance: Optional[float] = Field(
+        0.4,
+        description=(
+            "maximum distance between residues for salt bridge formation in angstrom"
+        ),
+    )
+
+    residue_pairs: Optional[dict] = Field(
+        None, description=("Identifies the amino acid residues for salt bridge")
+    )
+
+
+class SaltBridgeTool(BaseTool):
+    name = "SaltBridgeTool"
+    description = "A tool to find salt bridge in a protein trajectory"
+    args_schema = SaltBridgeToolInput
+    path_registry: Optional[PathRegistry]
+
+    def __init__(self, path_registry):
+        super().__init__()
+        self.path_registry = path_registry
+
+    def _run(
+        self, traj_file, top_file=None, threshold_distance=0.4, residue_pairs=None
+    ):
+        try:
+            # calls the salt bridge function
+            salt_bridge_function = SaltBridgeFunction(self.path_registry)
+            salt_bridge_function.find_salt_bridges(
+                traj_file, top_file, threshold_distance, residue_pairs
+            )
+            message = salt_bridge_function.get_results_string()
+        except Exception as e:
+            return f"Failed. {type(e).__name__}: {e}"
+        return "Succeeded. " + message

mdagent/tools/maketools.py

@@ -25,6 +25,7 @@
     RadiusofGyrationPerFrame,
     RadiusofGyrationPlot,
     RDFTool,
+    SaltBridgeTool,
     Scholar2ResultLLM,
     SetUpandRunFunction,
     SimulationOutputFigures,
@@ -69,6 +70,7 @@ def make_all_tools(
         RadiusofGyrationPerFrame(path_registry=path_instance),
         RadiusofGyrationPlot(path_registry=path_instance),
         RDFTool(path_registry=path_instance),
+        SaltBridgeTool(path_registry=path_instance),
         SetUpandRunFunction(path_registry=path_instance),
         SimulationOutputFigures(path_registry=path_instance),
         SmallMolPDB(path_registry=path_instance),

tests/__init__.py

@@ -0,0 +1 @@
+# This is an empty __init__.py file

tests/test_analysis/test_saltbridge_tools.py

@@ -0,0 +1,50 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from mdagent.tools.base_tools.analysis_tools.salt_bridge_tool import (
+    SaltBridgeFunction,
+    SaltBridgeTool,
+)
+
+
+@pytest.fixture
+def fake_path_registry():
+    # Mock PathRegistry to return a specific file path when asked
+    mock_registry = MagicMock()
+    mock_registry.get_mapped_path.side_effect = lambda x: f"/fake/path/{x}"
+    return mock_registry
+
+
+@pytest.fixture
+def mock_md_load():
+    with patch("mdtraj.load", autospec=True) as mock:
+        yield mock
+
+
+def test_saltbridge_tool_init(get_registry):
+    registry = get_registry("raw", False)
+    tool = SaltBridgeTool(path_registry=registry)
+    assert tool.name == "SaltBridgeTool"
+    assert tool.path_registry == registry
+
+
+def test_salt_bridge_function_init(get_registry):
+    path_registry = get_registry("raw", False)
+    sbf = SaltBridgeFunction(path_registry)
+    assert sbf.path_registry == path_registry
+
+
+def test_find_salt_bridges(fake_path_registry, mock_md_load):
+    sbf = SaltBridgeFunction(fake_path_registry)
+    sbf.find_salt_bridges("traj_file.dcd", "top_file.top")
+    mock_md_load.assert_called_once_with(
+        "/fake/path/traj_file.dcd", top="/fake/path/top_file.top"
+    )
+    # mock_md_compute_distances.assert_called()
+
+
+def test_salt_bridge_function_without_top(fake_path_registry, mock_md_load):
+    sbf = SaltBridgeFunction(fake_path_registry)
+    sbf.find_salt_bridges("traj_file.hdf5")
+    mock_md_load.assert_called_once_with("/fake/path/traj_file.hdf5")