Save only one optimization csv/png for each solving run

princeton-vl · Aug 26, 2024 · 5b785e6 · 5b785e6
1 parent 9732bae
commit 5b785e6
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 3 deletions.
diff --git a/infinigen/core/constraints/example_solver/annealing.py b/infinigen/core/constraints/example_solver/annealing.py
@@ -63,6 +63,7 @@ def __init__(
         self.last_eval_result = None
 
         self.eval_memo = {}
+        self.stats = []
 
     def save_stats(self, path):
         if len(self.stats) == 0:
@@ -76,7 +77,7 @@ def save_stats(self, path):
         fig, ax1 = plt.subplots()
         ax1.set_xlabel("Iteration")
         ax1.set_ylabel("Score", color="C0")
-        ax1.plot(df["curr_iteration"], df["loss"], color="C0")
+        ax1.plot(np.arange(len(df)), df["loss"], color="C0")
 
         # ax2 = ax1.twinx()
         # ax2.set_ylabel('Move Time', color='C1')
@@ -91,7 +92,6 @@ def save_stats(self, path):
 
     def reset(self, max_iters):
         self.curr_iteration = 0
-        self.stats = []
         self.curr_result = None
         self.best_loss = None
         self.eval_memo = {}

diff --git a/infinigen/core/constraints/example_solver/solve.py b/infinigen/core/constraints/example_solver/solve.py
@@ -189,7 +189,6 @@ def solve_objects(
         for j in ra:
             move_gen = self.choose_move_type(moves, j, n_steps)
             self.optim.step(consgraph, self.state, move_gen, filter_domain)
-        self.optim.save_stats(self.output_folder / f"optim_{desc}.csv")
 
         logger.info(
             f"Finished solving {desc_full}, added {len(self.state.objs) - n_start} "

diff --git a/infinigen_examples/generate_indoors.py b/infinigen_examples/generate_indoors.py
@@ -291,6 +291,8 @@ def solve_small():
 
     p.run_stage("solve_small", solve_small, use_chance=False, default=state)
 
+    solver.optim.save_stats(output_folder / "optim_records.csv")
+
     p.run_stage(
         "populate_assets", populate.populate_state_placeholders, state, use_chance=False
     )