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