Bax visualization extension

.pre-commit-config.yaml

@@ -2,7 +2,7 @@
 # See https://pre-commit.com/hooks.html for more hooks
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks.git
-    rev: v5.0.0
+    rev: v6.0.0
     hooks:
       - id: no-commit-to-branch
       - id: trailing-whitespace
@@ -26,8 +26,13 @@ repos:
         exclude: ^src/badger/tests/
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.12.2
+    rev: v0.15.16
     hooks:
       - id: ruff-check
         args: [--fix]
       - id: ruff-format
+  # - repo: https://github.com/pre-commit/mirrors-mypy
+  #   rev: v2.1.0
+  #   hooks:
+  #     - id: mypy
+  #       args: [--strict, --ignore-missing-imports]

pyproject.toml

@@ -6,9 +6,7 @@ build-backend = "setuptools.build_meta"
 name = "badger-opt"
 description = "Interface for optimization of arbitrary problems in Python."
 readme = "README.md"
-authors = [
-  { name = "Zhe Zhang", email = "zhezhang@slac.stanford.edu" },
-]
+authors = [{ name = "Zhe Zhang", email = "zhezhang@slac.stanford.edu" }]
 keywords = ["optimization", "machine learning", "GUI"]
 classifiers = [
     "Development Status :: 4 - Beta",
@@ -30,18 +28,15 @@ dependencies = [
     "pillow",
     "requests",
     "xopt>=3.0.0",
+    "bax-algorithms",
+
 ]
 dynamic = ["version"]
 [tool.setuptools_scm]
 version_file = "src/badger/_version.py"
 
 [project.optional-dependencies]
-dev = [
-    "pytest",
-    "pytest-cov",
-    "pytest-qt",
-    "pytest-mock"
-]
+dev = ["pytest", "pytest-cov", "pytest-qt", "pytest-mock", "pyqt5-stubs"]
 
 [project.urls]
 Homepage = "https://github.com/xopt-org/Badger"
@@ -60,12 +55,12 @@ include_package_data = true
 
 [tool.setuptools.packages.find]
 where = ["src"]
-include = [ "badger", ]
+include = ["badger"]
 namespaces = false
 
 [tool.ruff.lint]
 extend-select = ["TID252"] # Defaults + check imports
-ignore = ["E722"] # Until bare except blocks get fixed
+ignore = ["E722"]          # Until bare except blocks get fixed
 
 [tool.pytest.ini_options]
 addopts = "--cov=badger/"

src/badger/actions/__init__.py

@@ -4,9 +4,10 @@
 
 import os
 from importlib import metadata
+
 from badger.actions.doctor import check_n_config_paths
+from badger.settings import get_user_config_folder, init_settings
 from badger.utils import yprint
-from badger.settings import init_settings, get_user_config_folder
 
 
 def show_info(args):
@@ -43,6 +44,7 @@ def show_info(args):
     BADGER_LOGBOOK_ROOT = config_singleton.read_value("BADGER_LOGBOOK_ROOT")
     BADGER_ARCHIVE_ROOT = config_singleton.read_value("BADGER_ARCHIVE_ROOT")
     BADGER_LOG_DIRECTORY = config_singleton.read_value("BADGER_LOG_DIRECTORY")
+    BADGER_TEMP_DIRECTORY = config_singleton.read_value("BADGER_TEMP_DIRECTORY")
     BADGER_LOG_LEVEL = config_singleton.read_value("BADGER_LOG_LEVEL")
     BADGER_TENSOR_STRATEGY = config_singleton.read_value(
         "BADGER_PYTORCH_TENSOR_SHARING_STRATEGY"
@@ -59,6 +61,7 @@ def show_info(args):
         "logging directory": BADGER_LOG_DIRECTORY,
         "logging level": BADGER_LOG_LEVEL,
         "pytorch tensor sharing strategy": BADGER_TENSOR_STRATEGY,
+        "temporary directory": BADGER_TEMP_DIRECTORY,
         # 'plugin installation url': read_value('BADGER_PLUGINS_URL')
     }
 

src/badger/factory.py

@@ -9,26 +9,26 @@
 Also handles loading Markdown docs for the built-in documentation browser.
 """
 
-from typing import Any, TypedDict, cast, TYPE_CHECKING
-from badger.settings import init_settings
-from badger.utils import get_value_or_none
+import importlib
+import logging
+import os
+import re
+import sys
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, TypedDict, cast
+
+import yaml
+from xopt.generators import generators, get_generator_defaults
+
 from badger.errors import (
     BadgerConfigError,
-    BadgerInvalidPluginError,
     BadgerInvalidDocsError,
+    BadgerInvalidPluginError,
     BadgerPluginNotFoundError,
 )
-
 from badger.interface import Interface as BadgerInterface
-import sys
-import os
-import importlib
-import yaml
-import re
-from pathlib import Path
-from xopt.generators import generators, get_generator_defaults
-
-import logging
+from badger.settings import init_settings
+from badger.utils import get_value_or_none
 
 if TYPE_CHECKING:
     from badger.environment import Environment as BadgerEnvironment
@@ -44,7 +44,7 @@
     "time_dependent_upper_confidence_bound",
     "multi_fidelity",
     "nsga2",
-    "bax",
+    # "bax",
 ]
 
 

src/badger/gui/components/analysis_extensions.py

@@ -1,31 +1,36 @@
 """Dialog wrappers for analysis extensions (BO visualizer, Pareto front
 viewer). Each dialog receives live data updates from the run monitor."""
 
-from typing import Optional, cast
 import logging
+from typing import Optional, cast
 
-from PyQt5.QtCore import pyqtSignal
-from PyQt5.QtWidgets import QDialog, QVBoxLayout
+from PyQt5.QtCore import Qt, pyqtSignal
 from PyQt5.QtGui import QCloseEvent
+from PyQt5.QtWidgets import QVBoxLayout, QWidget
+from xopt import Generator
+from xopt.generators.bayesian.bayesian_generator import BayesianGenerator
+from xopt.generators.bayesian.mobo import MOBOGenerator
+
 from badger.gui.components.analysis_widget import AnalysisWidget
+from badger.gui.components.bax_visualizer.bax_widget import BaxWidget
 from badger.gui.components.bo_visualizer.bo_widget import BOPlotWidget
 from badger.gui.components.pf_viewer.pf_widget import ParetoFrontWidget
 from badger.routine import Routine
 
-from xopt.generators.bayesian.bayesian_generator import BayesianGenerator
-from xopt.generators.bayesian.mobo import MOBOGenerator
-from xopt import Generator
-
 logger = logging.getLogger(__name__)
 
 
-class AnalysisExtension(QDialog):
+class AnalysisExtension(QWidget):
     window_closed = pyqtSignal(object)
-    generator_type = Generator
-    widget = AnalysisWidget
+    generator_type: type[Generator]
+    widget: AnalysisWidget
 
-    def __init__(self, parent: Optional[QDialog] = None):
+    def __init__(self, parent: Optional[QWidget] = None):
         super().__init__(parent=parent)
+        # A parented QWidget is a child widget by default. Setting the Window
+        # flag keeps the palette as the owner (for lifetime/stacking) while
+        # still displaying this as a separate top-level window.
+        self.setWindowFlag(Qt.WindowType.Window)
 
     def update_window(self, routine: Routine) -> None:
         try:
@@ -62,8 +67,6 @@ def update_extension(
         This method should be implemented to handle the update logic for the extension.
         """
 
-        self.widget = cast(AnalysisWidget, self.widget)
-
         self.widget.isValidRoutine(routine)
 
         self.widget.update_routine(routine, self.generator_type)
@@ -73,7 +76,7 @@ def update_extension(
 
         self.widget.update_plots(requires_rebuild, interval=self.widget.update_interval)
 
-    def closeEvent(self, a0: Optional[QCloseEvent]) -> None:
+    def closeEvent(self, a0: QCloseEvent) -> None:
         self.window_closed.emit(self)
         super().closeEvent(a0)
 
@@ -82,7 +85,7 @@ class ParetoFrontViewer(AnalysisExtension):
     def __init__(
         self,
         routine: Routine,
-        parent: Optional[QDialog] = None,
+        parent: Optional[QWidget] = None,
     ):
         super().__init__(parent=parent)
 
@@ -97,12 +100,31 @@ class BOVisualizer(AnalysisExtension):
     def __init__(
         self,
         routine: Routine,
-        parent: Optional[QDialog] = None,
+        parent: Optional[QWidget] = None,
     ):
         super().__init__(parent=parent)
 
         self.initialize_extension(
-            extension_widget=BOPlotWidget(routine=routine),
+            extension_widget=BOPlotWidget(
+                routine=routine,
+            ),
             extension_name="Bayesian Optimization Visualizer",
-            generator_type=BayesianGenerator,
+            generator_type=cast(type[Generator], BayesianGenerator),
+        )
+
+
+class BaxVisualizer(AnalysisExtension):
+    def __init__(
+        self,
+        routine: Routine,
+        parent: Optional[QWidget] = None,
+    ):
+        super().__init__(parent=parent)
+
+        self.initialize_extension(
+            extension_widget=BaxWidget(
+                routine=routine,
+            ),
+            extension_name="Bax Visualizer",
+            generator_type=cast(type[Generator], BayesianGenerator),
         )

src/badger/gui/components/analysis_widget.py

@@ -1,22 +1,22 @@
 """Base class that all analysis extension widgets must implement.
 Defines the interface for receiving routine updates and rendering plots."""
 
+import logging
 from abc import abstractmethod
 from collections.abc import Callable
 from typing import Any, Optional
 
-from PyQt5.QtWidgets import QDialog
-from badger.gui.components.extension_utilities import HandledException
-from badger.routine import Routine
+from PyQt5.QtWidgets import QWidget
 from xopt import Generator
 
-import logging
-
+from badger.gui.components.extension_utilities import HandledException
+from badger.routine import Routine
+from badger.utils import create_archive_run_filename
 
 logger = logging.getLogger(__name__)
 
 
-class AnalysisWidget(QDialog):
+class AnalysisWidget(QWidget):
     routine: Routine
     generator: Generator
     parameters: dict[str, Any] = {}
@@ -30,7 +30,7 @@ class AnalysisWidget(QDialog):
     def __init__(
         self,
         routine: Routine,
-        parent: Optional[QDialog] = None,
+        parent: Optional[QWidget] = None,
     ):
         super().__init__(parent=parent)
         self.routine = routine
@@ -44,14 +44,6 @@ def initialize_widget(self) -> None:
         """
         raise NotImplementedError("initialize_widget method not implemented")
 
-    @abstractmethod
-    def requires_reinitialization(self) -> bool:
-        """
-        Check if the widget requires reinitialization.
-        This is used to determine if the widget needs to be reset or updated.
-        """
-        raise NotImplementedError("requires_reinitialization method not implemented")
-
     @abstractmethod
     def update_plots(self, requires_rebuild: bool, interval: int) -> None:
         """
@@ -76,6 +68,68 @@ def isValidRoutine(self, routine: Routine) -> None:
         """
         raise NotImplementedError("isValidRoutine method not implemented")
 
+    @abstractmethod
+    def reset_widget(self) -> None:
+        """
+        Reset the widget to its initial state.
+        This method should be implemented to clear the current state and prepare the widget for a new routine.
+        """
+        raise NotImplementedError("reset_widget method not implemented")
+
+    def requires_reinitialization(self) -> bool:
+        # Check if the extension needs to be reinitialized
+        logger.debug("Checking if AnalysisWidget needs to be reinitialized")
+
+        archive_name = create_archive_run_filename(self.routine)
+
+        logger.debug(f"Archive name: {archive_name}")
+
+        if not self.initialized:
+            logger.debug("Reset - Extension never initialized")
+            # Set up connections
+            logger.debug("Setting up connections")
+            self.setup_connections()
+            self.routine_identifier = archive_name
+            self.initialized = True
+            # Track the current data length so the growth check does not treat
+            # the first post-init update as a shrink and reinitialize again.
+            if self.routine.data is not None:
+                self.df_length = len(self.routine.data)
+            return True
+
+        if self.routine_identifier != archive_name:
+            logger.debug("Reset - Routine name has changed")
+            # Reset first: reset_widget() clears routine_identifier, so the new
+            # identifier must be assigned afterwards. Assigning before the reset
+            # would be clobbered back to "" and force a reinitialization on every
+            # subsequent update during the same run.
+            self.reset_widget()
+            self.routine_identifier = archive_name
+            # Sync the tracked data length to the new routine so the growth
+            # check below does not immediately treat the next update as a
+            # shrink (df_length is left at inf by reset_widget()).
+            if self.routine.data is not None:
+                self.df_length = len(self.routine.data)
+            return True
+
+        if self.routine.data is None:
+            logger.debug("Reset - No data available")
+
+            return True
+
+        previous_len = self.df_length
+        self.df_length = len(self.routine.data)
+        new_length = self.df_length
+
+        if previous_len > new_length:
+            logger.debug("Reset - Data length is smaller")
+            # Keep df_length at the current (smaller) length rather than resetting
+            # it to inf. Leaving it at inf would make every subsequent update look
+            # like a shrink and reinitialize the widget on a loop.
+            return True
+
+        return False
+
     def update_routine(self, routine: Routine, generator_type: type[Generator]) -> None:
         self.routine = routine
 

src/badger/gui/components/bax_visualizer/.gitignore

@@ -0,0 +1 @@
+bax_algorithms/