Bond angles dihedrals

mdagent/tools/base_tools/__init__.py

@@ -1,3 +1,14 @@
+from .analysis_tools.bond_angles_dihedrals_tool import (
+    ComputeAngles,
+    ComputeChi1,
+    ComputeChi2,
+    ComputeChi3,
+    ComputeChi4,
+    ComputeDihedrals,
+    ComputeOmega,
+    ComputePhi,
+    ComputePsi,
+)
 from .analysis_tools.plot_tools import SimulationOutputFigures
 from .analysis_tools.ppi_tools import PPIDistance
 from .analysis_tools.rdf_tool import RDFTool
@@ -23,12 +34,21 @@
 
 __all__ = [
     "CleaningToolFunction",
+    "ComputeAngles",
+    "ComputeChi1",
+    "ComputeChi2",
+    "ComputeChi3",
+    "ComputeChi4",
+    "ComputeDihedrals",
+    "ComputeOmega",
+    "ComputePhi",
+    "ComputePsi",
     "ListRegistryPaths",
     "MapPath2Name",
     "ModifyBaseSimulationScriptTool",
-    "ProteinName2PDBTool",
     "PackMolTool",
     "PPIDistance",
+    "ProteinName2PDBTool",
     "RadiusofGyrationAverage",
     "RadiusofGyrationPerFrame",
     "RadiusofGyrationPlot",

mdagent/tools/base_tools/analysis_tools/__init__.py

@@ -1,10 +1,30 @@
+from .bond_angles_dihedrals_tool import (
+    ComputeAngles,
+    ComputeChi1,
+    ComputeChi2,
+    ComputeChi3,
+    ComputeChi4,
+    ComputeDihedrals,
+    ComputeOmega,
+    ComputePhi,
+    ComputePsi,
+)
 from .plot_tools import SimulationOutputFigures
 from .ppi_tools import PPIDistance
 from .rgy import RadiusofGyrationAverage, RadiusofGyrationPerFrame, RadiusofGyrationPlot
 from .rmsd_tools import RMSDCalculator
 from .vis_tools import VisFunctions, VisualizeProtein
 
 __all__ = [
+    "ComputeAngles",
+    "ComputeChi1",
+    "ComputeChi2",
+    "ComputeChi3",
+    "ComputeChi4",
+    "ComputeDihedrals",
+    "ComputeOmega",
+    "ComputePhi",
+    "ComputePsi",
     "PPIDistance",
     "RadiusofGyrationAverage",
     "RadiusofGyrationPerFrame",

mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py

@@ -0,0 +1,270 @@
+import mdtraj as md
+from langchain.tools import BaseTool
+
+from mdagent.utils import PathRegistry, load_single_traj
+
+
+class ComputeAngles(BaseTool):
+    name = "compute_angles"
+    description = """Calculate the bond angles for the given sets of three atoms in
+    each snapshot, and provide a list of indices specifying which atoms are involved
+    in each bond angle calculation.."""
+
+    path_registry: PathRegistry | None = None
+
+    def __init__(self, path_registry: PathRegistry):
+        super().__init__()
+        self.path_registry = path_registry
+
+    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."
+
+            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."
+                )
+
+            return md.compute_angles(traj, angle_indices, periodic=True, opt=True)
+
+        except Exception as e:
+            return f"Failed. {type(e).__name__}: {e}"
+
+    async def _arun(self, traj_file, angle_indices, top_file=None):
+        raise NotImplementedError("Async version not implemented")
+
+
+class ComputeDihedrals(BaseTool):
+    name = "compute_dihedrals"
+    description = """Calculate the dihedral angles for the given groups of four atoms
+      in each snapshot, and provide a list of dihedral angles along with a list of
+      indices specifying which atoms are involved in each dihedral angle calculation."""
+
+    path_registry: PathRegistry | None = None
+
+    def __init__(self, path_registry: PathRegistry):
+        super().__init__()
+        self.path_registry = path_registry
+
+    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."
+
+            if (
+                not indices
+                or not isinstance(indices, list)
+                or not all(
+                    isinstance(tup, tuple) and all(isinstance(i, int) for i in tup)
+                    for tup in indices
+                )
+            ):
+                return (
+                    "Invalid indices. It should be a list of tuples, each containing"
+                    "atom indices as integers."
+                )
+
+            return md.compute_dihedrals(traj, indices, periodic=True, opt=True)
+
+        except Exception as e:
+            return f"Failed. {type(e).__name__}: {e}"
+
+    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"""
+
+    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 "Trajectory could not be loaded."
+            return md.compute_psi(traj, periodic=True, opt=True)
+
+        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")
+
+
+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
+
+    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 "Trajectory could not be loaded."
+            return md.compute_phi(traj, periodic=True, opt=True)
+
+        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")
+
+
+class ComputeChi1(BaseTool):
+    name = "compute_chi1"
+    description = """Calculate the chi1 angles (the first side chain torsion angle
+    formed between four atoms around the CA-CB axis) for each snapshot, providing a
+    list of chi1 angles and indices specifying the atoms involved in each chi1 angle
+    calculation."""
+
+    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 "Trajectory could not be loaded."
+            return md.compute_chi1(traj, periodic=True, opt=True)
+
+        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")
+
+
+class ComputeChi2(BaseTool):
+    name = "compute_chi2"
+    description = """Calculate the chi2 angles (the second side chain torsion angle
+    formed between four atoms around the CB-CG axis) for each snapshot, providing a
+    list of chi2 angles and a list of indices specifying the atoms involved in
+    calculating each chi2 angle."""
+
+    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 "Trajectory could not be loaded."
+            return md.compute_chi2(traj, periodic=True, opt=True)
+
+        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")
+
+
+class ComputeChi3(BaseTool):
+    name = "compute_chi3"
+    description = """Calculate the chi3 angles (the third side chain torsion angle
+    formed between four atoms around the CG-CD axis) for each snapshot in the
+    trajectory, providing a list of chi3 angles and indices specifying the atoms
+    involved in the calculation of each chi3 angle. Note: Only the residues ARG, GLN,
+      GLU, LYS, and MET have these atoms."""
+
+    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 "Trajectory could not be loaded."
+            return md.compute_chi3(traj, periodic=True, opt=True)
+
+        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")
+
+
+class ComputeChi4(BaseTool):
+    name = "compute_chi4"
+    description = """Calculate the chi4 angles (the fourth side chain torsion angle
+    formed between four atoms around the CD-CE or CD-NE axis) for each snapshot in the
+    trajectory, providing a list of indices specifying which atoms are involved in the
+    chi4 angle calculations. """
+
+    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 "Trajectory could not be loaded."
+            return md.compute_chi4(traj, periodic=True, opt=True)
+
+        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")
+
+
+class ComputeOmega(BaseTool):
+    name = "compute_omega"
+    description = """Calculate the omega angles for each snapshot in the trajectory,
+    providing a list of indices specifying which atoms are involved in the omega angle
+    calculations.."""
+
+    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 "Trajectory could not be loaded."
+            return md.compute_omega(traj, periodic=True, opt=True)
+
+        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")
+
+
+# class Ramachandran Plot

mdagent/tools/maketools.py

@@ -9,6 +9,15 @@
 
 from .base_tools import (
     CleaningToolFunction,
+    ComputeAngles,
+    ComputeChi1,
+    ComputeChi2,
+    ComputeChi3,
+    ComputeChi4,
+    ComputeDihedrals,
+    ComputeOmega,
+    ComputePhi,
+    ComputePsi,
     ListRegistryPaths,
     ModifyBaseSimulationScriptTool,
     PackMolTool,
@@ -46,6 +55,15 @@ def make_all_tools(
 
     # add base tools
     base_tools = [
+        ComputeAngles(path_registry=path_instance),
+        ComputeChi1(path_registry=path_instance),
+        ComputeChi2(path_registry=path_instance),
+        ComputeChi3(path_registry=path_instance),
+        ComputeChi4(path_registry=path_instance),
+        ComputeDihedrals(path_registry=path_instance),
+        ComputeOmega(path_registry=path_instance),
+        ComputePhi(path_registry=path_instance),
+        ComputePsi(path_registry=path_instance),
         CleaningToolFunction(path_registry=path_instance),
         ListRegistryPaths(path_registry=path_instance),
         ProteinName2PDBTool(path_registry=path_instance),

mdagent/utils/__init__.py

@@ -1,5 +1,15 @@
+from .data_handling import load_single_traj, load_traj_with_ref, save_plot, save_to_csv
 from .makellm import _make_llm
 from .path_registry import FileType, PathRegistry
 from .set_ckpt import SetCheckpoint
 
-__all__ = ["_make_llm", "PathRegistry", "FileType", "SetCheckpoint"]
+__all__ = [
+    "load_single_traj",
+    "load_traj_with_ref",
+    "save_plot",
+    "save_to_csv",
+    "_make_llm",
+    "FileType",
+    "PathRegistry",
+    "SetCheckpoint",
+]