Bond angles dihedrals

mdagent/tools/base_tools/__init__.py

@@ -8,6 +8,7 @@
     ComputeOmega,
     ComputePhi,
     ComputePsi,
+    RamachandranPlot,
 )
 from .analysis_tools.distance_tools import ContactsTool, DistanceMatrixTool
 from .analysis_tools.inertia import MomentOfInertia
@@ -66,6 +67,7 @@
     "RadiusofGyrationAverage",
     "RadiusofGyrationPerFrame",
     "RadiusofGyrationPlot",
+    "RamachandranPlot",
     "RDFTool",
     "RMSDCalculator",
     "Scholar2ResultLLM",

mdagent/tools/base_tools/analysis_tools/__init__.py

@@ -8,6 +8,7 @@
     ComputeOmega,
     ComputePhi,
     ComputePsi,
+    RamachandranPlot,
 )
 from .distance_tools import ContactsTool, DistanceMatrixTool
 from .inertia import MomentOfInertia
@@ -40,6 +41,7 @@
     "RadiusofGyrationAverage",
     "RadiusofGyrationPerFrame",
     "RadiusofGyrationPlot",
+    "RamachandranPlot",
     "RMSDCalculator",
     "SimulationOutputFigures",
     "SolventAccessibleSurfaceArea",

mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py

@@ -1,3 +1,4 @@
+import matplotlib.pyplot as plt
 import mdtraj as md
 from langchain.tools import BaseTool
 
@@ -20,19 +21,20 @@ def _run(self, traj_file, angle_indices, top_file=None):
         try:
             traj = load_single_traj(self.path_registry, traj_file, top_file)
             if not traj:
-                return "Trajectory could not be loaded."
+                return "Failed.Trajectory could not be loaded."
 
             if (
                 not angle_indices
                 or not isinstance(angle_indices, list)
                 or not all(len(indices) == 3 for indices in angle_indices)
             ):
                 return (
-                    "Invalid angle_indices. It should be a list of tuples, each "
-                    "containing three atom indices."
+                    "Failed. Invalid angle_indices. It should be a list of tuples, "
+                    "each containing three atom indices."
                 )
 
-            return md.compute_angles(traj, angle_indices, periodic=True, opt=True)
+            result = md.compute_angles(traj, angle_indices, periodic=True, opt=True)
+            return f"Succeeded. {result}"
 
         except Exception as e:
             return f"Failed. {type(e).__name__}: {e}"
@@ -57,7 +59,7 @@ def _run(self, traj_file, indices, top_file=None):
         try:
             traj = load_single_traj(self.path_registry, traj_file, top_file)
             if not traj:
-                return "Trajectory could not be loaded."
+                return " Failed. Trajectory could not be loaded."
 
             if (
                 not indices
@@ -67,12 +69,10 @@ def _run(self, traj_file, indices, top_file=None):
                     for tup in indices
                 )
             ):
-                return (
-                    "Invalid indices. It should be a list of tuples, each containing"
-                    "atom indices as integers."
-                )
+                return "Failed. Invalid indices. It should be a list of tuples."
 
-            return md.compute_dihedrals(traj, indices, periodic=True, opt=True)
+            result = md.compute_dihedrals(traj, indices, periodic=True, opt=True)
+            return f"Succeeded. {result}"
 
         except Exception as e:
             return f"Failed. {type(e).__name__}: {e}"
@@ -81,11 +81,10 @@ async def _arun(self, traj_file, indices, top_file=None):
         raise NotImplementedError("Async version not implemented")
 
 
-class ComputePsi(BaseTool):
-    name = "compute_psi"
-    description = """Calculate the psi angles for each snapshot, providing a list of
-    psi angles for each frame in the trajectory and a list of indices specifying the
-    atoms involved in calculating each psi angle"""
+class ComputePhi(BaseTool):
+    name = "compute_phi"
+    description = """This class calculates phi torsion angles and provides a list of phi
+      angles and indices specifying which atoms are involved in the calculations"""
 
     path_registry: PathRegistry | None = None
 
@@ -97,8 +96,10 @@ def _run(self, traj_file, top_file=None):
         try:
             traj = load_single_traj(self.path_registry, traj_file, top_file)
             if not traj:
-                return "Trajectory could not be loaded."
-            return md.compute_psi(traj, periodic=True, opt=True)
+                return " Failed. Trajectory could not be loaded."
+
+            result = md.compute_phi(traj, periodic=True, opt=True)
+            return f"Succeeded. {result}"
 
         except Exception as e:
             return f"Failed. {type(e).__name__}: {e}"
@@ -107,10 +108,11 @@ async def _arun(self, traj_file, top_file=None):
         raise NotImplementedError("Async version not implemented")
 
 
-class ComputePhi(BaseTool):
-    name = "compute_phi"
-    description = """This class calculates phi torsion angles and provides a list of phi
-      angles and indices specifying which atoms are involved in the calculations"""
+class ComputePsi(BaseTool):
+    name = "compute_psi"
+    description = """Calculate the psi angles for each snapshot, providing a list of
+    psi angles for each frame in the trajectory and a list of indices specifying the
+    atoms involved in calculating each psi angle"""
 
     path_registry: PathRegistry | None = None
 
@@ -122,8 +124,10 @@ def _run(self, traj_file, top_file=None):
         try:
             traj = load_single_traj(self.path_registry, traj_file, top_file)
             if not traj:
-                return "Trajectory could not be loaded."
-            return md.compute_phi(traj, periodic=True, opt=True)
+                return "Failed. Trajectory could not be loaded."
+
+            result = md.compute_psi(traj, periodic=True, opt=True)
+            return f"Succeeded. {result}"
 
         except Exception as e:
             return f"Failed. {type(e).__name__}: {e}"
@@ -149,8 +153,10 @@ def _run(self, traj_file, top_file=None):
         try:
             traj = load_single_traj(self.path_registry, traj_file, top_file)
             if not traj:
-                return "Trajectory could not be loaded."
-            return md.compute_chi1(traj, periodic=True, opt=True)
+                return "Failed. Trajectory could not be loaded."
+
+            result = md.compute_chi1(traj, periodic=True, opt=True)
+            return f"Succeeded. {result}"
 
         except Exception as e:
             return f"Failed. {type(e).__name__}: {e}"
@@ -176,8 +182,10 @@ def _run(self, traj_file, top_file=None):
         try:
             traj = load_single_traj(self.path_registry, traj_file, top_file)
             if not traj:
-                return "Trajectory could not be loaded."
-            return md.compute_chi2(traj, periodic=True, opt=True)
+                return "Failed. Trajectory could not be loaded."
+
+            result = md.compute_chi2(traj, periodic=True, opt=True)
+            return f"Succeeded. {result}"
 
         except Exception as e:
             return f"Failed. {type(e).__name__}: {e}"
@@ -204,8 +212,10 @@ def _run(self, traj_file, top_file=None):
         try:
             traj = load_single_traj(self.path_registry, traj_file, top_file)
             if not traj:
-                return "Trajectory could not be loaded."
-            return md.compute_chi3(traj, periodic=True, opt=True)
+                return "Failed. Trajectory could not be loaded."
+
+            result = md.compute_chi3(traj, periodic=True, opt=True)
+            return f"Succeeded. {result}"
 
         except Exception as e:
             return f"Failed. {type(e).__name__}: {e}"
@@ -231,8 +241,10 @@ def _run(self, traj_file, top_file=None):
         try:
             traj = load_single_traj(self.path_registry, traj_file, top_file)
             if not traj:
-                return "Trajectory could not be loaded."
-            return md.compute_chi4(traj, periodic=True, opt=True)
+                return "Failed. Trajectory could not be loaded."
+
+            result = md.compute_chi4(traj, periodic=True, opt=True)
+            return f"Succeeded. {result}"
 
         except Exception as e:
             return f"Failed. {type(e).__name__}: {e}"
@@ -257,8 +269,10 @@ def _run(self, traj_file, top_file=None):
         try:
             traj = load_single_traj(self.path_registry, traj_file, top_file)
             if not traj:
-                return "Trajectory could not be loaded."
-            return md.compute_omega(traj, periodic=True, opt=True)
+                return "Failed. Trajectory could not be loaded."
+
+            result = md.compute_omega(traj, periodic=True, opt=True)
+            return f"Succeeded. {result}"
 
         except Exception as e:
             return f"Failed. {type(e).__name__}: {e}"
@@ -267,4 +281,37 @@ async def _arun(self, traj_file, top_file=None):
         raise NotImplementedError("Async version not implemented")
 
 
-# class Ramachandran Plot
+class RamachandranPlot(BaseTool):
+    name = "ramachandran_plot"
+    description = """Generate a Ramachandran plot for the given trajectory, showing
+    the distribution of phi and psi angles for each frame."""
+
+    path_registry: PathRegistry | None = None
+
+    def __init__(self, path_registry: PathRegistry):
+        super().__init__()
+        self.path_registry = path_registry
+
+    def _run(self, traj_file, top_file=None):
+        try:
+            traj = load_single_traj(self.path_registry, traj_file, top_file)
+            if not traj:
+                return "Failed. Trajectory could not be loaded."
+
+            phi_indices, phi_angles = md.compute_phi(traj, periodic=True, opt=True)
+            psi_indices, psi_angles = md.compute_psi(traj, periodic=True, opt=True)
+
+            plt.figure(figsize=(10, 8))
+            plt.scatter(phi_angles.flatten(), psi_angles.flatten(), s=1, color="blue")
+            plt.xlabel("Phi Angles (radians)")
+            plt.ylabel("Psi Angles (radians)")
+            plt.title("Ramachandran Plot")
+            plt.grid(True)
+            plt.show()
+            return "Succeeded. Ramachandran plot generated."
+
+        except Exception as e:
+            return f"Failed. {type(e).__name__}: {e}"
+
+    async def _arun(self, traj_file, top_file=None):
+        raise NotImplementedError("Async version not implemented")

mdagent/tools/maketools.py

@@ -33,6 +33,7 @@
     RadiusofGyrationAverage,
     RadiusofGyrationPerFrame,
     RadiusofGyrationPlot,
+    RamachandranPlot,
     RDFTool,
     Scholar2ResultLLM,
     SetUpandRunFunction,
@@ -86,6 +87,7 @@ def make_all_tools(
         RadiusofGyrationAverage(path_registry=path_instance),
         RadiusofGyrationPerFrame(path_registry=path_instance),
         RadiusofGyrationPlot(path_registry=path_instance),
+        RamachandranPlot(path_registry=path_instance),
         RDFTool(path_registry=path_instance),
         SetUpandRunFunction(path_registry=path_instance),
         SimulationOutputFigures(path_registry=path_instance),