Upload models/gp_model.py with huggingface_hub

models/gp_model.py

@@ -0,0 +1,269 @@
+"""GPyTorch-based Gaussian Process models with physics-informed priors."""
+
+from typing import Callable, Optional, Tuple
+
+import torch
+from torch import Tensor
+import gpytorch
+from gpytorch.models import ExactGP
+from gpytorch.means import ConstantMean, ZeroMean
+from gpytorch.kernels import ScaleKernel, RBFKernel, MaternKernel
+from gpytorch.likelihoods import GaussianLikelihood
+from gpytorch.distributions import MultivariateNormal
+from gpytorch.mlls import ExactMarginalLogLikelihood
+
+from botorch.models.gpytorch import GPyTorchModel
+from botorch.posteriors.gpytorch import GPyTorchPosterior
+from botorch.models.transforms.input import Normalize
+from botorch.models.transforms.outcome import Standardize
+
+from physics_informed_bo.models.base import SurrogateModel
+from physics_informed_bo.models.physics_model import PhysicsMeanFunction
+
+
+class _ExactGPModel(ExactGP, GPyTorchModel):
+    """Core GPyTorch ExactGP model with BoTorch compatibility."""
+
+    _num_outputs = 1
+
+    def __init__(
+        self,
+        train_X: Tensor,
+        train_y: Tensor,
+        likelihood: GaussianLikelihood,
+        mean_module: Optional[gpytorch.means.Mean] = None,
+        kernel: str = "matern",
+        ard_num_dims: Optional[int] = None,
+    ):
+        super().__init__(train_X, train_y.squeeze(-1), likelihood)
+
+        self.mean_module = mean_module or ConstantMean()
+
+        if kernel == "rbf":
+            base_kernel = RBFKernel(ard_num_dims=ard_num_dims)
+        elif kernel == "matern":
+            base_kernel = MaternKernel(nu=2.5, ard_num_dims=ard_num_dims)
+        else:
+            raise ValueError(f"Unknown kernel: {kernel}. Use 'rbf' or 'matern'.")
+
+        self.covar_module = ScaleKernel(base_kernel)
+
+    def forward(self, X: Tensor) -> MultivariateNormal:
+        mean = self.mean_module(X)
+        covar = self.covar_module(X)
+        return MultivariateNormal(mean, covar)
+
+
+class StandardGP(SurrogateModel):
+    """Standard Gaussian Process model (no physics, pure data-driven).
+
+    Uses GPyTorch for the GP and is BoTorch-compatible for optimization.
+    """
+
+    def __init__(
+        self,
+        kernel: str = "matern",
+        noise_variance: float = 0.01,
+        learn_noise: bool = True,
+        normalize_inputs: bool = True,
+        standardize_outputs: bool = True,
+        device: str = "cpu",
+        dtype: torch.dtype = torch.float64,
+    ):
+        self.kernel = kernel
+        self.noise_variance = noise_variance
+        self.learn_noise = learn_noise
+        self.normalize_inputs = normalize_inputs
+        self.standardize_outputs = standardize_outputs
+        self.device = torch.device(device)
+        self.dtype = dtype
+        self._model = None
+        self._likelihood = None
+
+    def fit(
+        self,
+        X: Tensor,
+        y: Tensor,
+        training_iterations: int = 100,
+        lr: float = 0.1,
+    ) -> None:
+        """Fit the GP model by optimizing the marginal log likelihood."""
+        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)
+
+        self._likelihood = GaussianLikelihood()
+        if not self.learn_noise:
+            self._likelihood.noise = self.noise_variance
+            self._likelihood.noise_covar.raw_noise.requires_grad_(False)
+
+        self._model = _ExactGPModel(
+            train_X=X,
+            train_y=y,
+            likelihood=self._likelihood,
+            kernel=self.kernel,
+            ard_num_dims=X.shape[-1],
+        ).to(device=self.device, dtype=self.dtype)
+
+        self._optimize_hyperparameters(X, y, training_iterations, lr)
+
+    def _optimize_hyperparameters(
+        self, X: Tensor, y: Tensor, n_iter: int, lr: float
+    ) -> None:
+        """Optimize GP hyperparameters via type-II MLE."""
+        self._model.train()
+        self._likelihood.train()
+
+        optimizer = torch.optim.Adam(self._model.parameters(), lr=lr)
+        mll = ExactMarginalLogLikelihood(self._likelihood, self._model)
+
+        for _ in range(n_iter):
+            optimizer.zero_grad()
+            output = self._model(X)
+            loss = -mll(output, y.squeeze(-1))
+            loss.backward()
+            optimizer.step()
+
+        self._model.eval()
+        self._likelihood.eval()
+
+    def predict(self, X: Tensor) -> Tuple[Tensor, Tensor]:
+        X = X.to(device=self.device, dtype=self.dtype)
+        self._model.eval()
+        self._likelihood.eval()
+
+        with torch.no_grad(), gpytorch.settings.fast_pred_var():
+            posterior = self._likelihood(self._model(X))
+            mean = posterior.mean.unsqueeze(-1)
+            variance = posterior.variance.unsqueeze(-1)
+
+        return mean, variance
+
+    def posterior(self, X: Tensor):
+        self._model.eval()
+        self._likelihood.eval()
+        return self._model.posterior(X)
+
+    @property
+    def model(self):
+        """Access the underlying BoTorch-compatible GP model."""
+        return self._model
+
+
+class PhysicsInformedGP(SurrogateModel):
+    """GP with a physics model as the mean function.
+
+    The GP prior mean is set to the physics model predictions, so the GP
+    learns the residual (discrepancy) between the physics model and reality.
+    This is the core model of the platform.
+
+    Architecture:
+        f(x) = physics(x) + GP_residual(x)
+        where GP_residual ~ GP(0, k(x, x'))
+    """
+
+    def __init__(
+        self,
+        physics_fn: Callable[[Tensor], Tensor],
+        kernel: str = "matern",
+        physics_output_scale: float = 1.0,
+        learnable_physics_scale: bool = True,
+        noise_variance: float = 0.01,
+        learn_noise: bool = True,
+        device: str = "cpu",
+        dtype: torch.dtype = torch.float64,
+    ):
+        self.physics_fn = physics_fn
+        self.kernel = kernel
+        self.physics_output_scale = physics_output_scale
+        self.learnable_physics_scale = learnable_physics_scale
+        self.noise_variance = noise_variance
+        self.learn_noise = learn_noise
+        self.device = torch.device(device)
+        self.dtype = dtype
+        self._model = None
+        self._likelihood = None
+
+    def fit(
+        self,
+        X: Tensor,
+        y: Tensor,
+        training_iterations: int = 200,
+        lr: float = 0.05,
+    ) -> None:
+        """Fit the physics-informed GP model."""
+        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)
+
+        self._likelihood = GaussianLikelihood()
+        if not self.learn_noise:
+            self._likelihood.noise = self.noise_variance
+            self._likelihood.noise_covar.raw_noise.requires_grad_(False)
+
+        physics_mean = PhysicsMeanFunction(
+            physics_fn=self.physics_fn,
+            output_scale=self.physics_output_scale,
+            learnable_scale=self.learnable_physics_scale,
+        )
+
+        self._model = _ExactGPModel(
+            train_X=X,
+            train_y=y,
+            likelihood=self._likelihood,
+            mean_module=physics_mean,
+            kernel=self.kernel,
+            ard_num_dims=X.shape[-1],
+        ).to(device=self.device, dtype=self.dtype)
+
+        self._optimize_hyperparameters(X, y, training_iterations, lr)
+
+    def _optimize_hyperparameters(
+        self, X: Tensor, y: Tensor, n_iter: int, lr: float
+    ) -> None:
+        self._model.train()
+        self._likelihood.train()
+
+        optimizer = torch.optim.Adam(self._model.parameters(), lr=lr)
+        mll = ExactMarginalLogLikelihood(self._likelihood, self._model)
+
+        for _ in range(n_iter):
+            optimizer.zero_grad()
+            output = self._model(X)
+            loss = -mll(output, y.squeeze(-1))
+            loss.backward()
+            optimizer.step()
+
+        self._model.eval()
+        self._likelihood.eval()
+
+    def predict(self, X: Tensor) -> Tuple[Tensor, Tensor]:
+        X = X.to(device=self.device, dtype=self.dtype)
+        self._model.eval()
+        self._likelihood.eval()
+
+        with torch.no_grad(), gpytorch.settings.fast_pred_var():
+            posterior = self._likelihood(self._model(X))
+            mean = posterior.mean.unsqueeze(-1)
+            variance = posterior.variance.unsqueeze(-1)
+
+        return mean, variance
+
+    def posterior(self, X: Tensor):
+        self._model.eval()
+        self._likelihood.eval()
+        return self._model.posterior(X)
+
+    @property
+    def model(self):
+        return self._model
+
+    def get_residuals(self, X: Tensor, y: Tensor) -> Tensor:
+        """Compute residuals between physics predictions and observations."""
+        with torch.no_grad():
+            physics_pred = self.physics_fn(X)
+        return y.squeeze() - physics_pred