synbo 0.1.0__py3-none-any.whl

synbo/__init__.py

@@ -0,0 +1,22 @@
+"""Modern Reaction Optimization Framework.
+
+A sophisticated framework for optimizing chemical reactions using
+Bayesian Optimization with modern Python practices and rich output.
+"""
+
+from __future__ import annotations
+
+__version__ = "0.1.0"
+__author__ = "Zhenzhi Tan"
+__email__ = "zhenzhi-tan@outlook.com"
+
+# Core classes
+from .synbo import ReactionOptimizer
+from .initialize import Initializer
+from .optimize import Optimizer
+
+__all__ = [
+    "ReactionOptimizer",
+    "Initializer",
+    "Optimizer",
+]

synbo/algorithm/acq_function.py

@@ -0,0 +1,337 @@
+import numpy as np
+import gpytorch
+import torch
+from torch import Tensor
+from botorch.models import ModelListGP
+from botorch.acquisition.multi_objective.logei import qLogNoisyExpectedHypervolumeImprovement, qLogExpectedHypervolumeImprovement
+from botorch.acquisition.multi_objective.monte_carlo import qExpectedHypervolumeImprovement
+from botorch.utils.multi_objective.box_decompositions import NondominatedPartitioning
+from botorch.sampling.normal import SobolQMCNormalSampler
+
+from botorch.acquisition.monte_carlo import qUpperConfidenceBound, qExpectedImprovement
+from botorch.acquisition.logei import qLogNoisyExpectedImprovement
+from botorch.acquisition.objective import GenericMCObjective
+from botorch.utils.multi_objective.scalarization import get_chebyshev_scalarization
+from botorch.optim import optimize_acqf_discrete
+
+from synbo.utils.logger import console
+
+
+class BaseAcquisitionFunction:
+    """Base class for acquisition functions"""
+
+    def __init__(self, model: ModelListGP, sampler: SobolQMCNormalSampler):
+        self.model = model
+        self.sampler = sampler
+        self.acquisition_function = None
+        self.console = console
+
+    @property
+    def acq_func(self):
+        return self.acquisition_function
+
+    def _split_batch_eval_acqf(self, acq_func, X: Tensor, max_batch_size: int) -> Tensor:
+        """Helper to evaluate acquisition function in batches to avoid OOM."""
+        acq_values_list = []
+        with torch.no_grad():
+            for X_batches in X.split(max_batch_size):
+                acq_values = acq_func(X_batches)
+                acq_values_list.append(acq_values)
+        return torch.cat(acq_values_list, dim=0)
+
+    def optimize_discrete(
+        self,
+        acq_func,
+        q: int,
+        choices: Tensor,
+        unique: bool = True,
+        max_batch_size: int = 1024,
+        progress: object = None,
+        task: object = None,
+        min_distance: float = 1e-6,
+        exclude_points: Tensor = None,
+    ) -> tuple[Tensor, Tensor]:
+        """
+        Joint optimization for batch selection (EDBO-style).
+        This method selects q candidates simultaneously using BoTorch's optimize_acqf_discrete,
+        exactly as implemented in EDBO+.
+
+        Unlike the greedy sequence optimization, this approach considers the joint expected improvement
+        of all q candidates at once. This is the implementation used by EDBO+.
+
+        Args:
+            acq_func: Acquisition function to optimize (must support q-batch evaluation)
+            q: Number of candidates to select
+            choices: Candidate points tensor (n, D)
+            unique: Whether to ensure unique selections (default: True)
+            max_batch_size: Maximum batch size for evaluation (default: 128)
+            progress: Not used in EDBO implementation (kept for API compatibility)
+            task: Not used in EDBO implementation (kept for API compatibility)
+            min_distance: Not used in EDBO implementation (kept for API compatibility)
+            exclude_points: Not used in EDBO implementation (kept for API compatibility)
+
+        Returns:
+            tuple[Tensor, Tensor]:
+                - Selected candidates (q, D)
+                - Acquisition values for selected candidates (q,)
+        """
+
+        acq_result = self._new_optimize_acqf_discrete(
+            acq_function=acq_func,
+            choices=choices,
+            q=q,
+            unique=unique,
+            max_batch_size=max_batch_size,
+        )
+
+        selected_candidates = acq_result[0]  # (q, D)
+        acquisition_values = acq_result[1]  # (q,)
+
+        return selected_candidates, acquisition_values
+
+    def _new_optimize_acqf_discrete(self, acq_function, q, choices, max_batch_size, unique):
+
+        def _split_batch_eval_acqf(acq_function, X, max_batch_size):
+            return torch.cat([acq_function(X_) for X_ in X.split(max_batch_size)])
+
+        choices_batched = choices.unsqueeze(-2)
+        if q > 1:
+            candidate_list, acq_value_list = [], []
+            base_X_pending = acq_function.X_pending
+            for q_i in range(q):
+                with torch.no_grad():
+                    acq_values = _split_batch_eval_acqf(
+                        acq_function=acq_function,
+                        X=choices_batched,
+                        max_batch_size=max_batch_size,
+                    )
+
+                best_idx = torch.argmax(acq_values)
+                candidate_list.append(choices_batched[best_idx])
+                acq_value_list.append(acq_values[best_idx])
+                # set pending points
+                candidates = torch.cat(candidate_list, dim=-2)
+                acq_function.set_X_pending(torch.cat([base_X_pending, candidates], dim=-2) if base_X_pending is not None else candidates)
+                # need to remove choice from choice set if enforcing uniqueness
+                if unique:
+                    choices_batched = torch.cat([choices_batched[:best_idx], choices_batched[best_idx + 1 :]])
+
+            # Reset acq_func to previous X_pending state
+            acq_function.set_X_pending(base_X_pending)
+            return candidates, torch.stack(acq_value_list)
+
+        with torch.no_grad():
+            acq_values = _split_batch_eval_acqf(acq_function=acq_function, X=choices_batched, max_batch_size=max_batch_size)
+
+        best_idx = torch.argmax(acq_values)
+        return choices_batched[best_idx], acq_values[best_idx]
+
+
+class EHVIAcquisitionFunction(BaseAcquisitionFunction):
+    """Enhanced Expected Hypervolume Improvement acquisition function"""
+
+    def __init__(
+        self,
+        model: ModelListGP,
+        sampler: SobolQMCNormalSampler,
+        ref_point: torch.Tensor,
+        partitioning,
+        train_x,
+    ):
+        super().__init__(model, sampler)
+        self.ref_point = ref_point
+        self.partitioning = partitioning
+        self.acquisition_function = qLogNoisyExpectedHypervolumeImprovement(
+            model=model,
+            sampler=sampler,
+            ref_point=ref_point,
+            # partitioning=partitioning,
+            alpha=0.0,
+            incremental_nehvi=True,
+            X_baseline=train_x,
+            prune_baseline=True,
+        )
+
+        print(self.ref_point)
+
+    def optimize_acqf_discrete(
+        self,
+        q: int,
+        choices: Tensor,
+        max_batch_size: int = 1024,
+        unique: bool = True,
+        progress: object = None,
+        task: object = None,
+        min_distance: float = 1e-6,
+        exclude_points: Tensor = None,
+    ) -> tuple[Tensor, Tensor]:
+
+        return self.optimize_discrete(
+            acq_func=self.acquisition_function,
+            q=q,
+            choices=choices,
+            max_batch_size=max_batch_size,
+            unique=unique,
+            progress=progress,
+            task=task,
+            min_distance=min_distance,
+            exclude_points=exclude_points,
+        )
+
+
+class UCBAcquisitionFunction(BaseAcquisitionFunction):
+    """
+    Upper Confidence Bound acquisition function.
+    """
+
+    def __init__(
+        self,
+        model: ModelListGP,
+        sampler: SobolQMCNormalSampler,
+        beta: float = 2.0,
+        weights: Tensor = None,
+    ):
+        super().__init__(model, sampler)
+        self.beta = beta
+        objective = None
+        if weights is not None:
+            objective = GenericMCObjective(lambda Z, X: Z @ weights)
+
+        self.acquisition_function = qUpperConfidenceBound(
+            model=model,
+            beta=beta,
+            sampler=sampler,
+            objective=objective,
+        )
+
+    def optimize_acqf_discrete(
+        self,
+        q: int,
+        choices: Tensor,
+        max_batch_size: int = 1024,
+        unique: bool = True,
+        maximum_metrics: bool = True,
+        progress: object = None,
+        task: object = None,
+        min_distance: float = 1e-6,
+        exclude_points: Tensor = None,
+    ) -> tuple[Tensor, Tensor]:
+        # 直接调用基类的通用逻辑
+        return self.optimize_discrete(
+            acq_func=self.acquisition_function,
+            q=q,
+            choices=choices,
+            max_batch_size=max_batch_size,
+            unique=unique,
+            progress=progress,
+            task=task,
+            min_distance=min_distance,
+            exclude_points=exclude_points,
+        )
+
+
+class ParEGOAcquisitionFunction(BaseAcquisitionFunction):
+    """
+    Convert multiple objectives into single objectives using random Chebyshev scaling, and then apply EI.
+    """
+
+    def __init__(
+        self,
+        model: ModelListGP,
+        sampler: SobolQMCNormalSampler,
+        X_baseline: Tensor,
+        num_objectives: int,
+    ):
+        super().__init__(model, sampler)
+
+        weights = torch.randn(num_objectives).abs()
+        weights /= weights.sum()
+
+        with torch.no_grad():
+            posterior = model.posterior(X_baseline)
+            Y_baseline = posterior.mean
+
+        weights = weights.to(Y_baseline.device)
+
+        objective = self._get_chebyshev_objective(weights=weights, Y=Y_baseline)
+
+        scalarized_Y = objective(Y_baseline)
+        best_f = scalarized_Y.max()
+
+        self.acquisition_function = qExpectedImprovement(
+            model=model,
+            best_f=best_f,
+            sampler=sampler,
+            objective=objective,
+        )
+
+    def optimize_acqf_discrete(
+        self,
+        q: int,
+        choices: Tensor,
+        max_batch_size: int = 1024,
+        unique: bool = True,
+        maximum_metrics: bool = True,
+        progress: object = None,
+        task: object = None,
+        min_distance: float = 1e-6,
+        exclude_points: Tensor = None,
+    ) -> tuple[Tensor, Tensor]:
+        return self.optimize_discrete(
+            acq_func=self.acquisition_function,
+            q=q,
+            choices=choices,
+            max_batch_size=max_batch_size,
+            unique=unique,
+            progress=progress,
+            task=task,
+            min_distance=min_distance,
+            exclude_points=exclude_points,
+        )
+
+    @staticmethod
+    def _get_chebyshev_objective(weights: Tensor, Y: Tensor) -> GenericMCObjective:
+        """ """
+        scalarization_fn = get_chebyshev_scalarization(weights=weights, Y=Y)
+        return GenericMCObjective(scalarization_fn)
+
+
+class ParetoFrontCalculator:
+    """Class for calculating Pareto fronts"""
+
+    @staticmethod
+    def calculate_target_function(points: np.ndarray, progress: object, task: object) -> np.ndarray:
+        """
+        Calculate Pareto front for points in arbitrary dimensions
+
+        Args:
+            points: numpy array of shape (n_points, n_dimensions)
+
+        Returns:
+            numpy array of Pareto optimal points
+        """
+        if len(points) == 0:
+            return np.array([])
+        pareto_front = [points[0]]  # Initialize list of Pareto optimal points
+        for point in points[1:]:
+            progress.update(task, advance=1)
+            is_pareto = True
+            to_remove = []
+            # Compare with all points in current Pareto front
+            for i, pf_point in enumerate(pareto_front):
+                # Check if the current point dominates any existing Pareto point
+                if np.all(point >= pf_point) and np.any(point > pf_point):
+                    to_remove.append(i)
+                # Check if any existing Pareto point dominates the current point
+                elif np.all(point <= pf_point) and np.any(point < pf_point):
+                    is_pareto = False
+                    break
+
+            # Remove dominated points from Pareto front
+            for i in reversed(to_remove):
+                pareto_front.pop(i)
+
+            # Add current point if it's Pareto optimal
+            if is_pareto:
+                pareto_front.append(point)
+        return torch.tensor(np.array(pareto_front))

synbo/algorithm/bo_core.py

@@ -0,0 +1,257 @@
+from typing import List, Tuple
+import numpy as np
+import torch
+from rich.progress import Progress, TextColumn, BarColumn, MofNCompleteColumn, TimeRemainingColumn
+
+from botorch.models import ModelListGP
+from botorch.utils.multi_objective.box_decompositions import NondominatedPartitioning
+from botorch.sampling.normal import SobolQMCNormalSampler
+
+from synbo.utils.util_func import compute_hvi, generate_constraint_mask
+from synbo.utils.logger import console
+from synbo.algorithm.sg_model import (
+    BNNEnsembleSurrogateModel,
+    BayesianLinearSurrogateModel,
+    GPSurrogateModel,
+    RFSurrogateModel,
+    SklearnModelWrapper,
+)
+from synbo.algorithm.acq_function import (
+    EHVIAcquisitionFunction,
+    ParEGOAcquisitionFunction,
+    UCBAcquisitionFunction,
+    ParetoFrontCalculator,
+)
+
+import warnings
+from linear_operator.utils.cholesky import NumericalWarning
+
+warnings.filterwarnings("ignore", category=NumericalWarning)
+
+
+class DefaultBO:
+    def __init__(
+        self,
+        random_seed: int = 42,
+        surrogate_model: str = "GP",
+        acq_func: str = "EHVI",
+        device: torch.device = torch.device("cpu"),
+        accuracy: str = "medium",
+    ):
+        self.random_seed = random_seed
+        self.console = console
+
+        if accuracy == "medium":
+            self.mc_num_samples, self.max_batch_size = 256, 1024
+        self.device = device
+
+        if surrogate_model == "GP":
+            self.surrogate_model_class = GPSurrogateModel
+        elif surrogate_model == "RF":
+            self.surrogate_model_class = RFSurrogateModel
+        elif surrogate_model == "BNN":
+            self.surrogate_model_class = BNNEnsembleSurrogateModel
+        elif surrogate_model == "BayesianLinear":
+            self.surrogate_model_class = BayesianLinearSurrogateModel
+        else:
+            raise ValueError(f"Unknown surrogate model: {surrogate_model}")
+
+        if acq_func == "EHVI":
+            self.acquisition_function_class = EHVIAcquisitionFunction
+        elif acq_func == "UCB":
+            self.acquisition_function_class = UCBAcquisitionFunction
+        elif acq_func == "ParEGO":
+            self.acquisition_function_class = ParEGOAcquisitionFunction
+        else:
+            raise ValueError(f"Unknown acquisition function: {acq_func}")
+
+        self.target_evaluator = ParetoFrontCalculator()
+
+    def optimize(
+        self,
+        training_X: np.ndarray,
+        training_y: np.ndarray,
+        candidate_X: np.ndarray,
+        opt_metric_settings: List[dict],
+        batch_size: int,
+        total_name_arr: np.ndarray = None,
+        condition_types: List[str] = None,
+        total_desc_arr: np.ndarray = None,
+    ) -> Tuple[np.ndarray, List[str], np.ndarray, np.ndarray]:
+
+        training_X_t = torch.tensor(training_X).double().to(device=self.device)
+        training_y_t = torch.tensor(training_y).double()
+        training_y_t = self._weight_y(training_y_t, opt_metric_settings).to(device=self.device)
+        candidate_X_t = torch.tensor(candidate_X).double().to(device=self.device)
+
+        with Progress(
+            TextColumn("[bold cyan]{task.description}"),
+            BarColumn(bar_width=None),
+            MofNCompleteColumn(),
+            TimeRemainingColumn(),
+            console=self.console,
+        ) as progress:
+            num_models = training_y_t.shape[1]
+            task_train = progress.add_task(description="Training surrogate models", total=num_models, start=True)
+
+            models = []
+            for i in range(num_models):
+                # key = list(training_y_dict.keys())[i]
+                progress.log(f"Fitting the {i}th model...", style="yellow")
+
+                # Instantiate model with random_seed for reproducibility
+                if self.surrogate_model_class in [RFSurrogateModel, BNNEnsembleSurrogateModel]:
+                    model_i = self.surrogate_model_class(device=self.device, num_dims=training_X_t.shape[1], random_seed=self.random_seed)
+                else:
+                    model_i = self.surrogate_model_class(device=self.device, num_dims=training_X_t.shape[1])
+
+                if isinstance(model_i, GPSurrogateModel):
+                    model_i.fit(training_X_t, training_y_t[:, i].unsqueeze(-1))
+                    models.append(model_i.model)
+                else:
+                    wrapper = SklearnModelWrapper(model_i)
+                    wrapper.fit_surrogate(training_X_t, training_y_t[:, i].unsqueeze(-1))
+                    models.append(wrapper)
+
+                progress.update(task_train, advance=1)
+
+        self.global_model = ModelListGP(*models)
+
+        task_pareto = progress.add_task(description="Calculating Pareto frontiers", total=len(training_y) - 1)
+        training_y_np = training_y_t.cpu().numpy()
+        self.pareto_y = self.target_evaluator.calculate_target_function(training_y_np, progress, task_pareto).to(device=self.device)
+
+        y_min = training_y_t.min(dim=0).values
+        y_max = training_y_t.max(dim=0).values
+        y_range = y_max - y_min
+
+        ref_point_values = []
+        for i, omi in enumerate(opt_metric_settings):
+
+            ref_val = y_min[i] if y_range[i] > 0 else y_min[i] - 0.1
+            ref_point_values.append(ref_val)
+
+        self.ref_point = torch.tensor(ref_point_values, dtype=torch.double, device=self.device)
+
+        sampler = SobolQMCNormalSampler(sample_shape=torch.Size([self.mc_num_samples]), seed=self.random_seed)
+        if self.acquisition_function_class == EHVIAcquisitionFunction:
+            partitioning = NondominatedPartitioning(ref_point=self.ref_point, Y=self.pareto_y)
+            acq_func = self.acquisition_function_class(
+                model=self.global_model,
+                sampler=sampler,
+                ref_point=self.ref_point,
+                partitioning=partitioning,
+                train_x=torch.Tensor(training_X).to(self.device),
+            )
+        elif self.acquisition_function_class == UCBAcquisitionFunction:
+            weights = torch.tensor([d["metric_weight"] for d in opt_metric_settings], dtype=torch.double, device=self.device)
+            acq_func = self.acquisition_function_class(
+                model=self.global_model,
+                sampler=sampler,
+                beta=2.0,
+                weights=weights,
+            )
+        elif self.acquisition_function_class == ParEGOAcquisitionFunction:
+
+            acq_func = self.acquisition_function_class(
+                model=self.global_model, sampler=sampler, X_baseline=training_X_t, num_objectives=len(opt_metric_settings)
+            )
+        else:
+            raise ValueError(f"Unknown acquisition function class: {self.acquisition_function_class}")
+
+        task_acq_opt = progress.add_task(description="Optimizing acquisition function", total=batch_size)
+        self.acq_result, self.acq_value = acq_func.optimize_acqf_discrete(
+            q=batch_size,
+            choices=candidate_X_t,
+            max_batch_size=self.max_batch_size,
+            unique=True,
+            exclude_points=training_X_t,
+            min_distance=1e-6,
+            progress=progress,
+            task=task_acq_opt,
+        )
+
+        if self.device.type == "cuda":
+            best_samples = [res.cpu().numpy() for res in self.acq_result]
+        else:
+            best_samples = [res.numpy() for res in self.acq_result]
+
+        recommend_type = self._get_exploit_or_explore()
+        pred_mean, pred_std = self._get_predictions()
+
+        for i, d in enumerate(opt_metric_settings):
+            if d["opt_direct"] == "min":
+                pred_mean[:, i] = -pred_mean[:, i]
+
+        pred_mean = self._unweight_y(torch.tensor(pred_mean), opt_metric_settings).numpy()
+        pred_std = self._unweight_y(torch.tensor(pred_std), opt_metric_settings).numpy()
+        return best_samples, recommend_type, pred_mean, pred_std
+
+    def _get_exploit_or_explore(self) -> List[str]:
+        with torch.no_grad():
+            posterior = self.global_model.posterior(self.acq_result)
+            pred_mean = posterior.mean
+
+        # Handle different shapes from posterior.mean
+        # The posterior.mean from ModelListGP with multiple outputs has shape (batch_size, q)
+        # where each model contributes one dimension
+        # We need to reshape to (batch_size, n_outputs) where batch_size = n_points
+        if pred_mean.dim() == 2:
+            # Already in correct shape (n_points, n_outputs)
+            pass
+        elif pred_mean.dim() == 3:
+            # Shape: (1, n_points, n_outputs) -> (n_points, n_outputs)
+            pred_mean = pred_mean.squeeze(0)
+
+        # Ensure pred_mean is 2D
+        if pred_mean.dim() != 2:
+            raise ValueError(f"Unexpected pred_mean shape: {pred_mean.shape}")
+
+        hvi_values = torch.tensor([compute_hvi(pred_mean[i], self.pareto_y, self.ref_point) for i in range(pred_mean.shape[0])])
+        ehvi_values = self.acq_value.to(device="cpu")
+        exploit_scores = hvi_values / (ehvi_values + 1e-6)
+        explore_scores = 1 - exploit_scores
+
+        for i in range(self.acq_result.shape[0]):
+            self.console.log(
+                f"Point {i}: "
+                f"EHVI = {ehvi_values[i]:.3f}, "
+                f"HVI = {hvi_values[i]:.3f}, "
+                f"Exploit Score = {exploit_scores[i]:.3f}, "
+                f"Explore Score = {explore_scores[i]:.3f}"
+            )
+        return ["exploit" if exploit_scores[i] > explore_scores[i] else "explore" for i in range(self.acq_result.shape[0])]
+
+    def _get_predictions(self) -> Tuple[np.ndarray, np.ndarray]:
+        with torch.no_grad():
+            posterior = self.global_model.posterior(self.acq_result)
+            pred_mean = posterior.mean
+            pred_var = posterior.variance
+
+        # Handle different shapes from posterior
+        if pred_mean.dim() == 2:
+            # Already in correct shape (n_points, n_outputs)
+            pass
+        elif pred_mean.dim() == 3:
+            # Shape: (1, n_points, n_outputs) -> (n_points, n_outputs)
+            pred_mean = pred_mean.squeeze(0)
+            pred_var = pred_var.squeeze(0)
+
+        # Ensure correct shapes
+        if pred_mean.dim() != 2:
+            raise ValueError(f"Unexpected pred_mean shape: {pred_mean.shape}")
+
+        pred_mean = pred_mean.cpu().numpy()
+        pred_var = pred_var.cpu().numpy()
+        pred_std = np.sqrt(pred_var)
+        return pred_mean, pred_std
+
+    def _weight_y(self, training_y: torch.Tensor, opt_metric_settings: List[dict]) -> torch.Tensor:
+        weights = torch.tensor([d["metric_weight"] for d in opt_metric_settings])
+        training_y = training_y * weights
+        return training_y
+
+    def _unweight_y(self, training_y: torch.Tensor, opt_metric_settings: List[dict]) -> torch.Tensor:
+        weights = torch.tensor([d["metric_weight"] for d in opt_metric_settings])
+        training_y = training_y / weights
+        return training_y