From 7fd4c3a3799df496987318018d2edad10e82b0d7 Mon Sep 17 00:00:00 2001
From: Sam Cox <swrig30@ur.rochester.edu>
Date: Fri, 23 Feb 2024 17:07:01 -0800
Subject: [PATCH] added plot and traj rgy tools

---
 mdagent/tools/base_tools/__init__.py          |  8 ++-
 .../base_tools/analysis_tools/__init__.py     |  4 +-
 .../{analysis_props.py => rgy.py}             | 54 +++++++++++++++++++
 mdagent/tools/maketools.py                    |  4 ++
 4 files changed, 68 insertions(+), 2 deletions(-)
 rename mdagent/tools/base_tools/analysis_tools/{analysis_props.py => rgy.py} (66%)

diff --git a/mdagent/tools/base_tools/__init__.py b/mdagent/tools/base_tools/__init__.py
index 651938a6..b1911505 100644
--- a/mdagent/tools/base_tools/__init__.py
+++ b/mdagent/tools/base_tools/__init__.py
@@ -1,6 +1,10 @@
-from .analysis_tools.analysis_props import RadiusofGyrationAverage
 from .analysis_tools.plot_tools import SimulationOutputFigures
 from .analysis_tools.ppi_tools import PPIDistance
+from .analysis_tools.rgy import (
+    RadiusofGyrationAverage,
+    RadiusofGyrationPerFrame,
+    RadiusofGyrationPlot,
+)
 from .analysis_tools.rmsd_tools import RMSDCalculator
 from .analysis_tools.vis_tools import (
     CheckDirectoryFiles,
@@ -46,6 +50,8 @@
     "RMSDCalculator",
     "RemoveWaterCleaningTool",
     "RadiusofGyrationAverage",
+    "RadiusofGyrationPerFrame",
+    "RadiusofGyrationPlot",
     "Scholar2ResultLLM",
     "SerpGitTool",
     "SetUpAndRunTool",
diff --git a/mdagent/tools/base_tools/analysis_tools/__init__.py b/mdagent/tools/base_tools/analysis_tools/__init__.py
index c10d6e0d..04bbbb94 100644
--- a/mdagent/tools/base_tools/analysis_tools/__init__.py
+++ b/mdagent/tools/base_tools/analysis_tools/__init__.py
@@ -1,12 +1,14 @@
-from .analysis_props import RadiusofGyrationAverage
 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 CheckDirectoryFiles, VisFunctions, VisualizeProtein
 
 __all__ = [
     "PPIDistance",
     "RMSDCalculator",
+    "RadiusofGyrationPerFrame",
+    "RadiusofGyrationPlot",
     "SimulationOutputFigures",
     "CheckDirectoryFiles",
     "VisualizeProtein",
diff --git a/mdagent/tools/base_tools/analysis_tools/analysis_props.py b/mdagent/tools/base_tools/analysis_tools/rgy.py
similarity index 66%
rename from mdagent/tools/base_tools/analysis_tools/analysis_props.py
rename to mdagent/tools/base_tools/analysis_tools/rgy.py
index 6b8a568f..32a18a20 100644
--- a/mdagent/tools/base_tools/analysis_tools/analysis_props.py
+++ b/mdagent/tools/base_tools/analysis_tools/rgy.py
@@ -110,3 +110,57 @@ def _run(self, pdb_id: str) -> str:
     async def _arun(self, query: str) -> str:
         """Use the tool asynchronously."""
         raise NotImplementedError("custom_search does not support async")
+
+
+class RadiusofGyrationPerFrame(BaseTool):
+    name = "RadiusofGyrationPerFrame"
+    description = """This tool calculates the radius of gyration
+    at each frame of a given trajectory file. Give this tool the
+    protein ID (PDB ID) only. The tool will automatically find the necessary files.
+    The tool will save the radii of gyration to a csv file and
+    map it to the registry."""
+
+    path_registry: Optional[PathRegistry]
+
+    def __init__(self, path_registry):
+        super().__init__()
+        self.path_registry = path_registry
+
+    def _run(self, pdb_id: str) -> str:
+        """use the tool."""
+        try:
+            RGY = RadiusofGyration(self.path_registry)
+            return RGY.rad_gyration_per_frame(pdb_id)
+        except ValueError as e:
+            return str(e)
+
+    async def _arun(self, query: str) -> str:
+        """Use the tool asynchronously."""
+        raise NotImplementedError("custom_search does not support async")
+
+
+class RadiusofGyrationPlot(BaseTool):
+    name = "RadiusofGyrationPlot"
+    description = """This tool calculates the radius of gyration
+    at each frame of a given trajectory file and plots it.
+    Give this tool the protein ID (PDB ID) only.
+    The tool will automatically find the necessary files.
+    The tool will save the plot to a png file and map it to the registry."""
+
+    path_registry: Optional[PathRegistry]
+
+    def __init__(self, path_registry):
+        super().__init__()
+        self.path_registry = path_registry
+
+    def _run(self, pdb_id: str) -> str:
+        """use the tool."""
+        try:
+            RGY = RadiusofGyration(self.path_registry)
+            return RGY.plot_rad_gyration(pdb_id)
+        except ValueError as e:
+            return str(e)
+
+    async def _arun(self, query: str) -> str:
+        """Use the tool asynchronously."""
+        raise NotImplementedError("custom_search does not support async")
diff --git a/mdagent/tools/maketools.py b/mdagent/tools/maketools.py
index 8c966b08..e5041cde 100644
--- a/mdagent/tools/maketools.py
+++ b/mdagent/tools/maketools.py
@@ -23,6 +23,8 @@
     PPIDistance,
     ProteinName2PDBTool,
     RadiusofGyrationAverage,
+    RadiusofGyrationPerFrame,
+    RadiusofGyrationPlot,
     RMSDCalculator,
     Scholar2ResultLLM,
     SetUpandRunFunction,
@@ -87,6 +89,8 @@ def make_all_tools(
         SmallMolPDB(path_registry=path_instance),
         VisualizeProtein(path_registry=path_instance),
         RadiusofGyrationAverage(path_registry=path_instance),
+        RadiusofGyrationPerFrame(path_registry=path_instance),
+        RadiusofGyrationPlot(path_registry=path_instance),
         PPIDistance(),
         RMSDCalculator(),
         SetUpandRunFunction(path_registry=path_instance),