rgy png fix

ur-whitelab · Nov 4, 2024 · ab25b48 · ab25b48
1 parent 320a853
commit ab25b48
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/mdagent/tools/base_tools/analysis_tools/rgy.py b/mdagent/tools/base_tools/analysis_tools/rgy.py
@@ -66,9 +66,8 @@ def plot_rgy(self) -> str:
         plot_id = self.path_registry.get_fileid(
             file_name=plot_name, type=FileType.FIGURE
         )
-        if plot_name.endswith(".png"):
-            plot_name = plot_name.split(".png")[0]
         plot_path = f"{self.path_registry.ckpt_figures}/{plot_name}"
+        plot_path = plot_path if plot_path.endswith(".png") else plot_path + ".png"
         print("plot_path", plot_path)
         plt.plot(rg_per_frame)
         plt.xlabel("Frame")
@@ -78,12 +77,12 @@ def plot_rgy(self) -> str:
         plt.savefig(f"{plot_path}")
         self.path_registry.map_path(
             plot_id,
-            plot_path + ".png",
+            plot_path,
             description=f"Plot of radii of gyration over time for {self.traj_file}",
         )
         plt.close()
         plt.clf()
-        return "Plot saved as: " + f"{plot_name}.png with plot ID {plot_id}"
+        return "Plot saved as: " + f"{plot_name} with plot ID {plot_id}"
 
 
 class RadiusofGyrationAverage(BaseTool):

diff --git a/tests/test_analysis/test_rgy_tool.py b/tests/test_analysis/test_rgy_tool.py
@@ -26,3 +26,4 @@ def test_rgy_average(rgy):
 def test_plot_rgy(rgy):
     output = rgy.plot_rgy()
     assert "Plot saved as: " in output
+    assert ".png" in output