Add GATv2 attention-weight extraction

[dbdimitrov]

Extracts the learned GATv2 attention coefficients into adata.obsp-shaped, cell-by-cell sparse matrices — one per GAT layer, no averaging across layers.

Coefficients come from the official PyG API (GATv2Conv.forward(..., return_attention_weights=True)); alpha is never re-derived by hand.

API

def get_attention_weights(
    self,
    adata: AnnData | None = None,
    indices: list | None = None,
    batch_size: int = 512,
    num_neighbors: list[int] | None = None,   # default [-1] * n_layers
    give_mean: bool = True,
    key_added: str | None = None,
    allow_sampling: bool = False,
) -> dict[int, scipy.sparse.csr_matrix]

Returns {layer_index: csr_matrix}, each (n_obs, n_obs), float32. With indices a subset the shape is still (n_obs, n_obs); unselected rows are all-zero. key_added (opt-in) additionally writes adata.obsp[f"{key_added}_l{i}"]; the default never mutates adata.

Orientation: row = destination

A[i, j] = alpha(j -> i) — the weight cell i places on its neighbour j. Rows sum to 1 for non-isolated cells and to 0 for isolated ones. This matches the adj_t convention (row=dst, col=src) already used in GCNLayers.forward.

One finding worth flagging. The sparsity pattern is a subset of spatial_connectivities.T, not of spatial_connectivities itself. The data loader maps a nonzero spatial_connectivities[i, j] to a PyG edge i -> j (edge_index[0] = obsp.row, edge_index[1] = obsp.col), so it is cell j that attends over cell i. For a symmetric graph the two patterns coincide, but the kNN graphs from spatial_neighbors are not symmetric — on crc_210 out-degree is capped at GRAPH_K=20 while in-degree reaches 39. This is pre-existing message-passing behaviour, unchanged here, but it is now documented and pinned by a test on a deliberately asymmetric graph (a transposed implementation passes every other test in the section).

Determinism (give_mean)

Under condition_on_intrinsic=True, inference() feeds the reparameterised sample z.detach() into the GCN input. module.eval() disables dropout but not the draw, so attention would differ between two identical calls. give_mean=True (default) substitutes qzm. give_mean=False reproduces the sampled behaviour.

Seed-slice correctness

Under NeighborLoader, only edges whose destination is a seed node carry exact attention: a 1-hop node v's layer-1 alpha depends on the layer-0 outputs of v's neighbours, which need a 3-hop neighbourhood that was never sampled. So each batch keeps only dst_local < batch_size. Since every requested cell is a seed exactly once across the loader, every edge is emitted exactly once and the matrix is complete. A nnz == accumulated_entries check guards this: coo -> csr silently sums duplicates, so a seed-masking or n_id-remapping bug would inflate values rather than raise.

Extraction defaults to num_neighbors = [-1] * n_layers. A finite fan-out below the graph's max in-degree raises ValueError (naming the degree) unless allow_sampling=True — sampled coefficients still sum to 1 per row and therefore look correct.

Changes

• _spatial_encoder.py — return_attention_weights: bool = False on GCNLayers.forward and GraphEncoder.forward. When False, code path and return signature are identical to before; the encoder appends attentions as one trailing element only when True, leaving both return_dist shapes untouched. Non-GAT raises NotImplementedError.
• _cellina_gcn_module.py — new get_attention(). inference() is untouched; the spatial-encoder input construction both need is factored into a shared _build_spatial_input() helper so the two cannot drift.
• _edge_data_splitter.py — batches additionally expose n_id (PyG local -> global map for every sampled node, vs node_indices which is seeds only). create_inference_loader accepts a per-call num_neighbors override, so extraction forces the exact neighbourhood without mutating the shared cached splitter.
• _cellina_gcn_model.py — get_attention_weights().

batch_size=512 is a deliberate default: with [-1, -1] on a k=20 graph a batch expands to ~1+20+400 nodes per seed before dedup and the loader densifies features.

Tests

tests/test_cellina_gcn.py, all CPU/synthetic, built on the existing fixtures:

1. test_attention_shape_and_pattern — n_layers entries, (n_obs, n_obs), pattern within the graph; also asserts key_added=None leaves adata.obsp untouched and key_added="…" adds one entry per layer.
2. test_attention_orientation — strictly asymmetric graph; pattern must equal the connectivity transpose exactly.
3. test_attention_rows_sum_to_one — rows sum to 1 (non-isolated) / 0 (isolated), every layer.
4. test_attention_batching_invariance — batch_size=8 vs batch_size=n_obs elementwise allclose, on a condition_on_intrinsic=True model. Also contrasts give_mean=True reproducibility against give_mean=False drift.
5. test_attention_raises_for_non_gat — NotImplementedError at both model and encoder level.
6. test_attention_raises_on_sampling — ValueError on a finite fan-out below max in-degree; passes with allow_sampling=True; wrong-length fan-out also rejected.
7. test_attention_subset_indices — subset keeps full shape and matches the all-cells result on selected rows.

$ pytest tests/ -q
57 passed, 45 warnings in 9.11s

Full suite green on both GPU and CPU (CUDA_VISIBLE_DEVICES=""), including all 50 pre-existing tests — the default-off kwarg disturbs nothing.

Read-only smoke test — crc_210 GAT checkpoint

epoch=19-step=8322-vae_loss_validation=306.0075378417969, adata rebuilt exactly as application/02_microenv_gat.ipynb does (raw crc_210.h5ad, preprocess_crc, spatial_neighbors(bandwidth=100/0.12028, max_neighbours=20), drop-NaN on typ_clean). 5000 cells via indices, batch_size=256, single GPU.

data + graph ready in 44s: (552241, 2000)
n_obs = 552241 | n_edges = 11044820 | max in-degree = 39
loaded checkpoint; condition_on_intrinsic = False | convolution_type = gat | n_layers = 2

=== extraction: 6.3s wall | peak VRAM 0.74 GB ===
layers returned: [0, 1]
connectivity graph symmetric? False
layer 0: shape=(552241, 552241) nnz=99821 dtype=float32
  row sums on selected non-isolated cells: min=1.000000 max=1.000000 (n=5000)
  row sums on unselected cells: max=0
  nonzeros outside spatial_connectivities.T: 0
  attention asymmetry: max|M - M.T| = 1
layer 1: shape=(552241, 552241) nnz=99821 dtype=float32
  row sums on selected non-isolated cells: min=1.000000 max=1.000000 (n=5000)
  row sums on unselected cells: max=0
  nonzeros outside spatial_connectivities.T: 0
  attention asymmetry: max|M - M.T| = 1
determinism layer 0: max |call1 - call2| = 0
determinism layer 1: max |call1 - call2| = 1.04e-07
adata.obsp unchanged (key_added=None): True

nnz = 99821 for 5000 seeds is ~19.96 in-edges/cell, as expected for GRAPH_K=20. The 1.04e-07 residual is CUDA scatter float non-associativity, not sampling — this checkpoint has condition_on_intrinsic=False, so give_mean is inert for it; the give_mean path is covered by the synthetic condition_on_intrinsic=True test instead. Note the checkpoint's own num_neighbors=[20, 20] is below this graph's max in-degree of 39, so extraction defaulting to [-1, -1] (rather than inheriting the model's fan-out) is what makes these coefficients exact.

Nothing was written: no adata mutation, no files touched under application/.

🤖 Generated with Claude Code

https://claude.ai/code/session_01SLYXz8HRt6GUjyLjsdWKNm

---

Independent audit

An independent audit checked this implementation against the PyTorch Geometric API. No bugs were found in the extraction logic.

Strongest evidence:

• Orientation (the risk: a transposed result still yields row sums of 1.0 on a symmetric graph, so row sums prove nothing). Confirmed three ways — hand-computed GATv2 alpha on a strictly asymmetric 5-node graph matched .coo() = (dst, src, alpha) to 1e-6 while the transposed hypothesis matched nothing; MessagePassing._collect maps edge_index_i = row for SparseTensor; and on an asymmetric fixture 484/498 entries fall outside spatial_connectivities and 0 fall outside its transpose.
• End-to-end: on a strictly asymmetric 200-cell graph, batched extraction at batch_size=17 reproduced a full-graph reference to 1.2e-07 with an identical sparsity pattern — validating orientation, n_id remapping, seed masking and completeness at once.
• Mutation testing: 8/8 mutants killed (transpose output, drop seed mask, drop/swap n_id map, off-by-one mask, mask on src, ignore give_mean, transpose the SparseTensor construction). test_attention_orientation is load-bearing — without it four of these survive every other test.
• _build_spatial_input non-regression: both commits extracted via git archive, identical seeds and init, 40/40 inference() output tensors bitwise identical across four configs.

Fixed in this PR as a result

• get_attention_weights now restores the module's prior train/eval mode (try/finally).
• Boolean masks in indices raise a named error rather than being reported as duplicated seeds.
• Two docstring claims corrected — see below.

Pre-existing issues found, deliberately NOT fixed here

Both are reachable via get_latent_representation() on rebuttal and belong in their own PR:

1. adata= is silently ignored by _make_data_loader: it is validated and then discarded, with the loader always rebuilt from the registered training data. Passing an adata with empty connectivities and all-zero X returns the training-adata result bitwise. The docstring here now says so explicitly.
2. A batch whose sampled subgraph has zero edges aborts the process with SIGFPE inside GATv2Conv (GAT-specific; GCNConv handles the same empty SparseTensor fine). Repro needs no cellina code:

   adj = SparseTensor(row=torch.zeros(0, dtype=torch.long),
                      col=torch.zeros(0, dtype=torch.long), sparse_sizes=(6, 6))
   GATv2Conv(8, 4, add_self_loops=False, concat=False).eval()(x, adj)  # -> SIGFPE

   This makes the documented isolated-cell behaviour unreachable in that corner case; the docstring now carries the caveat. Zero rows are produced correctly whenever the crash is not triggered.

Separately worth noting for whoever picks those up: get_latent_representation() never calls module.eval(), so GCNLayers' nn.Dropout(p=0.1) is live there. Existing results are unaffected, because scvi.model.base.BaseModelClass.load calls model.module.eval() and analyses load from a checkpoint — the exposure is the train-then-immediately-infer path.

Not verified by the audit

CPU only (the < 1e-5 reproducibility threshold in test_attention_batching_invariance is CPU-calibrated and may be tight on CUDA); the crc_210 smoke-test numbers were not re-run; and nothing pins the PyG version, though a changed SparseTensor return contract would raise AttributeError on .coo() rather than fail silently.

[codecov]

Welcome to Codecov 🎉

Once you merge this PR into your default branch, you're all set! Codecov will compare coverage reports and display results in all future pull requests.

Thanks for integrating Codecov - We've got you covered ☂️