pisa

tangermeme.pisa.pisa(model: ~torch.nn.modules.module.Module, X: ~torch.Tensor, args: tuple | None = None, batch_size: int = 32, references: ~collections.abc.Callable[[...], ~typing.Any] | ~torch.Tensor = <function dinucleotide_shuffle>, n_shuffles: int = 20, return_references: bool = False, hypothetical: bool = False, warning_threshold: float = 0.001, additional_nonlinear_ops: dict | None = None, print_convergence_deltas: bool = False, raw_outputs: bool = False, device: str | ~torch.device | None = None, random_state: int | ~numpy.random.mtrand.RandomState | None = None, verbose: bool = False) → Tensor | AttributionReferencesResult

An implementation of Pairwise Influence by Sequence Attribution (PISA).

PISA is a method for deciphering how each input in an example influences each output. For sequence-based machine learning models this usually means how each nucleotide influences the predictions made at each output. This can be useful for understanding precisely what parts of a predicted profile are influenced by the presence of certain motifs or other changes in nucleotide composition.

In practice, PISA involves running DeepLIFT/SHAP once for each output position. Given a fixed input and model, one might suspect that the Jacobian could be calculated simply to speed up DeepLIFT/SHAP, but this seems to not be the case. This is, in part, due to DeepLIFT/SHAP overriding some of the gradient operations in ways that do not play nicely with some of PyTorch’s built-in functions.

Parameters

model: torch.nn.Module: A PyTorch model to use for making predictions. These models can take in any number of inputs and make any number of outputs. The additional inputs must be specified in the args parameter.
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.
batch_size: int, optional: The number of sequence-reference pairs to pass through DeepLiftShap at a time. Importantly, this is not the number of elements in X that are processed simultaneously (alongside ALL their references) but the total number of X-reference pairs that are processed. This means that if you are in a memory-limited setting where you cannot process all references for even a single sequence simultaneously that the work is broken down into doing only a few references at a time. Default is 32.
references: func or torch.Tensor, optional: If a function is passed in, this function is applied to each sequence with the provided random state and number of shuffles. This function should serve to transform a sequence into some form of signal-null background, such as by shuffling it. If a torch.Tensor is passed in, that tensor must have shape (len(X), n_shuffles, *X.shape[1:]), in that for each sequence a number of shuffles are provided. Default is the function dinucleotide_shuffle.
n_shuffles: int, optional: The number of shuffles to use if a function is given for references. If a torch.Tensor is provided, this number is ignored. Default is 20.
return_references: bool, optional: Whether to return the references that were generated during this process. Only use if references is not a torch.Tensor. Default is False.
hypothetical: bool, optional: Whether to return attributions for all possible characters at each position (True) or only for the character that is actually in the sequence (False). When False, the per-character attributions are multiplied by the one-hot encoded input so that only the observed character has a non-zero attribution at each position. Default is False.
warning_threshold: float, optional: A threshold on the convergence delta that will always raise a warning if the delta is larger than it. Normal deltas are in the range of 1e-6 to 1e-8. Note that convergence deltas are calculated on the gradients prior to the aggr_func being applied to them. Default is 0.001.
additional_nonlinear_ops: dict or None, optional: If additional nonlinear ops need to be added to the dictionary of operations that can be handled by DeepLIFT/SHAP, pass a dictionary here where the keys are class types and the values are the name of the function that handle that sort of class. Make sure that the signature matches those of _nonlinear and _maxpool above. This can also be used to overwrite the hard-coded operations by passing in a dictionary with overlapping key names. If None, do not add any additional operations. Default is None.
print_convergence_deltas: bool, optional: Whether to print the convergence deltas for each example when using DeepLiftShap. Default is False.
raw_outputs: bool, optional: Whether to return the raw outputs from the method – in this case, the multipliers for each example-reference pair – or the processed attribution values. Default is False.
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.
random_state: int or None or numpy.random.RandomState, optional: The random seed to use to ensure determinism. If None, the process is not deterministic. Default is None.
verbose: bool, optional: Whether to display a progress bar. Default is False.

Returns

attributions: torch.tensor: If raw_outputs=False (default), the per-output attribution values with shape (X.shape[0], n_outputs, X.shape[1], X.shape[2]). If raw_outputs=True, the multipliers for each example-reference pair with shape (X.shape[0], n_shuffles, n_outputs, X.shape[1], X.shape[2]).
references: torch.tensor, optional: The references used for each input sequence, with the shape (n_input_sequences, n_shuffles, 4, length). Only returned if return_references = True.

Notes

n_outputs is inferred from the last axis of the model’s output (predict(model, X[:1]).shape[-1]). Profile-style models that return (batch, n_tasks, length) will therefore be attributed over length only — to attribute over tasks instead, wrap the model in a small adapter that transposes or flattens the output.

The returned attributions and references tensors are produced via torch.stack over per-example accumulators and may live on the input device rather than CPU in some return paths; explicit .cpu() is advisable when consuming results on a non-CUDA host.