torchcell.models.equivariant_cell_graph_transformer.PerMetaboliteHead

class PerMetaboliteHead(hidden_dim: int, output_dim: int = 1, dropout: float = 0.1)

Bases: Module

Per-metabolite readout head via GPR pooling of gene embeddings.

Metabolism enters as a REPRESENTATION ANNOTATION (not an attention prior): perturbed gene embeddings are pooled over the genes catalyzing each reaction (gene->gpr->reaction incidence), then pooled over the reactions each metabolite participates in (metabolite<-reaction incidence), producing a per-metabolite vector. Metabolites are NEVER promoted to encoder nodes – the encoder keeps fixed N gene nodes. Used for MetabolitePhenotype (mapped to Yeast9 metabolite ids by the cell_graph metabolite node ordering).

graph_level == “metabolism” selects this head.

TODO(ws7): the pooling here is a mean over the catalyzing genes / participating reactions using the sha-normalized incidence built from cell_graph. Stoichiometry (edge weights on the metabolite<->reaction hyperedge) is ignored for now; a future version can weight the reaction->metabolite pool by |stoichiometric coefficient|.

forward(H_genes_pert: Tensor, gpr_incidence_T: Tensor, mr_incidence: Tensor) → Tensor

Forward pass of the per-metabolite head.

Parameters:

• H_genes_pert – [batch, N, d] equivariant perturbed gene embeddings.

• gpr_incidence_T – sparse [R, N] reaction<-gene incidence (row-normalized so each reaction row is a mean over its catalyzing genes).

• mr_incidence – sparse [M, R] metabolite<-reaction incidence (row-normalized so each metabolite row is a mean over its participating reactions).

Returns:

[batch, M] if output_dim == 1, else [batch, M, output_dim].

Return type:

predictions