Upload models/multi_fidelity.py with huggingface_hub

models/multi_fidelity.py

@@ -0,0 +1,142 @@
+"""Multi-fidelity model: physics model as low-fidelity, experiments as high-fidelity."""
+
+from typing import Callable, Optional, Tuple
+
+import torch
+from torch import Tensor
+from botorch.models import SingleTaskMultiFidelityGP
+from botorch.models.transforms.outcome import Standardize
+
+from physics_informed_bo.models.base import SurrogateModel
+
+
+class MultiFidelitySurrogate(SurrogateModel):
+    """Multi-fidelity BO model using physics as cheap low-fidelity source.
+
+    Uses BoTorch's SingleTaskMultiFidelityGP to jointly model:
+    - Low-fidelity: physics model predictions (cheap, approximate)
+    - High-fidelity: experimental observations (expensive, accurate)
+
+    The model learns the correlation between fidelities to transfer
+    knowledge from physics to improve experimental predictions.
+    """
+
+    def __init__(
+        self,
+        physics_fn: Callable[[Tensor], Tensor],
+        fidelity_dim: int = -1,
+        device: str = "cpu",
+        dtype: torch.dtype = torch.float64,
+    ):
+        """
+        Args:
+            physics_fn: Physics model function (low-fidelity source).
+            fidelity_dim: Column index for the fidelity indicator.
+            device: Torch device.
+            dtype: Torch dtype.
+        """
+        self.physics_fn = physics_fn
+        self.fidelity_dim = fidelity_dim
+        self.device = torch.device(device)
+        self.dtype = dtype
+        self._model = None
+
+    def build_multi_fidelity_data(
+        self,
+        X_experiment: Tensor,
+        y_experiment: Tensor,
+        X_physics_grid: Optional[Tensor] = None,
+        n_physics_points: int = 100,
+    ) -> Tuple[Tensor, Tensor]:
+        """Combine physics predictions and experimental data into multi-fidelity dataset.
+
+        Args:
+            X_experiment: Experimental inputs (n_exp, d).
+            y_experiment: Experimental outputs (n_exp, 1).
+            X_physics_grid: Optional grid for physics evaluations.
+            n_physics_points: Number of physics evaluation points if no grid given.
+
+        Returns:
+            X_mf: Combined inputs with fidelity column (n_total, d+1).
+            y_mf: Combined outputs (n_total, 1).
+        """
+        d = X_experiment.shape[-1]
+
+        # Generate physics data
+        if X_physics_grid is None:
+            X_physics_grid = torch.rand(
+                n_physics_points, d, device=self.device, dtype=self.dtype
+            )
+
+        with torch.no_grad():
+            y_physics = self.physics_fn(X_physics_grid).unsqueeze(-1)
+
+        # Add fidelity column: 0 = low (physics), 1 = high (experiment)
+        fidelity_low = torch.zeros(len(X_physics_grid), 1, device=self.device, dtype=self.dtype)
+        fidelity_high = torch.ones(len(X_experiment), 1, device=self.device, dtype=self.dtype)
+
+        X_physics_mf = torch.cat([X_physics_grid, fidelity_low], dim=-1)
+        X_experiment_mf = torch.cat([X_experiment, fidelity_high], dim=-1)
+
+        X_mf = torch.cat([X_physics_mf, X_experiment_mf], dim=0)
+        y_mf = torch.cat([y_physics, y_experiment], dim=0)
+
+        return X_mf, y_mf
+
+    def fit(
+        self,
+        X: Tensor,
+        y: Tensor,
+        training_iterations: int = 200,
+        lr: float = 0.05,
+    ) -> None:
+        """Fit the multi-fidelity GP model.
+
+        X should include a fidelity column (use build_multi_fidelity_data).
+        """
+        X = X.to(device=self.device, dtype=self.dtype)
+        y = y.to(device=self.device, dtype=self.dtype)
+        if y.dim() == 1:
+            y = y.unsqueeze(-1)
+
+        d = X.shape[-1]
+        fidelity_col = d - 1 if self.fidelity_dim == -1 else self.fidelity_dim
+
+        self._model = SingleTaskMultiFidelityGP(
+            train_X=X,
+            train_Y=y,
+            data_fidelities=[fidelity_col],
+            outcome_transform=Standardize(m=1),
+        ).to(device=self.device, dtype=self.dtype)
+
+        # Optimize hyperparameters
+        from botorch.fit import fit_gpytorch_mll
+        from gpytorch.mlls import ExactMarginalLogLikelihood
+
+        mll = ExactMarginalLogLikelihood(self._model.likelihood, self._model)
+        fit_gpytorch_mll(mll)
+
+    def predict(self, X: Tensor) -> Tuple[Tensor, Tensor]:
+        """Predict at high fidelity (fidelity=1)."""
+        X = X.to(device=self.device, dtype=self.dtype)
+
+        # Add high-fidelity indicator
+        if X.shape[-1] == self._model.train_inputs[0].shape[-1] - 1:
+            fidelity_col = torch.ones(len(X), 1, device=self.device, dtype=self.dtype)
+            X = torch.cat([X, fidelity_col], dim=-1)
+
+        posterior = self._model.posterior(X)
+        return posterior.mean, posterior.variance
+
+    def posterior(self, X: Tensor):
+        X = X.to(device=self.device, dtype=self.dtype)
+        if X.shape[-1] == self._model.train_inputs[0].shape[-1] - 1:
+            fidelity_col = torch.ones(
+                *X.shape[:-1], 1, device=self.device, dtype=self.dtype
+            )
+            X = torch.cat([X, fidelity_col], dim=-1)
+        return self._model.posterior(X)
+
+    @property
+    def model(self):
+        return self._model