discontinuum.engines.gpytorch

discontinuum.engines.gpytorch#

Data transformations to improve optimization

Classes

LatentGPyTorch([model_config])

MarginalGPyTorch([model_config])

NoOpMean()

class discontinuum.engines.gpytorch.LatentGPyTorch(model_config: Dict | None = None)#
fit(covariates, target=None)#

Fit model to data.

Parameters:

  • covariates (Dataset) – Covariates for training.

  • target (Dataset) – Target data for training.

  • kwargs (dict) – Additional keyword arguments.

class discontinuum.engines.gpytorch.MarginalGPyTorch(model_config: Dict | None = None)#
build_model(X, y, **kwargs) → ExactGP#

Creates an instance of pm.Model based on provided data and model_config, and attaches it to self.

The subclass method must instantiate self.model and self.likelihood.

Raises:

NotImplementedError –

fit(covariates: Dataset, target: Dataset, target_unc: Dataset = None, iterations: int = 100, optimizer: str | None = None, learning_rate: float = None, early_stopping: bool = False, patience: int = 60, scheduler: bool = True, resume: bool = False, penalty_callback: Callable[[], torch.Tensor] | None = None, penalty_weight: float = 0.0)#

Fit the model to data.

Parameters:

  • covariates (Dataset) – Covariates for training.

  • target (Dataset) – Target data for training.

  • target_unc (Dataset, optional) – Uncertainty on target data.

  • iterations (int, optional) – Number of iterations for optimization. The default is 100.

  • optimizer (str, optional) – Optimization method. Supported: “adam”, “adamw”. If None, uses the optimizer stored in a loaded checkpoint, otherwise defaults to “adam”.

  • learning_rate (float, optional) – Learning rate for optimization. If None, uses the value from checkpoint if available, otherwise a conservative default.

  • early_stopping (bool, optional) – Whether to use early stopping. The default is False.

  • patience (int, optional) – Number of iterations to wait without improvement before stopping. The default is 60.

  • scheduler (bool, optional) – Whether to use a learning rate scheduler. The default is True.

  • penalty_callback (callable, optional) – A zero-argument callable that returns a scalar torch.Tensor penalty to be added to the optimization objective each iteration. This enables model-specific regularization such as monotonicity constraints.

  • penalty_weight (float, optional) – Multiplier applied to the penalty value when added to the objective. Default 0.0.

classmethod load(f: str | PathLike[str] | IO[bytes], covariates: Dataset, target: Dataset, target_unc: Dataset | None = None) → MarginalGPyTorch#

Load a model from a checkpoint file and initialize for resume.

Parameters:

  • filepath (str) – Path to a checkpoint saved by save().

  • covariates (xarray.Dataset) – Training covariates used to rebuild the data manager.

  • target (xarray.DataArray or Dataset) – Training target used to rebuild the data manager.

  • target_unc (xarray.DataArray or Dataset, optional) – Training uncertainty used to rebuild the data manager.

predict(covariates: Dataset, diag=True, pred_noise=False) → Tuple[DataArray, DataArray]#

Uses the fitted model to make predictions on new data.

The input and output are in the original data space.

Parameters:

  • covariates (Dataset) – Covariates for prediction.

  • diag (bool, optional) – Return only the diagonal of the covariance matrix. The default is True.

  • pred_noise (bool, optional) – Include measurement uncertainty in the prediction. The default is False.

Returns:

  • target (DataArray) – Target prediction.

  • se (DataArray) – Standard error of the prediction.

predict_grid(covariate: str, coord: str = None, t_step: int = 12)#

Predict on a grid of points.

Parameters:

  • covariate (str) – Covariate dimension to predict on.

  • coord (str, optional) – Coordinate of the covariate dimension to predict on. The default is the first coordinate of the covariate.

  • t_step (int, optional) – Number of grid points per step in coord units. The default is 12.

sample(covariates, n=1000) → DataArray#

Sample from the posterior distribution of the model.

Parameters:

  • covariates (Dataset) – Covariates for prediction.

  • n (int, optional) – Number of samples to draw.

save(f: str | PathLike[str] | IO[bytes], optimizer_obj: torch.optim.Optimizer | None = None, scheduler: torch.optim.lr_scheduler._LRScheduler | None = None, extra: Dict | None = None) → None#

Save the model to a checkpoint file (weights + optimizer/scheduler).

Parameters:

  • filepath (str) – Path to save the checkpoint (.pt/.pth).

  • optimizer_obj (torch.optim.Optimizer, optional) – Optimizer whose state will be saved. If None, uses the last one from fit().

  • scheduler (torch.optim.lr_scheduler, optional) – LR scheduler whose state will be saved. If None, uses the last one from fit().

  • extra (dict, optional) – Extra metadata to include in the checkpoint.

class discontinuum.engines.gpytorch.NoOpMean#
forward(x)#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Contents