added plot and traj rgy tools

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",

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",

mdagent/tools/base_tools/analysis_tools/analysis_props.py → 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")

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),