Upload experiment/designer.py with huggingface_hub

experiment/designer.py

@@ -0,0 +1,241 @@
+"""ExperimentDesigner: the main entry point for designing experiments."""
+
+from typing import Callable, Dict, List, Optional, Tuple
+
+import torch
+from torch import Tensor
+
+from physics_informed_bo.config import OptimizationConfig, OptimizerBackend
+from physics_informed_bo.experiment.parameter_space import ParameterSpace
+from physics_informed_bo.models.hybrid_model import HybridSurrogate
+from physics_informed_bo.priors.prior_manager import PriorManager
+from physics_informed_bo.priors.data_prior import DataPrior
+from physics_informed_bo.priors.physics_prior import PhysicsPrior
+from physics_informed_bo.optimizers.factory import create_optimizer
+from physics_informed_bo.optimizers.base_optimizer import BaseOptimizer
+
+
+class ExperimentDesigner:
+    """High-level API for physics-informed Bayesian experiment design.
+
+    This is the main user-facing class. It orchestrates:
+    1. Parameter space definition
+    2. Physics and data prior management
+    3. Surrogate model selection and fitting
+    4. Acquisition function optimization
+    5. Experiment suggestion
+
+    Example:
+        designer = ExperimentDesigner(
+            parameter_space=space,
+            physics_fn=arrhenius_model,
+            initial_data=(X_init, y_init),
+        )
+
+        # Get next experiment suggestions
+        next_experiments = designer.suggest(n=3)
+
+        # After running experiments, update with results
+        designer.update(X_new, y_new)
+    """
+
+    def __init__(
+        self,
+        parameter_space: ParameterSpace,
+        physics_fn: Optional[Callable[[Tensor], Tensor]] = None,
+        initial_data: Optional[Tuple[Tensor, Tensor]] = None,
+        config: Optional[OptimizationConfig] = None,
+        physics_constraints: Optional[List[Dict]] = None,
+    ):
+        """
+        Args:
+            parameter_space: The experimental parameter space.
+            physics_fn: Optional physics model function.
+            initial_data: Optional tuple of (X, y) initial observations.
+            config: Optimization configuration. Defaults to sensible settings.
+            physics_constraints: Optional list of physics constraint dicts.
+        """
+        self.parameter_space = parameter_space
+        self.config = config or OptimizationConfig()
+
+        # Set up physics prior
+        physics_prior = None
+        if physics_fn is not None:
+            physics_prior = PhysicsPrior(physics_fn=physics_fn)
+            if physics_constraints:
+                for c in physics_constraints:
+                    physics_prior.add_constraint(**c)
+
+        # Set up data prior
+        data_prior = DataPrior()
+        if initial_data is not None:
+            X_init, y_init = initial_data
+            if y_init.dim() == 1:
+                y_init = y_init.unsqueeze(-1)
+            data_prior.X = X_init
+            data_prior.y = y_init
+            data_prior.feature_names = parameter_space.parameter_names
+
+        # Prior manager
+        self.prior_manager = PriorManager(
+            physics_prior=physics_prior,
+            data_prior=data_prior,
+        )
+
+        # Build surrogate
+        self._surrogate: Optional[HybridSurrogate] = None
+        self._optimizer: Optional[BaseOptimizer] = None
+        self._iteration = 0
+
+        # Initialize if we have enough data
+        self._initialize()
+
+    def _initialize(self) -> None:
+        """Initialize surrogate model and optimizer."""
+        try:
+            mode = self.prior_manager.recommend_surrogate_mode()
+        except ValueError:
+            # Not enough data or physics model
+            return
+
+        self._surrogate = self.prior_manager.build_surrogate(
+            mode=mode,
+            kernel="matern",
+            noise_variance=self.config.noise_variance,
+            device=self.config.device,
+        )
+
+        # Set up optimizer
+        self._optimizer = create_optimizer(self.config)
+        self._optimizer.set_surrogate(self._surrogate)
+        self._optimizer.set_bounds(self.parameter_space.bounds)
+
+        if self.prior_manager.physics_prior:
+            self._optimizer.set_physics_prior(self.prior_manager.physics_prior)
+
+    def suggest(self, n: int = 1) -> Tensor:
+        """Suggest the next n experiments to run.
+
+        If not enough data exists for GP-based suggestion, falls back to:
+        1. Physics-guided sampling (if physics model available)
+        2. Latin Hypercube sampling (space-filling design)
+
+        Args:
+            n: Number of experiments to suggest.
+
+        Returns:
+            Tensor of shape (n, d) with suggested parameter values.
+        """
+        self._iteration += 1
+
+        # Not enough data for BO: use initial design
+        if self._surrogate is None or self.prior_manager.data_prior.n_observations < 3:
+            return self._initial_design(n)
+
+        # Re-fit surrogate with latest data
+        data = self.prior_manager.data_prior
+        if data.n_observations >= 3:
+            self._surrogate.fit(data.X, data.y)
+            self._optimizer.set_surrogate(self._surrogate)
+
+        # Suggest via optimizer
+        candidates = self._optimizer.suggest(
+            n_candidates=n,
+            X_observed=data.X,
+            y_observed=data.y,
+        )
+
+        return candidates
+
+    def _initial_design(self, n: int) -> Tensor:
+        """Generate initial design points when insufficient data for BO.
+
+        Uses physics model to prioritize promising regions if available.
+        """
+        if self.prior_manager.physics_prior is not None:
+            # Sample candidates and pick those with best physics predictions
+            n_candidates = max(n * 20, 200)
+            candidates = self.parameter_space.sample_latin_hypercube(n_candidates)
+
+            # Filter by physics constraints
+            candidates = self.prior_manager.physics_prior.get_feasible_subset(candidates)
+            if len(candidates) < n:
+                candidates = self.parameter_space.sample_latin_hypercube(n_candidates)
+
+            # Rank by physics model prediction
+            with torch.no_grad():
+                physics_scores = self.prior_manager.physics_prior.evaluate(candidates)
+
+            # Select top-n diverse points (greedy furthest-point selection)
+            selected = self._select_diverse_top_k(candidates, physics_scores, n)
+            return selected
+        else:
+            return self.parameter_space.sample_latin_hypercube(n)
+
+    def _select_diverse_top_k(
+        self, X: Tensor, scores: Tensor, k: int, top_fraction: float = 0.3
+    ) -> Tensor:
+        """Select k diverse points from the top-scoring candidates."""
+        # Pre-filter to top fraction
+        n_top = max(k * 3, int(len(X) * top_fraction))
+        top_idx = scores.argsort(descending=True)[:n_top]
+        X_top = X[top_idx]
+
+        # Greedy furthest-point selection for diversity
+        selected_idx = [0]
+        for _ in range(k - 1):
+            dists = torch.cdist(X_top, X_top[selected_idx]).min(dim=1).values
+            next_idx = dists.argmax().item()
+            selected_idx.append(next_idx)
+
+        return X_top[selected_idx]
+
+    def update(self, X_new: Tensor, y_new: Tensor) -> None:
+        """Update the designer with new experimental observations.
+
+        Args:
+            X_new: New input observations (n, d).
+            y_new: New output observations (n, 1) or (n,).
+        """
+        if y_new.dim() == 1:
+            y_new = y_new.unsqueeze(-1)
+
+        self.prior_manager.update_with_observations(X_new, y_new)
+
+        # Re-initialize if we now have enough data
+        if self._surrogate is None and self.prior_manager.data_prior.n_observations >= 3:
+            self._initialize()
+
+    def get_best(self, maximize: bool = True) -> Dict:
+        """Get the best observation so far."""
+        X_best, y_best = self.prior_manager.data_prior.get_best(maximize)
+        params = self.parameter_space.to_dict(X_best.unsqueeze(0))[0]
+        return {"parameters": params, "objective": float(y_best)}
+
+    def predict(self, X: Tensor) -> Tuple[Tensor, Tensor]:
+        """Get surrogate model predictions at X."""
+        if self._surrogate is None:
+            if self.prior_manager.physics_prior:
+                pred = self.prior_manager.physics_prior.evaluate(X)
+                return pred.unsqueeze(-1), torch.ones_like(pred.unsqueeze(-1)) * 0.1
+            raise RuntimeError("No surrogate model fitted yet.")
+        return self._surrogate.predict(X)
+
+    def model_quality(self) -> Dict:
+        """Assess current surrogate model quality."""
+        if self._surrogate is None:
+            return {"status": "no_model"}
+        return self._surrogate.physics_model_quality()
+
+    def summary(self) -> Dict:
+        """Get a summary of the current optimization state."""
+        return {
+            "iteration": self._iteration,
+            "n_observations": self.prior_manager.data_prior.n_observations,
+            "prior_summary": self.prior_manager.summary(),
+            "model_quality": self.model_quality(),
+            "parameter_space": {
+                "dimension": self.parameter_space.dimension,
+                "parameters": self.parameter_space.parameter_names,
+            },
+        }