dev: update visualization code

src/evox/monitors/eval_monitor.py

@@ -169,7 +169,7 @@ def get_best_fitness(self):
     def get_history(self):
         return [self.opt_direction * fit for fit in self.fitness_history]
 
-    def plot(self, **kwargs):
+    def plot(self, problem_pf=None, **kwargs):
         if not self.fitness_history:
             warnings.warn("No fitness history recorded, return None")
             return
@@ -182,9 +182,9 @@ def plot(self, **kwargs):
         if n_objs == 1:
             return plot.plot_obj_space_1d(self.fitness_history, **kwargs)
         elif n_objs == 2:
-            return plot.plot_obj_space_2d(self.fitness_history, **kwargs)
+            return plot.plot_obj_space_2d(self.fitness_history, problem_pf, **kwargs)
         elif n_objs == 3:
-            return plot.plot_obj_space_3d(self.fitness_history, **kwargs)
+            return plot.plot_obj_space_3d(self.fitness_history, problem_pf, **kwargs)
         else:
             warnings.warn("Not supported yet.")
 

src/evox/monitors/pop_monitor.py

@@ -66,7 +66,7 @@ def post_step(self, state):
             fitness = jax.device_put(fitness, self.host)
         self.fitness_history.append(fitness)
 
-    def plot(self, **kwargs):
+    def plot(self, problem_pf=None, **kwargs):
         if not self.fitness_history:
             warnings.warn("No fitness history recorded, return None")
             return
@@ -79,9 +79,9 @@ def plot(self, **kwargs):
         if n_objs == 1:
             return plot.plot_obj_space_1d(self.fitness_history, **kwargs)
         elif n_objs == 2:
-            return plot.plot_obj_space_2d(self.fitness_history, **kwargs)
+            return plot.plot_obj_space_2d(self.fitness_history, problem_pf, **kwargs)
         elif n_objs == 3:
-            return plot.plot_obj_space_3d(self.fitness_history, **kwargs)
+            return plot.plot_obj_space_3d(self.fitness_history, problem_pf, **kwargs)
         else:
             warnings.warn("Not supported yet.")
 
@@ -90,4 +90,3 @@ def get_population_history(self):
 
     def get_fitness_history(self):
         return self.fitness_history
-

src/evox/vis_tools/plot.py

@@ -158,7 +158,16 @@ def plot_obj_space_1d(fitness_history):
             go.Scatter(x=generation, y=max_fitness, mode="lines", name="Max"),
             go.Scatter(x=generation, y=median_fitness, mode="lines", name="Median"),
             go.Scatter(x=generation, y=avg_fitness, mode="lines", name="Average"),
-        ]
+        ],
+        layout=go.Layout(
+            legend={
+                "x": 1,
+                "y": 1,
+                "xanchor": "auto",
+                "xanchor": "auto",
+            },
+            margin={"l": 0, "r": 0, "t": 0, "b": 0},
+        ),
     )
 
     return fig