diff --git a/src/cellina/_cellina_gcn_model.py b/src/cellina/_cellina_gcn_model.py index dc53016..5b6cbef 100644 --- a/src/cellina/_cellina_gcn_model.py +++ b/src/cellina/_cellina_gcn_model.py @@ -3,6 +3,7 @@ from typing import List, Optional, Union import numpy as np +import scipy.sparse as sp import torch from anndata import AnnData from scvi import REGISTRY_KEYS @@ -164,7 +165,7 @@ def _get_cached_splitter(self, batch_size): return self._cached_splitter def _make_data_loader(self, adata=None, indices=None, batch_size=None, shuffle=False, - x_spatial_layer=None): + x_spatial_layer=None, num_neighbors=None): adata = self._validate_anndata(adata) if adata is not None else self.adata if batch_size is None: @@ -187,6 +188,7 @@ def _make_data_loader(self, adata=None, indices=None, batch_size=None, shuffle=F batch_size=batch_size, shuffle=shuffle, x_spatial_override=override, + num_neighbors=num_neighbors, ) def _make_counterfactual_loader( @@ -609,6 +611,234 @@ def get_latent_representation( return torch.cat(latent).numpy() + @torch.inference_mode() + def get_attention_weights( + self, + adata: Optional[AnnData] = None, + indices: Optional[list] = None, + batch_size: int = 512, + num_neighbors: Optional[List[int]] = None, + give_mean: bool = True, + key_added: Optional[str] = None, + allow_sampling: bool = False, + ) -> dict: + """ + Extract the learned GATv2 attention coefficients as cell-by-cell matrices. + + Requires ``convolution_type="gat"``. One matrix is returned per GAT layer of the + spatial encoder; they are **not** averaged, because each layer attends over a + different representation. + + Orientation + ----------- + ``A[i, j] = alpha(j -> i)``, i.e. **row = destination**: the weight cell ``i`` + places on its spatial neighbour ``j``. This follows the ``adj_t`` convention used + internally by the graph encoder. Consequences: + + - Rows sum to 1 for every cell that has at least one in-neighbour, and to 0 for + isolated cells (and for cells not covered by ``indices``). Caveat: a mini-batch + whose sampled subgraph happens to contain *no* edges at all aborts the process + inside :class:`~torch_geometric.nn.GATv2Conv` (SIGFPE). This is a pre-existing + upstream limitation, reachable the same way via + :meth:`get_latent_representation`, and in practice needs ``indices`` to select + only isolated cells. + - The matrix is **asymmetric** even when the connectivity graph is symmetric: + ``alpha(j -> i)`` and ``alpha(i -> j)`` are normalised over different + neighbourhoods. + - Its sparsity pattern is a subset of + ``adata.obsp["spatial_connectivities"].T`` (self-loops excluded), *not* of the + connectivity matrix itself. The data loader maps a nonzero + ``spatial_connectivities[i, j]`` to a message ``i -> j``, so it is cell ``j`` + that attends over cell ``i``. For a symmetric connectivity graph the two + patterns coincide; for the asymmetric kNN graphs produced by + :func:`~cellina._spatial_utils.spatial_neighbors` they do not. + - The shape is always ``(n_obs, n_obs)``, even when ``indices`` selects a subset; + unselected rows are all-zero. + + Parameters + ---------- + adata + Must be the AnnData the model was set up with; defaults to it. A *different* + object is not honoured: the loader is always rebuilt from the registered + training data, so passing another adata only changes the output shape and + where ``key_added`` writes, not which graph or counts are used. + indices + Integer positions of the cells to extract attention for. Defaults to all + cells. Must be unique; boolean masks are not accepted. + batch_size + Seed cells per mini-batch. The default of 512 is deliberately below the usual + inference default: with an exact ``[-1, -1]`` fan-out on a k=20 graph each + seed pulls in ~1 + 20 + 400 nodes before dedup and the loader densifies their + features, so a larger batch risks running out of memory. + num_neighbors + Fan-out per hop. Defaults to ``[-1] * n_layers`` (the exact, full + neighbourhood), which is what makes the coefficients equal to those a + full-graph forward pass would produce. Its length must equal the encoder's + number of layers, since layer ``i``'s attention at a seed needs an + ``i + 1``-hop neighbourhood. + give_mean + If True (default), the posterior mean ``qzm`` replaces the sampled ``z`` in + the GCN input when the model was built with ``condition_on_intrinsic=True``. + This makes extraction deterministic; ``give_mean=False`` reproduces the + sampled behaviour of :meth:`get_latent_representation` and will differ + between calls. + key_added + If given, additionally store each layer's matrix in + ``adata.obsp[f"{key_added}_l{i}"]``. ``None`` (default) never mutates + ``adata``. + allow_sampling + Permit a finite ``num_neighbors`` below the graph's maximum in-degree. Off by + default, because sampling silently yields approximate coefficients whose rows + still sum to 1 and therefore look correct. + + Returns + ------- + ``dict[int, scipy.sparse.csr_matrix]`` keyed by layer index (0-based), each of + shape ``(n_obs, n_obs)`` and dtype ``float32``. + + Notes + ----- + Under neighbour sampling only edges pointing *into* a seed node carry exact + attention (a 1-hop node's layer-1 coefficients would need a hop that was never + sampled), so non-seed destinations are dropped. Every requested cell is a seed + exactly once across the loader, so each edge is still emitted exactly once and + the result is complete. + """ + self._check_if_trained(warn=False) + adata = self._validate_anndata(adata) + + encoder = self.module.s_encoder.encoder + if encoder.convolution_type != "gat": + raise NotImplementedError( + "get_attention_weights() requires convolution_type='gat', but this model " + f"was built with convolution_type='{encoder.convolution_type}'. Other " + "graph convolutions do not learn per-edge attention coefficients." + ) + + n_layers = len(encoder.gcn_layers) + + if num_neighbors is None: + num_neighbors = [-1] * n_layers + num_neighbors = list(num_neighbors) + if len(num_neighbors) != n_layers: + raise ValueError( + f"num_neighbors must have one entry per GAT layer (n_layers={n_layers}), " + f"got length {len(num_neighbors)}: {num_neighbors}. Layer i's attention at " + "a seed cell requires an (i + 1)-hop neighbourhood." + ) + + splitter = self._get_cached_splitter(batch_size) + graph_edge_index = splitter.pyg_data.edge_index.numpy() + in_degrees = np.bincount(graph_edge_index[1], minlength=adata.n_obs) + max_in_degree = int(in_degrees.max()) if in_degrees.size else 0 + + finite_fanouts = [f for f in num_neighbors if f >= 0] + if finite_fanouts and min(finite_fanouts) < max_in_degree and not allow_sampling: + raise ValueError( + f"num_neighbors={num_neighbors} samples fewer neighbours than the graph's " + f"maximum in-degree ({max_in_degree}), so the returned coefficients would " + "be a renormalised approximation rather than the true attention. Use " + "num_neighbors=None for the exact full neighbourhood, or pass " + "allow_sampling=True to accept the approximation." + ) + + if indices is None: + indices = np.arange(adata.n_obs) + indices = np.asarray(indices) + if indices.dtype == bool: + raise ValueError( + "indices must be integer cell positions, not a boolean mask; convert " + "with np.flatnonzero(mask). A mask would otherwise be read as the " + "positions 0/1 and rejected as duplicated." + ) + if np.unique(indices).size != indices.size: + raise ValueError( + "indices must be unique; duplicated seed cells would silently sum their " + "attention coefficients in the returned matrix." + ) + + scdl = self._make_data_loader( + adata=adata, + indices=indices, + batch_size=batch_size, + num_neighbors=num_neighbors, + ) + + # Restore the caller's mode afterwards: `get_latent_representation` does not set + # eval() itself, so leaving the module in eval would silently change its output + # (GCNLayers carries an nn.Dropout) depending on whether this ran first. + was_training = self.module.training + self.module.eval() + + dst_parts = [[] for _ in range(n_layers)] + src_parts = [[] for _ in range(n_layers)] + val_parts = [[] for _ in range(n_layers)] + + try: + for tensors in scdl: + node_batch = tensors["node_batch"] + n_id = node_batch["n_id"].cpu().numpy() + seed_size = int(node_batch["batch_size"]) + + attentions = self.module.get_attention( + x=node_batch["X"], + batch_index=node_batch["batch_label"], + edge_index=node_batch["edge_index"], + batch_size=seed_size, + x_spatial=node_batch.get("x_spatial"), + give_mean=give_mean, + ) + if len(attentions) != n_layers: + raise RuntimeError( + f"Expected {n_layers} attention tensors, got {len(attentions)}." + ) + + for layer, att in enumerate(attentions): + dst_local, src_local, alpha = att.coo() + # heads is hardcoded to 1 in _make_conv_layer, but stay defensive. + if alpha.dim() > 1: + alpha = alpha.mean(dim=-1) if alpha.size(-1) > 1 else alpha.squeeze(-1) + # Seed nodes are the first `seed_size` local ids; only edges into them + # have a fully materialised receptive field at this depth. + keep = dst_local < seed_size + dst_parts[layer].append(n_id[dst_local[keep].cpu().numpy()]) + src_parts[layer].append(n_id[src_local[keep].cpu().numpy()]) + val_parts[layer].append(alpha[keep].float().cpu().numpy()) + finally: + self.module.train(was_training) + + n_obs = adata.n_obs + result = {} + for layer in range(n_layers): + if val_parts[layer]: + dst = np.concatenate(dst_parts[layer]) + src = np.concatenate(src_parts[layer]) + vals = np.concatenate(val_parts[layer]).astype(np.float32, copy=False) + else: + dst = np.zeros(0, dtype=np.int64) + src = np.zeros(0, dtype=np.int64) + vals = np.zeros(0, dtype=np.float32) + + mat = sp.coo_matrix( + (vals, (dst, src)), shape=(n_obs, n_obs), dtype=np.float32 + ).tocsr() + # coo -> csr *sums* duplicates; a mismatch here would mean an edge was + # emitted twice (broken seed masking or duplicated seeds) and silently + # inflated rather than raising. + if mat.nnz != vals.size: + raise RuntimeError( + f"Layer {layer}: built {vals.size} attention entries but the matrix " + f"has {mat.nnz} nonzeros; duplicate (destination, source) pairs were " + "summed. This indicates a seed-masking or index-mapping bug." + ) + result[layer] = mat + + if key_added is not None: + for layer, mat in result.items(): + adata.obsp[f"{key_added}_l{layer}"] = mat + + return result + def get_marginal_ll( self, adata: Optional[AnnData] = None, diff --git a/src/cellina/_cellina_gcn_module.py b/src/cellina/_cellina_gcn_module.py index 8568a8b..188af27 100644 --- a/src/cellina/_cellina_gcn_module.py +++ b/src/cellina/_cellina_gcn_module.py @@ -225,6 +225,28 @@ def _get_generative_input(self, tensors, inference_outputs): "batch_index": batch_index, } + def _build_spatial_input(self, x_, x_spatial, z_like): + """Assemble the GCN input from log-counts and (optionally) the intrinsic latent. + + Single source of truth shared by :meth:`inference` and :meth:`get_attention` so + the two cannot drift. + + Parameters + ---------- + x_ + ``log(1 + x)`` of the count matrix; used as the spatial features when + ``x_spatial`` is not provided. + x_spatial + Optional raw alternative spatial features (log1p-transformed here). + z_like + Intrinsic latent to condition on when ``condition_on_intrinsic`` is True. + Always detached before concatenation. + """ + x_spatial_ = torch.log(1 + x_spatial) if x_spatial is not None else x_ + if self.condition_on_intrinsic: + return torch.cat([x_spatial_, z_like.detach()], dim=-1) + return x_spatial_ + @auto_move_data def inference( self, @@ -239,12 +261,7 @@ def inference( qzm, qzv, z = self.z_encoder(x_, batch_index) - x_spatial_ = torch.log(1 + x_spatial) if x_spatial is not None else x_ - - if self.condition_on_intrinsic: - spatial_input = torch.cat([x_spatial_, z.detach()], dim=-1) - else: - spatial_input = x_spatial_ + spatial_input = self._build_spatial_input(x_, x_spatial, z) qsm, qsv, s, neighbor_means = self.s_encoder( spatial_input, edge_index, batch_index, @@ -294,6 +311,52 @@ def inference( return outputs + @auto_move_data + def get_attention( + self, + x, + batch_index, + edge_index, + batch_size, + x_spatial=None, + give_mean: bool = True, + ): + """Return the raw per-layer GAT attention coefficients for one graph batch. + + Mirrors :meth:`inference`'s input preparation exactly (``log(1 + x)``, the + ``x_spatial`` fallback and the ``condition_on_intrinsic`` concatenation) but runs + only the spatial encoder and returns attention instead of latents. + + Parameters + ---------- + give_mean + If True (default), feed the posterior mean ``qzm`` rather than the + reparameterised sample ``z`` into the GCN input under + ``condition_on_intrinsic=True``. ``module.eval()`` disables dropout but not + the reparameterisation draw, so ``give_mean=False`` makes attention differ + between two otherwise identical calls. + + Returns + ------- + ``list[SparseTensor]``, one per GAT layer, over the *sampled* nodes of this batch + (``row = destination``, ``col = source``). Only edges whose destination is one of + the first ``batch_size`` (seed) nodes carry exact coefficients. + """ + x_ = torch.log(1 + x) + + qzm, _, z = self.z_encoder(x_, batch_index) + z_like = qzm if give_mean else z + + spatial_input = self._build_spatial_input(x_, x_spatial, z_like) + + *_, attentions = self.s_encoder( + spatial_input, edge_index, batch_index, + batch_size=batch_size, + return_neighbor_means=False, + return_attention_weights=True, + ) + return attentions + @auto_move_data def generative(self, shifted, library, batch_index): px_scale, _, px_rate, px_dropout = self.decoder("gene", shifted, library, batch_index) diff --git a/src/cellina/_edge_data_splitter.py b/src/cellina/_edge_data_splitter.py index 28deddc..fd16273 100644 --- a/src/cellina/_edge_data_splitter.py +++ b/src/cellina/_edge_data_splitter.py @@ -42,6 +42,10 @@ def _node_batch_to_dict(self, node_batch): 'X': _gather_dense_rows(self.x_sparse, nid), 'edge_index': node_batch.edge_index, 'node_indices': node_batch.input_id, + # Local -> global node id map for *every* sampled node (seeds first), as + # opposed to 'node_indices' which covers the seed nodes only. Needed to map + # per-edge quantities (e.g. GAT attention) back onto adata's cell indexing. + 'n_id': node_batch.n_id, 'batch_size': node_batch.batch_size, 'batch_label': node_batch.batch_labels, REGISTRY_KEYS.LABELS_KEY: node_batch.labels, @@ -169,10 +173,14 @@ def load_spatial_store(self, layer): ) return x_sp.tocsr() - def _make_neighbor_loader(self, node_indices, batch_size, shuffle, drop_last): + def _make_neighbor_loader(self, node_indices, batch_size, shuffle, drop_last, + num_neighbors=None): + # ``num_neighbors=None`` uses this splitter's configured fan-out. An explicit + # value overrides it for this loader only -- the splitter is cached and shared, + # so it must never be mutated in place. return NeighborLoader( self.pyg_data, - num_neighbors=self.num_neighbors, + num_neighbors=self.num_neighbors if num_neighbors is None else num_neighbors, input_nodes=torch.tensor(node_indices, dtype=torch.long), batch_size=batch_size, shuffle=shuffle, @@ -198,15 +206,20 @@ def test_dataloader(self): return None def create_inference_loader(self, indices, batch_size=None, shuffle=False, - x_spatial_override=None): + x_spatial_override=None, num_neighbors=None): """Create a node-only loader for inference. ``x_spatial_override`` lets callers reuse this splitter's cached graph and base X store while swapping in a different spatial feature store (e.g. a perturbation ``cf_layer``), avoiding a full splitter rebuild. + + ``num_neighbors`` overrides the splitter's fan-out for this loader only (e.g. the + exact ``[-1] * n_layers`` neighbourhood required by attention extraction); the + cached splitter itself is left untouched. """ batch_size = batch_size or self.batch_size - node_loader = self._make_neighbor_loader(indices, batch_size, shuffle, drop_last=False) + node_loader = self._make_neighbor_loader(indices, batch_size, shuffle, drop_last=False, + num_neighbors=num_neighbors) spatial_store = ( x_spatial_override if x_spatial_override is not None else self._x_spatial_sparse ) diff --git a/src/cellina/_spatial_encoder.py b/src/cellina/_spatial_encoder.py index 618ebd0..3cf59ce 100644 --- a/src/cellina/_spatial_encoder.py +++ b/src/cellina/_spatial_encoder.py @@ -143,7 +143,31 @@ def forward( edge_index: torch.Tensor, *cat_list: int, cont: torch.Tensor | None = None, + return_attention_weights: bool = False, ): + """Run the graph convolutions. + + Parameters + ---------- + return_attention_weights + If ``True`` (``convolution_type="gat"`` only), also return the per-layer + attention coefficients. The default (``False``) leaves the return value and + the computation path exactly as they are without this argument. + + Returns + ------- + ``x`` when ``return_attention_weights=False`` (default), else + ``(x, attentions)`` where ``attentions`` is a ``list[SparseTensor]`` of length + ``n_layers``. Each ``SparseTensor`` holds the layer's attention coefficients with + the ``adj_t`` convention used here: ``row = destination``, ``col = source``, so + ``.coo()`` yields ``(dst, src, alpha)`` and ``alpha`` sums to 1 per destination. + """ + if return_attention_weights and self.convolution_type != "gat": + raise NotImplementedError( + "return_attention_weights=True is only supported for " + f"convolution_type='gat', got '{self.convolution_type}'." + ) + one_hot_cat_list = [] cont_list = [cont] if cont is not None else [] cat_list = cat_list or [] @@ -172,8 +196,17 @@ def forward( else: adj = edge_index + attentions = [] if return_attention_weights else None + for i, gcn_layer in enumerate(self.gcn_layers): - x = gcn_layer(x, adj) + if return_attention_weights: + # PyG >=2.4: with a SparseTensor input GATv2Conv returns + # (out, adj.set_value(alpha, layout='coo')) -- a SparseTensor, not the + # (edge_index, alpha) tuple returned for dense edge_index inputs. + x, att = gcn_layer(x, adj, return_attention_weights=True) + attentions.append(att) + else: + x = gcn_layer(x, adj) if self.cov_layers[i] is not None and cov_list: cov = torch.cat(cov_list, dim=-1) x = x + self.cov_layers[i](cov.float()) @@ -182,6 +215,8 @@ def forward( if self.dropout is not None: x = self.dropout(x) + if return_attention_weights: + return x, attentions return x @@ -263,8 +298,33 @@ def forward( *cat_list: int, batch_size: int, return_neighbor_means: bool = False, + return_attention_weights: bool = False, ): - q = self.encoder(x, edge_index, *cat_list) + """Encode node features into the spatial latent space. + + Parameters + ---------- + return_attention_weights + If ``True`` (``convolution_type="gat"`` only), append the per-layer GAT + attention coefficients as one extra trailing element to whichever return + tuple this encoder normally produces. The default (``False``) leaves both + existing return shapes (``return_dist`` True/False) untouched. + + Returns + ------- + ``(dist, latent, neighbor_means)`` if ``self.return_dist`` else + ``(q_m, q_v, latent, neighbor_means)``; with ``return_attention_weights=True`` a + trailing ``attentions`` element (``list[SparseTensor]``, one per GAT layer, + ``row = destination``) is appended. Note the attentions cover *all* sampled nodes, + not only the first ``batch_size`` seed nodes -- only rows whose destination is a + seed node carry exact coefficients under neighbour sampling. + """ + if return_attention_weights: + q, attentions = self.encoder( + x, edge_index, *cat_list, return_attention_weights=True + ) + else: + q = self.encoder(x, edge_index, *cat_list) q_seed = q[:batch_size] if self.bn is not None: q_seed = self.bn(q_seed) @@ -277,6 +337,11 @@ def forward( dist = Normal(q_m, q_v.sqrt()) latent = self.z_transformation(dist.rsample()) - if self.return_dist: - return dist, latent, neighbor_means - return q_m, q_v, latent, neighbor_means + out = ( + (dist, latent, neighbor_means) + if self.return_dist + else (q_m, q_v, latent, neighbor_means) + ) + if return_attention_weights: + return (*out, attentions) + return out diff --git a/tests/test_cellina_gcn.py b/tests/test_cellina_gcn.py index 7f2bcb2..ff753ad 100644 --- a/tests/test_cellina_gcn.py +++ b/tests/test_cellina_gcn.py @@ -700,6 +700,269 @@ def trained_model(adata_with_spatial): return model, adata_with_spatial +# ── GAT attention extraction ────────────────────────────────────────────────── + +@pytest.fixture +def trained_gat_model(adata_with_spatial): + """GAT model with condition_on_intrinsic=True (the give_mean-sensitive path).""" + CellinaGCN.setup_anndata( + adata_with_spatial, + batch_key="batch", + labels_key="cell_labels", + domains_key="domain", + spatial_connectivities_key="spatial_connectivities", + ) + model = CellinaGCN( + adata_with_spatial, + n_latent=5, + discriminator_lambda=0.0, + classifier_lambda=0.0, + convolution_type="gat", + condition_on_intrinsic=True, + ) + model.train(max_epochs=1, train_size=0.8) + return model, adata_with_spatial + + +def test_attention_shape_and_pattern(trained_gat_model): + """One matrix per GAT layer, (n_obs, n_obs), pattern within the connectivity graph.""" + model, adata = trained_gat_model + obsp_before = set(adata.obsp.keys()) + + att = model.get_attention_weights(batch_size=64) + + assert sorted(att) == list(range(model.n_layers)) + # Cell i attends over j only where the loader created a message j -> i, i.e. where + # spatial_connectivities[j, i] != 0 (see test_attention_orientation). + conn_t = (adata.obsp["spatial_connectivities"] != 0).T.tocsr() + for layer, mat in att.items(): + assert sp.issparse(mat) + assert mat.shape == (adata.n_obs, adata.n_obs) + assert mat.dtype == np.float32 + outside = (mat != 0) > conn_t + assert outside.nnz == 0, f"layer {layer} attends outside the spatial graph" + + # key_added=None must not touch adata. + assert set(adata.obsp.keys()) == obsp_before + + # key_added writes one obsp entry per layer. + model.get_attention_weights(batch_size=64, key_added="gat_att") + for layer in att: + assert f"gat_att_l{layer}" in adata.obsp + assert adata.obsp[f"gat_att_l{layer}"].shape == (adata.n_obs, adata.n_obs) + + +def test_attention_orientation(adata_with_spatial): + """On a strictly asymmetric graph, pin down which way round the matrix is. + + ``adata_with_spatial`` is symmetric, so it cannot distinguish ``A`` from ``A.T``; a + transposed implementation would pass every other test in this section. + """ + n_obs = adata_with_spatial.n_obs + rows, cols = [], [] + for i in range(n_obs): + for k in (1, 2, 3): + rows.append(i) + cols.append((i + k) % n_obs) + conn = sp.csr_matrix((np.ones(len(rows)), (rows, cols)), shape=(n_obs, n_obs)) + assert abs(conn - conn.T).nnz > 0, "graph must be asymmetric for this test to bite" + adata_with_spatial.obsp["spatial_connectivities"] = conn + + CellinaGCN.setup_anndata( + adata_with_spatial, + batch_key="batch", + labels_key="cell_labels", + domains_key="domain", + spatial_connectivities_key="spatial_connectivities", + ) + model = CellinaGCN( + adata_with_spatial, n_latent=5, discriminator_lambda=0.0, classifier_lambda=0.0 + ) + model.train(max_epochs=1, train_size=0.8) + + att = model.get_attention_weights(batch_size=64) + conn_bool = (conn != 0) + for layer, mat in att.items(): + pattern = (mat != 0) + # spatial_connectivities[i, j] != 0 becomes a message i -> j, so j attends over i + # and the attention pattern is the *transpose* of the connectivity pattern. + assert (pattern != conn_bool.T.tocsr()).nnz == 0, ( + f"layer {layer}: attention pattern is not the transpose of the connectivity " + "graph; the destination/source orientation has flipped" + ) + np.testing.assert_allclose( + np.asarray(mat.sum(axis=1)).ravel(), 1.0, rtol=1e-4, atol=1e-4 + ) + + +def test_attention_rows_sum_to_one(trained_gat_model): + """Row = destination, so each non-isolated cell's incoming weights sum to 1.""" + model, adata = trained_gat_model + conn = adata.obsp["spatial_connectivities"] + in_degree = np.asarray((conn != 0).sum(axis=0)).ravel() + non_isolated = in_degree > 0 + assert non_isolated.any() + + att = model.get_attention_weights(batch_size=64) + for layer, mat in att.items(): + row_sums = np.asarray(mat.sum(axis=1)).ravel() + np.testing.assert_allclose( + row_sums[non_isolated], 1.0, rtol=1e-4, atol=1e-4, + err_msg=f"layer {layer} rows do not sum to 1", + ) + if (~non_isolated).any(): + np.testing.assert_allclose(row_sums[~non_isolated], 0.0, atol=1e-6) + + +def test_attention_batching_invariance(trained_gat_model): + """Extraction must not depend on batch size (catches n_id / seed-slice bugs).""" + model, adata = trained_gat_model + + small = model.get_attention_weights(batch_size=8) + full = model.get_attention_weights(batch_size=adata.n_obs) + + assert sorted(small) == sorted(full) + for layer in small: + np.testing.assert_allclose( + small[layer].toarray(), full[layer].toarray(), rtol=1e-4, atol=1e-6, + err_msg=f"layer {layer} differs between batch sizes", + ) + + # give_mean=True (default) makes repeated calls reproducible even though + # condition_on_intrinsic=True feeds the intrinsic latent into the GCN input; only + # float non-associativity (GPU scatter order) may differ. give_mean=False samples z + # and must therefore move the coefficients by orders of magnitude more. + again = model.get_attention_weights(batch_size=8) + sampled_a = model.get_attention_weights(batch_size=8, give_mean=False) + sampled_b = model.get_attention_weights(batch_size=8, give_mean=False) + for layer in small: + deterministic_drift = np.abs( + small[layer].toarray() - again[layer].toarray() + ).max() + sampled_drift = np.abs( + sampled_a[layer].toarray() - sampled_b[layer].toarray() + ).max() + assert deterministic_drift < 1e-5, ( + f"layer {layer} is not reproducible with give_mean=True " + f"(max diff {deterministic_drift})" + ) + assert sampled_drift > 100 * max(deterministic_drift, 1e-9), ( + f"layer {layer}: give_mean=False was expected to vary between calls, but " + f"moved only {sampled_drift}" + ) + + +def test_attention_raises_for_non_gat(adata_with_spatial): + """Only GATv2 learns per-edge attention; anything else must refuse loudly.""" + CellinaGCN.setup_anndata( + adata_with_spatial, + batch_key="batch", + labels_key="cell_labels", + domains_key="domain", + spatial_connectivities_key="spatial_connectivities", + ) + model = CellinaGCN( + adata_with_spatial, + n_latent=5, + discriminator_lambda=0.0, + classifier_lambda=0.0, + convolution_type="gcn", + ) + model.train(max_epochs=1, train_size=0.8) + + with pytest.raises(NotImplementedError, match="gcn"): + model.get_attention_weights(batch_size=64) + + # The encoder-level flag guards the same way. + from cellina._spatial_encoder import GCNLayers + layers = GCNLayers(n_in=4, n_out=4, n_layers=1, convolution_type="gcn") + with pytest.raises(NotImplementedError, match="gcn"): + layers( + torch.zeros(3, 4), + torch.tensor([[0, 1], [1, 2]], dtype=torch.long), + return_attention_weights=True, + ) + + +def test_attention_raises_on_sampling(trained_gat_model): + """A finite fan-out below the max in-degree is approximate: opt in explicitly.""" + model, adata = trained_gat_model + conn = adata.obsp["spatial_connectivities"] + max_in_degree = int(np.asarray((conn != 0).sum(axis=0)).max()) + assert max_in_degree > 1 + + fanout = [1] * model.n_layers + with pytest.raises(ValueError, match=str(max_in_degree)): + model.get_attention_weights(batch_size=64, num_neighbors=fanout) + + att = model.get_attention_weights( + batch_size=64, num_neighbors=fanout, allow_sampling=True + ) + assert sorted(att) == list(range(model.n_layers)) + for mat in att.values(): + assert mat.shape == (adata.n_obs, adata.n_obs) + + # Wrong-length fan-out is an error regardless of allow_sampling. + with pytest.raises(ValueError, match="one entry per GAT layer"): + model.get_attention_weights( + batch_size=64, num_neighbors=[-1] * (model.n_layers + 1), allow_sampling=True + ) + + +def test_attention_subset_indices(trained_gat_model): + """A subset of indices keeps the full shape and matches the all-cells result.""" + model, adata = trained_gat_model + subset = np.arange(0, adata.n_obs, 7) + + full = model.get_attention_weights(batch_size=64) + partial = model.get_attention_weights(batch_size=64, indices=subset) + + mask = np.zeros(adata.n_obs, dtype=bool) + mask[subset] = True + for layer, mat in partial.items(): + assert mat.shape == (adata.n_obs, adata.n_obs) + row_sums = np.asarray(mat.sum(axis=1)).ravel() + np.testing.assert_allclose(row_sums[~mask], 0.0, atol=1e-6) + np.testing.assert_allclose( + mat[subset].toarray(), full[layer][subset].toarray(), rtol=1e-4, atol=1e-6, + ) + + +def test_attention_restores_training_mode(trained_gat_model, monkeypatch): + """Extraction must not leave the module in eval(): get_latent_representation does + not set the mode itself, so a leaked eval() silently changes its dropout behaviour.""" + model, _ = trained_gat_model + + for mode in (True, False): + model.module.train(mode) + model.get_attention_weights(batch_size=64) + assert model.module.training is mode + + # Restored even when extraction raises inside the batch loop. + def _boom(*args, **kwargs): + raise RuntimeError("boom") + + model.module.train(True) + monkeypatch.setattr(model.module, "get_attention", _boom) + with pytest.raises(RuntimeError, match="boom"): + model.get_attention_weights(batch_size=64) + assert model.module.training is True + + +def test_attention_rejects_boolean_mask(trained_gat_model): + """A boolean mask must be named as such, not misreported as duplicated indices.""" + model, adata = trained_gat_model + mask = np.zeros(adata.n_obs, dtype=bool) + mask[:10] = True + + with pytest.raises(ValueError, match="boolean mask"): + model.get_attention_weights(batch_size=64, indices=mask) + + # The documented conversion works. + att = model.get_attention_weights(batch_size=64, indices=np.flatnonzero(mask)) + assert sorted(att) == list(range(model.n_layers)) + + def test_make_perturbed_expression(): X = np.arange(1, 10, dtype=float).reshape(3, 3) adata = AnnData(X=X.copy())