"""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, }, }