utils.attribution_utils

utils.attribution_utils#

Utilities for prediction, attribution, and visualization of nucleotide models.

This module provides small, practical helpers for working with PyTorch models that consume one-hot encoded nucleotide sequences and (optionally) per-position structure features. It covers three workflows:

Prediction helpers: forward pass + optional logits→probabilities conversion.
Attribution helpers: gradient-based explanations via igrads (IG / Grad×Input).
Visualization helpers: sequence-logo style plots via logomaker.

Who this is for#

Researchers/engineers who already have a trained PyTorch model and want a lightweight utility layer for (a) inference, (b) attribution, and (c) quick visualization.

Dependencies#

torch, torch.nn.functional
igrads (required for attribution)
matplotlib (required for plotting)
pandas, logomaker (required for logo plots)

Data conventions#

Nucleotide channel order

All logo/attribution plotting assumes channel order: ['A', 'C', 'G', 'U'] corresponding to columns 0..3.

Sequence encoding

The included encoder uses:

base2int = {'A': 0, 'C': 1, 'G': 2, 'T': 3}

This means the 4th channel is produced from 'T'. If your sequences are RNA and contain 'U', you should either pre-convert U→T before encoding, or extend the mapping (recommended):

base2int = {'A': 0, 'C': 1, 'G': 2, 'U': 3, 'T': 3}

Tensor shapes

One-hot sequence: - unbatched: (L, 4) (float32) - batched: (B, L, 4) (float32)
Structure features (optional, attribution workflow): - unbatched: (L, S) (float32) - batched: (B, L, S) (float32)

Model interface expectations#

This module supports two common model styles:

Sequence-only prediction models (used by predict / predict_from_sequence)
Expected forward signature:
model(one_hot) -> dict[str, torch.Tensor]
Output key naming convention (used by _to_probs): - keys containing '_profile' → softmax(dim=1) - keys containing '_mixing_coefficient' → sigmoid

If to_probs=True and a key does not match either rule, a ValueError is raised.
Sequence + structure attribution models (used by attribution)
Expected forward signature:
model((one_hot, structure)) -> torch.Tensor
attribution will add a batch dimension and call:
(inputs.unsqueeze(0), structure.unsqueeze(0))

Important

Target selection for attributions

igrads.integrated_gradients(...) / igrads.grad_x_input(...) require a clear scalar objective or a target specification that matches your output shape.

The current implementation passes target_mask=pred (the raw model output). This is only appropriate when your model output is already scalar-per-example or when your igrads version interprets target_mask in a compatible way.

If your model returns multi-dimensional outputs (profiles, multi-task heads, etc.), you will usually need to adapt the objective, e.g.:

pick an index (task/class) and reduce to a scalar
sum/mean over positions for profile outputs

Example pattern (conceptual):

pred = model(inputs)              # shape: (B, ...)
score = pred[..., idx].sum()      # scalar objective
# then run IG/Grad×Input against 'score' (depending on igrads API)

Notes and caveats#

Device placement: Inputs must be on the same device as the model (CPU/GPU).
Checkpoint loading: torch.load uses pickle; only load checkpoints from trusted sources.

Usage examples#

Prediction from raw sequence (sequence-only models)

model = load_model(MyModel(), "checkpoint.pt")
pred = model.predict_from_sequence("ACGTACGT...", to_probs=True)

Attribution for sequence + structure models

seq = sequence2onehot("ACGT...").to(device)          # (L, 4)
struct = torch.zeros(seq.shape[0], S).to(device)     # (L, S)
attrs_seq, attrs_struct = attribution(seq, struct, model, atype="IG", steps=50)

Visualization

fig1 = visualize_attribution_only(attrs_seq)
fig2 = visualize_track_attribution(track, attrs_seq, sequence="ACGT...", title="Example")

Functions

`attribution`(inputs, structure, model[, ...])	Compute sequence/structure attributions for a PyTorch model.
`load_model`(model, filepath, **kwargs)	Load a PyTorch model state dict from disk and attach convenience methods.
`make_attribution_figure`(a, ax)	Plot an attribution matrix (L x 4) as a sequence logo on a given Matplotlib axis.
`predict`(inputs, model[, to_probs])	Run a model forward pass on already-prepared inputs and optionally convert logits to probabilities.
`predict_from_sequence`(sequences, model, **kwargs)	Convenience wrapper: encode sequence(s) to one-hot, then call predict.
`sequence2int`(sequence)	Convert a nucleotide sequence string into integer indices using base2int.
`sequence2onehot`(sequence)	Convert a single sequence string into a one-hot encoded tensor.
`sequences2inputs`(sequences)	Convert one or more sequences into a batch of one-hot encoded tensors.
`visualize_attribution_only`(attribution)	Visualize only an attribution matrix (L x 4) as a logo.
`visualize_track_attribution`(track, attribution)	Visualize a predicted 1D track together with an attribution logo (and optionally the input sequence logo).

utils.attribution_utils.attribution(inputs, structure, model, atype='IG', steps=50)[source]#

Compute sequence/structure attributions for a PyTorch model.

Parameters:

inputs (torch.Tensor) – One-hot encoded sequence of shape (L, 4), dtype float.
structure (torch.Tensor) – Structure features of shape (L, S), dtype float.
model (torch.nn.Module) – A model whose forward accepts (inputs, structure) as a tuple and returns a prediction tensor.
atype ({“IG”, “grad_x_input”}, default=”IG”) – Attribution method: Integrated Gradients (“IG”) or Grad×Input (“grad_x_input”).
steps (int, default=50) – Number of IG interpolation steps. Only used when atype=”IG”.

Returns:

Any Attributions returned by igrads.* for (inputs, structure). Typically a tuple of tensors with shapes matching the inputs.

Raises:

ValueError – If atype is not supported.

utils.attribution_utils.make_attribution_figure(a, ax)[source]#

Plot an attribution matrix (L x 4) as a sequence logo on a given Matplotlib axis.

This function converts an attribution matrix into a pandas DataFrame with columns corresponding to nucleotide channels and then uses logomaker to render a logo. It also draws a horizontal baseline at y=0 and removes the bottom spine.

Parameters:

a (array-like) – Attribution matrix with shape (L, 4), where L is sequence length. Channel order is assumed to match [‘A’, ‘C’, ‘G’, ‘U’]. Values can be signed (e.g., importance scores) or non-negative (e.g., probabilities).
ax (matplotlib.axes.Axes) – Axis on which the logo will be drawn.

Returns:

None. The plot is drawn in-place on ax.

Requires:

logomaker installed and importable.

Visual conventions:

shade_below/fade_below highlight negative contributions by default.
a y=0 baseline is drawn in red.

utils.attribution_utils.visualize_track_attribution(track, attribution, sequence=None, title=None)[source]#

Visualize a predicted 1D track together with an attribution logo (and optionally the input sequence logo).

The output figure stacks panels vertically:

track panel (line plot)
attribution panel (logo)
optional sequence panel (logo of one-hot encoding)

Parameters:

track (array-like) – 1D signal of length L (e.g., predicted binding signal along the sequence). Must be plottable by Matplotlib (list, np.ndarray, torch.Tensor, etc.).
attribution (array-like or torch.Tensor) – Attribution matrix with shape (L, 4). If a torch.Tensor, it will be detached to CPU numpy. Channel order should match [‘A’,’C’,’G’,’U’] (see make_attribution_figure).
sequence (str, optional) – Optional RNA sequence of length L. If provided, it will be converted to one-hot and shown as a logo in the third panel. IMPORTANT: sequence2onehot (as written below) supports bases in base2int. Unknown bases will lead to failure in one-hot (see notes in module docstring).
title (str, optional) – Title for the top track panel.

Returns:

matplotlib.figure.Figure – The created figure instance.

utils.attribution_utils.visualize_attribution_only(attribution)[source]#

Visualize only an attribution matrix (L x 4) as a logo.

Parameters:: attribution (array-like or torch.Tensor) – Attribution matrix with shape (L, 4).
Returns:: matplotlib.figure.Figure – The created figure instance.

utils.attribution_utils.predict(inputs, model, to_probs=True)[source]#

Run a model forward pass on already-prepared inputs and optionally convert logits to probabilities.

Assumes model(inputs) returns a dictionary mapping output-name -> tensor.

Parameters:

inputs (torch.Tensor) –
Model inputs, typically one-hot encoded sequences of shape:
- (B, L, 4) for batch size B, length L, 4 channels.
The exact expected shape depends on your model implementation.
model (torch.nn.Module) – PyTorch model that returns a dict of outputs.
to_probs (bool) – If True, convert each output tensor from logits to probabilities using _to_probs based on the output key naming convention.

Returns:

dict[str, torch.Tensor] – Model outputs; either raw logits (to_probs=False) or probabilities (to_probs=True).

Raises:

ValueError – If to_probs=True and an output key is not recognized by _to_probs.

utils.attribution_utils.sequence2int(sequence)[source]#

Convert a nucleotide sequence string into integer indices using base2int.

Parameters:: sequence (str) – RNA sequence. Expected characters are keys of base2int (default: A/C/G/T).
Returns:: list[int] – Integer-encoded sequence of length L.

Warning

Unknown bases are mapped to 999 by default, which will later break one-hot encoding (torch.nn.functional.one_hot requires values < num_classes). Ensure sequences contain only valid bases before calling downstream helpers.

utils.attribution_utils.sequences2inputs(sequences)[source]#

Convert one or more sequences into a batch of one-hot encoded tensors.

Parameters:: sequences (str or list[str]) – If str, treated as a single sequence (length L). If list[str], treated as a batch of sequences (all should have equal length L for stacking).
Returns:: torch.Tensor – One-hot tensor of shape (B, L, 4), dtype float32, where B is batch size (1 if input is a single string).
Raises:: RuntimeError / ValueError – If sequences contain invalid bases (mapped to 999), one_hot will fail. If sequences have inconsistent lengths, tensor construction may fail.

utils.attribution_utils.sequence2onehot(sequence)[source]#

Convert a single sequence string into a one-hot encoded tensor.

Parameters:: sequence (str) – Single RNA sequence of length L.
Returns:: torch.Tensor – One-hot tensor of shape (L, 4), dtype float32.
Raises:: RuntimeError / ValueError – If sequence contains invalid bases (mapped to 999), one_hot will fail.

utils.attribution_utils.predict_from_sequence(sequences, model, **kwargs)[source]#

Convenience wrapper: encode sequence(s) to one-hot, then call predict.

Parameters:

sequences (str or list[str]) – Sequence(s) to predict on. - If str: returns outputs with batch dimension squeezed. - If list[str]: returns batched outputs.
model (torch.nn.Module) – PyTorch model that accepts one-hot inputs and returns a dict of outputs.
**kwargs – Passed through to predict, e.g. to_probs=False.

Returns:

dict[str, torch.Tensor] – Prediction dictionary. - If input is a single sequence string, each tensor will be squeezed from (1, …) to (…). - If input is a list, tensors remain batched.

utils.attribution_utils.load_model(model, filepath, **kwargs)[source]#

Load a PyTorch model state dict from disk and attach convenience methods.

Parameters:

model (torch.nn.Module) – Instantiated model object with the same architecture as the saved checkpoint.
filepath (str) – Path to a file produced by torch.save(model.state_dict(), filepath).
**kwargs – Reserved for future extensions (currently unused). You may use this to pass torch.load kwargs in your own fork (e.g., map_location), but as written it is not forwarded.

Returns:

torch.nn.Module – The same model instance with loaded parameters, patched methods, and set to eval() mode.

Important

This uses torch.load(filepath) directly. If you need CPU/GPU mapping, you may want to modify to torch.load(filepath, map_location=…).

utils.attribution_utils

Contents

utils.attribution_utils#

Who this is for#

Dependencies#

Data conventions#

Model interface expectations#

Notes and caveats#

Usage examples#