File size: 1,665 Bytes
058f8a4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 | """Abstract base classes for surrogate models."""
from abc import ABC, abstractmethod
from typing import Optional, Tuple
import torch
from torch import Tensor
class SurrogateModel(ABC):
"""Abstract base class for all surrogate models in the platform.
A surrogate model provides predictions (mean + uncertainty) and can be
updated with new observations.
"""
@abstractmethod
def predict(self, X: Tensor) -> Tuple[Tensor, Tensor]:
"""Return posterior mean and variance at input locations X.
Args:
X: Input tensor of shape (n, d).
Returns:
mean: Predicted mean of shape (n, 1).
variance: Predicted variance of shape (n, 1).
"""
@abstractmethod
def fit(self, X: Tensor, y: Tensor) -> None:
"""Fit/update the surrogate model with observed data.
Args:
X: Training inputs of shape (n, d).
y: Training targets of shape (n, 1).
"""
@abstractmethod
def posterior(self, X: Tensor):
"""Return the full posterior distribution at X (for BoTorch compatibility).
Args:
X: Input tensor of shape (batch, n, d).
"""
def condition_on_observations(self, X: Tensor, y: Tensor) -> "SurrogateModel":
"""Return a new model conditioned on additional observations.
Default implementation refits the model. Subclasses can override
for fantasy-based conditioning.
"""
raise NotImplementedError(
"Fantasy conditioning not implemented for this model. Use fit() instead."
)
|