From f439a7855625b13b80d72aba11d64d6ee61d63ea Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 18 Dec 2025 04:01:10 +0000
Subject: [PATCH 1/2] Initial plan


From b141f3fcd09b97df3679a8c4e9a9fd0b27bada5d Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 18 Dec 2025 04:05:16 +0000
Subject: [PATCH 2/2] Refactor: Define OUTPUT_DIR constant for visualization
 output directory

Co-authored-by: josephdviviano <4142570+josephdviviano@users.noreply.github.com>
---
 src/gfn/gym/diffusion_sampling.py | 23 +++++++++++++----------
 1 file changed, 13 insertions(+), 10 deletions(-)

diff --git a/src/gfn/gym/diffusion_sampling.py b/src/gfn/gym/diffusion_sampling.py
index 4eb50ea3..93a83882 100644
--- a/src/gfn/gym/diffusion_sampling.py
+++ b/src/gfn/gym/diffusion_sampling.py
@@ -27,6 +27,9 @@
 #   - Exit action trigger: t + dt >= 1.0 - dt * TERMINAL_TIME_EPS (next step reaches terminal)
 TERMINAL_TIME_EPS = 1e-2
 
+# Default output directory for saving visualizations
+OUTPUT_DIR = "output"
+
 
 ###############################
 ### Target energy functions ###
@@ -407,8 +410,8 @@ def visualize(
         if show:
             plt.show()
         else:
-            os.makedirs("output", exist_ok=True)
-            plt.savefig(f"output/{prefix}simple_gmm.png")
+            os.makedirs(OUTPUT_DIR, exist_ok=True)
+            plt.savefig(f"{OUTPUT_DIR}/{prefix}simple_gmm.png")
 
         plt.close()
 
@@ -479,8 +482,8 @@ def visualize(
         if show:
             plt.show()
         else:
-            os.makedirs("output", exist_ok=True)
-            fig.savefig(f"output/{prefix}gmm25.png")
+            os.makedirs(OUTPUT_DIR, exist_ok=True)
+            fig.savefig(f"{OUTPUT_DIR}/{prefix}gmm25.png")
         plt.close()
 
 
@@ -565,8 +568,8 @@ def visualize(
         if show:
             plt.show()
         else:
-            os.makedirs("output", exist_ok=True)
-            fig.savefig(f"output/{prefix}posterior9of25.png")
+            os.makedirs(OUTPUT_DIR, exist_ok=True)
+            fig.savefig(f"{OUTPUT_DIR}/{prefix}posterior9of25.png")
         plt.close()
 
 
@@ -670,8 +673,8 @@ def visualize(
         if show:
             plt.show()
         else:
-            os.makedirs("output", exist_ok=True)
-            fig.savefig(f"output/{prefix}funnel.png")
+            os.makedirs(OUTPUT_DIR, exist_ok=True)
+            fig.savefig(f"{OUTPUT_DIR}/{prefix}funnel.png")
 
         plt.close()
 
@@ -830,8 +833,8 @@ def visualize(
         if show:
             plt.show()
         else:
-            os.makedirs("output", exist_ok=True)
-            fig.savefig(f"output/{prefix}manywell.png")
+            os.makedirs(OUTPUT_DIR, exist_ok=True)
+            fig.savefig(f"{OUTPUT_DIR}/{prefix}manywell.png")
 
         plt.close()