saturation_mutagenesis

tangermeme.saturation_mutagenesis.saturation_mutagenesis(model: Module, X: Tensor, args: tuple | None = None, start: int = 0, end: int = -1, batch_size: int = 32, target: int | slice | None = None, hypothetical: bool = False, raw_outputs: bool = False, dtype: str | dtype | None = None, device: str | device | None = None, verbose: bool = False, func: Callable[[...], Any] | None = None) → Tensor | SaturationMutagenesisRawResult

Performs in-silico saturation mutagenesis on a set of sequences.

This function will perform in-silico saturation mutagenesis on a set of sequences and return the predictions on the original sequences and each of the sequences with an edit distance of one on them.

By default, this function will aggregate these predictions into an attribution value. For each single-character substitution, the change in prediction (perturbed minus original) is computed and then mean-subtracted across the alphabet axis so the values at each position are centered on zero; when target=None these are additionally averaged across all of the model’s outputs. The values are not Euclidean distances and are not Z-score normalized – preserving the raw magnitude lets you compare the importance of different regions head-to-head. Unless hypothetical=True, the result is then projected onto the observed bases by multiplying by the one-hot X. This aggregation assumes that the model returns only a single tensor. This tensor can have multiple outputs, e.g., be of shape (batch_size, n_targets) where n_targets > 1, but the model cannot return multiple tensors.

If you simply want the predictions before and after the substitutions without the method turning those into attributions because, perhaps, you want to define your own aggregation method, you can use raw_outputs=True.

Parameters

model: torch.nn.Module

The PyTorch model to use to make predictions.

X: torch.tensor, shape=(-1, len(alphabet), length)

A set of one-hot encoded sequences to calculate attribution values for.

args: tuple or None, optional

An optional set of additional arguments to pass into the model. If provided, each element in the tuple or list is one input to the model and the element must be formatted to be the same batch size as X. If None, no additional arguments are passed into the forward function. Default is None.

start: int, optional

The start of where to begin making perturbations to the sequence. Default is 0.

end: int, optional

The end of where to make perturbations to the sequence. Positive values follow standard Python slice semantics (non-inclusive). Negative values are remapped via end = length + 1 + end so that end=-1 maps to length and the last nucleotide is included. Note this differs from Python’s X[start:-1] (which would drop the last element). Default is -1, meaning the entire sequence.

batch_size: int, optional

The number of examples to make predictions for at a time. Default is 32.

target: int or slice or None, optional

Whether to focus on a single output/slice of outputs from the model when calculating attributions rather than the entire set of outputs. If None, use all targets when calculating attributions. Default is None.

hypothetical: bool, optional

Whether to return attributions for all possible characters at each position or only for the character that is actually in the sequence. Only matters when raw_outputs=False. Default is False.

raw_outputs: bool, optional

Whether to return the raw outputs from the method – in this case, the predictions from the reference sequence and from each of the perturbations – or the processed attribution values. Default is False.

dtype: str or torch.dtype or None, optional

The dtype to use with mixed precision autocasting. If None, use the dtype of the model. This allows you to use int8 to represent large data sets and only convert batches to the higher precision, saving memory. Default is None.

device: str or torch.device or None, optional

The device to move the model and batches to when making predictions. If None, use CUDA when available and fall back to CPU otherwise. Default is None.

verbose: bool, optional

Whether to display a progress bar during predictions. Default is False.

func: function or None, optional

A function to apply to a batch of predictions after they have been made, forwarded to predict. It is applied identically to the reference predictions and to every perturbation, so attributions are computed on the post-processed values. Use it to, e.g., select a single output head or apply a final non-linearity. If None, do nothing. Default is None.

Returns

attr: torch.Tensor

Processed attribution values. For each example and each position, the difference between the model’s prediction on the perturbed sequence and on the original sequence is computed, then mean-subtracted across the alphabet axis so the values at each position are centered on zero.

– or, if raw_outputs=True –

y0: torch.Tensor or list/tuple of torch.Tensors

The outputs from the model for the reference sequences.

y_hat: torch.Tensor or list/tuple of torch.Tensors

The outputs from the model for each of the perturbed sequences.