code and unit test for salt bridge tool

mdagent/tools/base_tools/analysis_tools/salt_bridge_tool.py

@@ -0,0 +1,104 @@
+import os
+from typing import Optional
+
+import matplotlib.pyplot as plt
+import mdtraj as md
+import numpy as np
+from langchain.tools import BaseTool
+
+from mdagent.utils import FileType, 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.includes_top = [".h5", ".lh5", ".pdb"] # those are the files we need for this analysis
+        self.paired_salt_bridges=[] #stores paired salt bridges
+        self.unpaired_residues=set() #store unpaired residues
+
+    def find_salt_bridges(self, traj_file, top_file , threshold_distance=0.4, residue_pairs=None):
+        salt_bridges = []
+        # load trajectory using MDTraj
+        traj = md.load(traj_file, top= top_file)
+        if residue_pairs is None:
+            residue_pairs = [("ARG", "ASP"), ("ARG", "GLU"), ("LYS", "ASP"), ("LYS", "GLU")]
+
+        for pair in residue_pairs:
+            donor_residues = traj.topology.select(f'residue_name == "{pair[0]}"')
+            acceptor_residues = traj.topology.select(f'residue_name == "{pair[1]}"')
+
+            for donor_idx in donor_residues:
+                for acceptor_idx in acceptor_residues:
+                    distances = md.compute_distances(traj, [[donor_idx, acceptor_idx]])
+                    if any(d <= threshold_distance for d in distances):
+                        salt_bridges.append((donor_idx, acceptor_idx))
+
+ # Check if the donor and acceptor form a salt bridge
+                        if any(d <= threshold_distance for d in distances):
+                            # If yes, remove them from the unpaired set
+                            self.unpaired_residues.discard(donor_idx)  # Remove donor from unpaired residues set
+                            self.unpaired_residues.discard(acceptor_idx)  # Remove acceptor from unpaired residues set
+                        else:
+                            # If not, add them to the unpaired set
+                            self.unpaired_residues.add(donor_idx)  # Add donor to unpaired residues set
+                            self.unpaired_residues.add(acceptor_idx)  # Add acceptor to unpaired residues set
+        print("Salt bridges found:")
+        for bridge in salt_bridges:
+            print(
+                f"Residue {traj.topology.atom(bridge[0]).residue.index + 1} ({traj.topology.atom(bridge[0]).residue.name}) - "
+                f"Residue {traj.topology.atom(bridge[1]).residue.index + 1} ({traj.topology.atom(bridge[1]).residue.name})"
+            )
+
+            #Print unpaired residues
+        print("Unpaired_residues:")
+        salt_bridge_function = salt_bridge_tool.salt_bridge_function
+        for residue_idx in salt_bridge_function.unpaired_residues:
+            print(f"Residue {traj.topology.atom(residue_idx).residue.index + 1} ({traj.topology.atom(residue_idx).residue.name})")
+
+        return salt_bridges, list(self.unpaired_residues), list(residue_pairs)
+
+
+
+class SaltBridgeTool(BaseTool):
+        name = "salt_bridge_tool"
+        description = "A tool to find salt bridge in a protein trajectory"
+
+     def __init__(self, path_registry):
+         self.salt_bridge_function = SaltBridgeFunction(path_registry)
+
+    def _run(self, traj_file, top_file, threshold_distance=0.4, residue_pairs=None):
+         # Load trajectory using MDTraj
+        traj = md.load(traj_file, top=top_file)
+        #calls the salt bridge function
+        salt_bridges = [self.salt_bridge_function.find_salt_bridges(traj_file, top_file,
+        threshold_distance, residue_pairs)]
+        return salt_bridges
+
+    def _agg_result(self, result):
+        return result
+
+    def _call__(self, traj_file, top_file, threshold_distance=0.4, residue_pairs=None):
+        result = self._run(traj_file, top_file, threshold_distance, residue_pairs)
+        return self._agg_result(result)
+
+
+
+
+# Print identified salt bridges
+
+# this need to be moved under class somewhere.  move it 
+
+print("Salt bridges found:")
+for bridge in salt_bridges:
+    print(
+        f"Residue {traj.topology.atom(bridge[0]).residue.index + 1} ({traj.topology.atom(bridge[0]).residue.name}) - "
+        f"Residue {traj.topology.atom(bridge[1]).residue.index + 1} ({traj.topology.atom(bridge[1]).residue.name})"
+    )
+
+    #Print unpaired residues
+print("Unpaired_residues:")
+salt_bridge_function = salt_bridge_tool.salt_bridge_function
+for residue_idx in salt_bridge_function.unpaired_residues:
+    print(f"Residue {traj.topology.atom(residue_idx).residue.index + 1} ({traj.topology.atom(residue_idx).residue.name})")

tests/__init__.py

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

tests/test_saltbridge_tools.py

@@ -0,0 +1,90 @@
+from unittest.mock import Mock
+import mdtraj as md 
+import pytest
+
+from mdagent.tools.base_tools.analysis_tools.salt_bridge_tool import SaltBridgeFunction
+from mdagent.utils import PathRegistry
+
+# define test data
+
+traj_file = "test_trajectory.dcd"
+top_file = "test_topology.pdb"
+
+#Test saltbridgefunction class
+
+def test_salt_bridge_function():
+    #initalize SaltBridgeFunction
+    salt_bridge_function = SaltBridgeFunction(path_registry=None)
+
+#load traj using MDtraj 
+    traj = md.load(traj_file, top=top_file)
+
+#perform salt bridge analysis 
+
+    salt_bridges, unpaired_residues, residue_pairs = salt_bridge_function.find.salt_bridges(traj, top=top_file)
+
+#check to see if we get a list in our results
+
+    assert isinstance(salt_bridges, list)
+
+#Check to make sure residue pairs cant be changed ( tuple)
+
+    for pair in residue_pairs:
+         assert isinstance(pair, tuple) 
+
+# check to make sure unpaired residue is assigned a number to identify it by the number (integer)
+# to ensure the code works
+
+    for residue in unpaired_residues:
+        assert isinstance( residue, int)
+
+# Finally, run the test
+    if __name__ == "__main__":
+        test_salt_bridge_function()
+
+
+
+@pytest.fixture
+def mock_traj():
+    # Create a mock trajectory object
+    traj = Mock()
+    traj.topology.select.return_value = [1, 2, 3]  # Mocking residue selections
+    return traj
+
+
+def test_find_salt_bridges(mock_traj):
+    # Instantiate SaltBridgeFunction
+    salt_bridge_function = SaltBridgeFunction(path_registry=None)
+
+    # Call find_salt_bridges method
+    traj = md.load(traj_file,top=top_file)  
+    salt_bridges = salt_bridge_function.find_salt_bridges(traj=mock_traj)
+
+    # Assert that salt_bridges contain the expected values
+    expected_salt_bridges = [
+        (1, 2),
+        (1, 3),
+        (2, 1),
+        (2, 3),
+        (3, 1),
+        (3, 2),
+    ]  # Example expected output
+    assert salt_bridges == expected_salt_bridges
+
+
+import pytest
+
+
+def test_count_salt_bridges():
+    salt_bridge_function = SaltBridgeFunction(path_registry=None)
+    salt_bridges = [
+        (1, 2),
+        (1, 3),
+        (2, 1),
+        (2, 3),
+        (3, 1),
+        (3, 2),
+    ]  # Example salt bridges
+    count = salt_bridge_function.count_salt_bridges(salt_bridges)
+    assert count == 6  # Example expected count
+