optimizers/bofire_optimizer.py

"""BoFire optimizer backend for physics-informed BO."""

from typing import Callable, Dict, List, Optional, Tuple

import torch
from torch import Tensor

from physics_informed_bo.config import OptimizationConfig
from physics_informed_bo.optimizers.base_optimizer import BaseOptimizer


class BoFireOptimizer(BaseOptimizer):
    """BoFire backend for chemistry/materials-focused Bayesian optimization.

    BoFire is designed for real-world experimental design in chemistry
    and materials science. It supports:
    - Complex parameter spaces (continuous, categorical, molecular)
    - Mixture constraints (sum-to-one)
    - Multi-objective optimization with Pareto fronts
    - Integration with domain-specific descriptors

    The physics model is incorporated as a prior mean function in BoFire's
    surrogate model specification.
    """

    def __init__(self, config: OptimizationConfig):
        super().__init__(config)
        self._domain = None
        self._strategy = None
        self._experiments_df = None

    def setup_domain(
        self,
        parameters: Dict[str, Dict],
        objectives: Dict[str, Dict],
        constraints: Optional[List[Dict]] = None,
    ) -> None:
        """Set up a BoFire domain with physics-informed features.

        Args:
            parameters: Dict of parameter specifications.
                Example: {"temp": {"type": "continuous", "bounds": (300, 500)}}
            objectives: Dict of objective specifications.
                Example: {"yield": {"type": "maximize", "weight": 1.0}}
            constraints: Optional list of constraint specifications.
                Example: [{"type": "linear", "features": ["x1", "x2"], "coeffs": [1, 1], "rhs": 1}]
        """
        try:
            from bofire.data_models.domain.api import Domain, Inputs, Outputs
            from bofire.data_models.features.api import (
                ContinuousInput,
                ContinuousOutput,
                CategoricalInput,
            )
            from bofire.data_models.objectives.api import MaximizeObjective, MinimizeObjective
            from bofire.data_models.constraints.api import (
                LinearInequalityConstraint,
                LinearEqualityConstraint,
            )
        except ImportError:
            raise ImportError(
                "BoFire is required for BoFireOptimizer. "
                "Install with: pip install bofire"
            )

        # Build input features
        input_features = []
        self._feature_names = []

        for name, spec in parameters.items():
            self._feature_names.append(name)
            if spec["type"] == "continuous":
                lb, ub = spec["bounds"]
                input_features.append(
                    ContinuousInput(key=name, bounds=(float(lb), float(ub)))
                )
            elif spec["type"] == "categorical":
                input_features.append(
                    CategoricalInput(key=name, categories=spec["categories"])
                )

        # Build output features
        output_features = []
        for name, spec in objectives.items():
            if spec.get("type", "maximize") == "maximize":
                obj = MaximizeObjective(w=spec.get("weight", 1.0))
            else:
                obj = MinimizeObjective(w=spec.get("weight", 1.0))
            output_features.append(ContinuousOutput(key=name, objective=obj))

        # Build constraints
        bofire_constraints = []
        if constraints:
            for c in constraints:
                if c["type"] == "linear_inequality":
                    bofire_constraints.append(
                        LinearInequalityConstraint(
                            features=c["features"],
                            coefficients=c["coeffs"],
                            rhs=c["rhs"],
                        )
                    )
                elif c["type"] == "linear_equality":
                    bofire_constraints.append(
                        LinearEqualityConstraint(
                            features=c["features"],
                            coefficients=c["coeffs"],
                            rhs=c["rhs"],
                        )
                    )

        self._domain = Domain(
            inputs=Inputs(features=input_features),
            outputs=Outputs(features=output_features),
            constraints=bofire_constraints if bofire_constraints else None,
        )

    def setup_strategy(self, strategy_type: str = "sobo") -> None:
        """Set up the BoFire optimization strategy.

        Args:
            strategy_type: One of 'sobo' (single-objective), 'mobo' (multi-objective),
                          'qehvi' (q-Expected Hypervolume Improvement).
        """
        try:
            from bofire.data_models.strategies.api import SoboStrategy, QehviStrategy
            from bofire.data_models.acquisition_functions.api import qEI, qNEI
            import bofire.strategies.api as strategies
        except ImportError:
            raise ImportError("BoFire is required. Install with: pip install bofire")

        if self._domain is None:
            raise RuntimeError("Call setup_domain() before setup_strategy().")

        if strategy_type == "sobo":
            strategy_data = SoboStrategy(domain=self._domain, acquisition_function=qEI())
        elif strategy_type in ("mobo", "qehvi"):
            strategy_data = QehviStrategy(domain=self._domain)
        else:
            raise ValueError(f"Unsupported strategy type: {strategy_type}")

        self._strategy = strategies.map(strategy_data)

    def suggest(
        self,
        n_candidates: int = 1,
        X_observed: Optional[Tensor] = None,
        y_observed: Optional[Tensor] = None,
    ) -> Tensor:
        """Suggest next experiments using BoFire."""
        if self._strategy is None:
            raise RuntimeError("Call setup_domain() and setup_strategy() first.")

        import pandas as pd

        # Tell strategy about existing experiments
        if self._experiments_df is not None:
            self._strategy.tell(self._experiments_df)

        candidates_df = self._strategy.ask(n_candidates)
        candidates = torch.tensor(
            candidates_df[self._feature_names].values, dtype=torch.float64
        )

        # Filter through physics constraints
        candidates = self._filter_feasible(candidates)
        return candidates[:n_candidates]

    def update(self, X_new: Tensor, y_new: Tensor) -> None:
        """Update BoFire with new observations."""
        import pandas as pd

        data = {}
        for i, name in enumerate(self._feature_names):
            data[name] = X_new[:, i].numpy()

        # Assume single objective for now
        output_keys = [f.key for f in self._domain.outputs.features]
        for i, key in enumerate(output_keys):
            if y_new.dim() > 1 and y_new.shape[1] > i:
                data[key] = y_new[:, i].numpy()
            else:
                data[key] = y_new.squeeze().numpy()

        new_df = pd.DataFrame(data)
        if self._experiments_df is None:
            self._experiments_df = new_df
        else:
            self._experiments_df = pd.concat(
                [self._experiments_df, new_df], ignore_index=True
            )