torchcell.models.equivariant_cell_graph_transformer.GlobalHead

class GlobalHead(hidden_dim: int, output_dim: int, use_gene_pool: bool = True, dropout: float = 0.1, param_dim: int = 1)

Bases: Module

Whole-cell readout head for global phenotypes.

Maps the CLS token (whole-cell representation), optionally concatenated with a mean pool over the equivariant perturbed gene embeddings, to a configurable output vector. Used for CalMorph morphology (501-D) AND scalar VisualScore (e.g. beta-carotene, output_dim=1).

graph_level == “global” selects this head.

forward(h_CLS: Tensor, H_genes_pert: Tensor) → Tensor

Forward pass of the global head.

Parameters:

• h_CLS – [d] whole-cell CLS representation.

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

Returns:

[batch_size, output_dim] global phenotype predictions, or [batch_size, output_dim, param_dim] when param_dim > 1 (distributional).

Return type:

predictions