GATv2 attention extraction + attention_by_group, tutorial section 3.1

[dbdimitrov]

Summary

Extracts the learned GATv2 attention coefficients from CellinaGCN and summarises them by cell group. The rebuttal-branch design is re-applied to current main from scratch.

API

• CellinaGCN.get_attention_weights(adata=None, indices=None, batch_size=512, num_neighbors=None, key_added=None, allow_sampling=False) -> dict[int, csr_matrix]
  • Returns one (n_obs, n_obs) float32 CSR per GAT layer, keyed by 0-based layer index. Rows not in indices are all-zero.
  • Coefficients are exact by default: the fan-out is [-1] * n_layers, and only edges into seed cells are kept (dst_local < batch_size). Each cell is a seed exactly once, so every edge is written exactly once, and an nnz == entries check guards this.
  • A finite fan-out below the graph's max in-degree raises unless allow_sampling=True. num_neighbors must have one entry per layer, and indices must be unique integers.
  • The module's train/eval mode is restored afterwards. If there are multiple heads, they are averaged (defensive; heads is 1 today).
  • key_added writes adata.obsp[f"{key_added}_l{i}"]. By default adata is never mutated.
• CellinaGCN.attention_by_group(groupby, layer=-1, normalize=True, **kwargs) -> DataFrame
  • For each receiver group and sender group it computes log2(mean attention share / mean neighbour share), averaged over receiver cells that have neighbours.
  • The neighbour share is computed on the same edge set as the attention, so 0 is what proximity alone predicts.
  • normalize=False returns the mean attention share instead, whose rows sum to 1.
  • kwargs are forwarded to get_attention_weights. This reproduces the statistic in the rebuttal notebook application/05_analyse_attention.ipynb.
• CellinaGCNModule.get_attention(...) runs only the spatial encoder, with the same spatial input as inference(): log1p(x_spatial) if given, else log1p(x).
• Plumbing:
  • GCNLayers.forward and GraphEncoder.forward gain return_attention_weights=False. The default path and return signatures are unchanged, and non-GAT convolutions raise NotImplementedError.
  • Loader batches now carry n_id.
  • create_inference_loader / _make_neighbor_loader accept a per-loader num_neighbors override. The cached splitter is never mutated.

Orientation

A[i, j] = alpha(j -> i): the row is the destination (the attending cell) and the column is the source (its neighbour). Rows sum to 1 for cells that have neighbours. The splitter symmetrises the connectivity graph, so the sparsity pattern equals the connectivity pattern minus self-loops. A itself is still asymmetric.

Deliberately dropped relative to the rebuttal implementation

• give_mean / condition_on_intrinsic / _build_spatial_input / z_encoder call in get_attention: main has no condition_on_intrinsic, so the GCN input never depends on z. Extraction is deterministic without these, and the extra encoder pass would only cost time.
• The long "pattern is the transpose of the connectivity" caveat: main symmetrises the graph in the splitter, so the attention pattern equals the connectivity pattern and the caveat no longer applies.
• The inference() rewrite: main's inference() (count/library encoders on seed rows only) is left untouched. get_attention mirrors its spatial input directly.

Tests (tests/test_cellina_gcn.py)

• Shape, dtype and pattern (connectivity minus diagonal), plus key_added behaviour (no mutation by default).
• Orientation on an asymmetric connectivity (pattern == symmetrised graph minus diagonal, rows sum to 1).
• Rows sum to 1, and to 0 for an isolated cell.
• Batching invariance (batch_size 8 vs n_obs).
• Raises for non-GAT, at both the model and the encoder level.
• Raises when sampling below the max in-degree, passes with allow_sampling=True, and rejects a wrong-length num_neighbors.
• Subset indices: full shape is kept, rows match the full run, and duplicate / boolean indices are rejected.
• attention_by_group:
  • shape and labels
  • normalize=False rows sum to 1
  • uniform attention (monkeypatched) gives all-zero log2
  • layer=-1 == layer=n_layers-1

Tutorial

docs/tutorial_gat.ipynb gains section 3.1 What the GAT attends to: one explanatory markdown cell, an enrichment heatmap, and a short read-out. The whole notebook was re-executed end to end on GPU with this branch.

🤖 Generated with Claude Code

[codecov]

Codecov Report

❌ Patch coverage is 99.01961% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 90.73%. Comparing base (c77b614) to head (bc47e7d).

Files with missing lines	Patch %	Lines
src/cellina/_cellina_gcn_model.py	98.73%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main      #40      +/-   ##
==========================================
+ Coverage   90.10%   90.73%   +0.63%     
==========================================
  Files          10       10              
  Lines        1394     1489      +95     
==========================================
+ Hits         1256     1351      +95     
  Misses        138      138              

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.