GP Model

A Gaussian Process (GP) model implementation for protein sequence prediction. This model uses a BoTorch SingleTaskGP backbone with a configurable kernel to model sequence fitness, returning both predicted means and variances (uncertainty estimates). The trained backbone is exposed via the botorch_model property for use with native BoTorch acquisition functions.

GPTrainConfig defaults to normalise_inputs_strategy="minmax" and standardise_outputs=True: GP kernels measure distances between inputs, so min-max scaling to [0, 1] improves marginal log-likelihood optimisation, and output standardisation improves numerical stability.

class alf_tools.models.gp.FeaturizerConfig(featurizer_type='precomputed', custom_featurizer=None, flatten_one_hot=True)

Bases: object

Configuration for sequence featurization.

Parameters:

• featurizer_type (Literal['one_hot', 'custom', 'precomputed']) – Type of featurization to use. - ‘one_hot’: One-hot encoding of amino acids - ‘custom’: User-provided featurization function - ‘precomputed’: Features are pre-computed and provided directly

• custom_featurizer (Optional[Callable[[list[str]], Float[Tensor, 'batch_size n_features']]]) – Custom featurization function. Only used when featurizer_type=’custom’. Should take list of sequences and return torch.Tensor of shape (batch_size, num_features).

• flatten_one_hot (bool) – Whether to flatten one-hot encoded sequences from (batch_size, alphabet_size, seq_length) to (batch_size, alphabet_size * seq_length). Only used for one_hot encoding.

custom_featurizer: Callable[[list[str]], Float[Tensor, 'batch_size n_features']] | None = None

featurizer_type: Literal['one_hot', 'custom', 'precomputed'] = 'precomputed'

flatten_one_hot: bool = True

class alf_tools.models.gp.GPModel(name='gp_model', model_config=None, train_config=None, featurizer_config=None, alphabet='ARNDCQEGHILKMFPSTWYV', device=None)

Bases: BaseModel

Gaussian Process model for sequence fitness prediction.

Uses BoTorch’s SingleTaskGP backbone for efficient GP inference with flexible featurisation, kernel selection, and uncertainty quantification.

Inputs are featurised via the configurable featuriser (one-hot, custom, or precomputed) and optionally passed through an external InputNormaliser and OutputStandardiser fitted at train time. The likelihood and kernel priors are configured through GPModelConfig, while fitting is controlled by GPTrainConfig, which exposes three optimiser options (‘adam’, ‘lbfgs’, and the robust fit_gpytorch_mll-based ‘scipy’) and defaults to float64.

predict() returns the noise-inclusive predictive mean and variance. The trained backbone is exposed via the botorch_model property for use with native BoTorch acquisition functions; its posterior(X) gives the noise-free latent posterior.

property botorch_model: SingleTaskGP

Get the underlying BoTorch model.

Returns:

Trained BoTorch SingleTaskGP model.

Raises:

RuntimeError – If the model has not been trained yet.

feature_dim: int | None

featurise(inputs)

Convert inputs to feature tensors.

Parameters:

inputs (LabelledCandidates | list[Candidate]) – Either LabelledCandidates or list of Candidates to featurize.

Return type:

Float[Tensor, 'batch_size n_features']

Returns:

Feature tensor of shape (batch_size, num_features).

Raises:

ValueError – If input type is invalid or featurizer_type is unsupported.

get_epoch_metrics()

Return per-epoch training metrics recorded during hyperparameter optimisation.

Returns:

List of metrics for each epoch, including training loss and any additional metrics.The length of the list corresponds to the number of training iterations completed.

Return type:

list[SurrogateEpochMetrics]

get_hyperparameters()

Get current GP hyperparameters.

Returns:

• noise: Likelihood noise

• lengthscale: Kernel lengthscale(s)

• outputscale: Kernel output scale

• mean_constant: Mean function constant (if applicable)

Return type:

dict[str, Any]

Raises:

RuntimeError – If model is not trained.

get_training_summary_metrics()

Return training metrics including learned hyperparameters.

Returns:

• final_mll: Final marginal log likelihood

• final_loss: Final optimisation loss

• num_iterations: Number of optimisation iterations completed

• final_train_*: Metrics on the training set (e.g. final_train_mse, final_train_spearman)

• final_val_*: Metrics on the validation set, present only if validation data was provided

Learned hyperparameters (noise, lengthscale, outputscale) are available separately via get_hyperparameters.

Return type:

dict[str, float | int | number]

gp_model: SingleTaskGP | None

likelihood: GaussianLikelihood | None

predict(candidate_points)

Make predictions with uncertainty quantification.

Returns the noise-inclusive predictive variance (latent variance plus likelihood noise). The noise-free latent posterior is available via botorch_model.posterior(X).

Parameters:

candidate_points (list[Candidate]) – List of candidates to predict fitness for.

Return type:

Predictions

Returns:

Predictions containing predicted fitness means and variances.

Raises:

• RuntimeError – If the model is not trained.

• ValueError – If the input is invalid, not 2-D after featurisation, or its feature dimension does not match the training data.

problem_type: ProblemType

sample(*args, **kwargs)

Sample candidate points from the GP posterior.

This could be used to generate new sequences by: 1. Sampling from the GP in feature space 2. Decoding features back to sequences (requires inverse featurization)

Note

This is challenging for discrete sequence spaces and may not be practically useful. Consider raising NotImplementedError.

Return type:

list[Candidate]

setup(dataset)

Validate that the dataset uses REGRESSION problem type.

Parameters:

dataset (BaseDataset) – The dataset this model will be trained on.

Raises:

ValueError – If the dataset problem_type is not REGRESSION.

Return type:

None

train(train_data, val_data=None)

Train the GP model by optimizing hyperparameters.

Parameters:

• train_data (LabelledCandidates) – Training data containing sequences and oracle values.

• val_data (LabelledCandidates | None) – Validation data (used for monitoring, not for training).

Raises:

RuntimeError – If GP fitting produces a NaN/Inf loss.

Return type:

None

Note

For exact GPs, all training data is used for predictions. Validation data is only used for logging validation metrics during training.

This method is not thread-safe. Concurrent calls to train() and predict() on the same instance will produce undefined behaviour.

train_x: Float[Tensor, 'n_samples n_features'] | None

train_y: Float[Tensor, 'n_samples'] | None

training_metrics: dict[str, float | int | number]

class alf_tools.models.gp.GPModelConfig(kernel_type='rbf', matern_nu=2.5, ard=True, mean_type='constant', lengthscale_prior=<factory>, lengthscale_constraint=None, outputscale_prior=None, noise_prior=<factory>, noise_constraint=None, build_kernel_fn=None)

Bases: object

Configuration for Gaussian Process model.

Parameters:

• kernel_type (Literal['rbf', 'matern', 'linear', 'polynomial', 'rbf_linear', 'custom']) – Type of kernel to use. Supported: ‘rbf’, ‘matern’, ‘linear’, ‘polynomial’, ‘rbf_linear’ (RBF + Linear).

• matern_nu (float) – Smoothness parameter for Matern kernel (0.5, 1.5, or 2.5). Only used when kernel_type=’matern’.

• ard (bool) – Whether to use Automatic Relevance Determination (separate lengthscale per dimension).

• mean_type (Literal['constant', 'zero']) – Type of mean function (‘constant’ or ‘zero’).

• lengthscale_prior (dict | None) – Prior for the kernel lengthscale, as a _target_ dict. Defaults to LogNormal(sqrt(2), sqrt(3)) (Hvarfner et al. 2024 fixed form). Set to None to use no prior.

• lengthscale_constraint (dict | None) – Constraint on the kernel lengthscale, as a _target_ dict. Defaults to None (no constraint).

• outputscale_prior (dict | None) – Prior for the kernel output scale, as a _target_ dict. Defaults to None.

• noise_prior (dict | None) – Prior for the likelihood noise, as a _target_ dict. Defaults to LogNormal(-4.0, 1.0), matching SingleTaskGP’s default likelihood noise prior (Hvarfner et al. 2024). Set to None for a prior-free likelihood.

• noise_constraint (dict | None) – Constraint on the likelihood noise, as a _target_ dict. Defaults to None (a GreaterThan(1e-4) fallback is applied internally).

• build_kernel_fn (Optional[Callable[..., Kernel]]) – Optional custom function to build the kernel. If provided, used instead of the default kernel construction logic. Not serializable — for advanced Python-only use. Should return a gpytorch.kernels.Kernel.

ard: bool = True

build_kernel_fn: Callable[[...], Kernel] | None = None

kernel_type: Literal['rbf', 'matern', 'linear', 'polynomial', 'rbf_linear', 'custom'] = 'rbf'

lengthscale_constraint: dict | None = None

lengthscale_prior: dict | None

matern_nu: float = 2.5

mean_type: Literal['constant', 'zero'] = 'constant'

noise_constraint: dict | None = None

noise_prior: dict | None

outputscale_prior: dict | None = None

class alf_tools.models.gp.GPTrainConfig(learning_rate=0.01, log_frequency=10, normalise_inputs_strategy='minmax', standardise_outputs=True, label_dtype=None, num_iterations=100, optimizer_type='adam', max_attempts=5, dtype='float64', early_stopping_patience=None, early_stopping_delta=0.0001)

Bases: BaseTrainConfig

Configuration for Gaussian Process training.

Overrides normalise_inputs_strategy to minmax because GP kernels measure distances between inputs; scaling continuous features to [0, 1] improves marginal log-likelihood optimisation.

Parameters:

• normalise_inputs_strategy (Literal[‘minmax’, ‘zscore’] | None) – Which input normalisation to apply before training. Defaults to minmax; GP kernels measure distances so scaling continuous features to [0, 1] improves MLL.

• standardise_outputs (bool) – Whether to apply Z-score standardisation to output labels before training. Defaults to True; standardisation can improve training. For constant labels, std is clamped to a minimum value to avoid division by zero, but standardisation may not be meaningful.

• num_iterations (int) – Number of optimisation iterations.

• optimizer_type (Literal[‘adam’, ‘lbfgs’, ‘scipy’]) – Type of optimizer to use (‘adam’, ‘lbfgs’ or ‘scipy’). ‘scipy’ fits via BoTorch’s fit_gpytorch_mll with scipy L-BFGS-B.

• max_attempts (int) – Maximum fitting attempts on numerical failure. Only used when optimizer_type='scipy'.

• dtype (str) – Working dtype for features, labels, and the GP, as a string (e.g. ‘float32’, ‘float64’). Defaults to ‘float64’.

• early_stopping_patience (int | None) – Number of iterations without improvement before stopping. If None, no early stopping is used. Only used for ‘adam’ and ‘lbfgs’.

• early_stopping_delta (float) – Minimum change in loss to qualify as an improvement. Only used for ‘adam’ and ‘lbfgs’.

• learning_rate (float) – Inherited from BaseTrainConfig. Default overridden to 0.01.

• log_frequency (int) – Inherited from BaseTrainConfig. Default: 10.

• label_dtype (TorchDtype | None) – Inherited from BaseTrainConfig. None (the default) falls back to dtype. If set, it must match dtype — SingleTaskGP requires features and labels to share one dtype, so a mismatch raises ValueError at model construction.

dtype: str = 'float64'

early_stopping_delta: float = 0.0001

early_stopping_patience: int | None = None

learning_rate: float = 0.01

max_attempts: int = 5

normalise_inputs_strategy: Literal['minmax', 'zscore'] | None = 'minmax'

num_iterations: int = 100

optimizer_type: Literal['adam', 'lbfgs', 'scipy'] = 'adam'

standardise_outputs: bool = True