torchcell.models.equivariant_cell_graph_transformer.EquivariantPerturbationTransform

class EquivariantPerturbationTransform(hidden_dim: int, num_heads: int = 8, dropout: float = 0.1, residual: str = 'postln', num_layers: int = 1, ffn_mult: int = 4, extra_layer_ffn_mult: int | None = None, hadamard: str = 'off', null_sink: bool = False, null_sink_bias_init: float = -4.0, null_sink_trainable: bool = True, null_sink_magnitude_match: bool = False)[source]

Bases: Module

Equivariant perturbation transformation that preserves per-gene structure.

Unlike the standard perturbation head which collapses to a summary vector, this module transforms ALL gene embeddings based on perturbation context, maintaining the [batch, N, d] shape for downstream equivariant tasks.

Implements Type I Virtual Instrument: H_genes → H_genes_pert

forward(H_genes: Tensor, perturbation_indices: Tensor, batch_assignment: Tensor, batch_size: int) → tuple[Tensor, Tensor][source]

Apply equivariant perturbation transformation.

Parameters:

  • H_genes – [N, d] - wildtype gene embeddings

  • perturbation_indices – [total_pert_genes] - indices of perturbed genes

  • batch_assignment – [total_pert_genes] - batch index for each perturbed gene

  • batch_size – Number of genotypes in the batch. Passed explicitly (the model reads batch.num_graphs) because a genotype with no perturbed gene has no entry in batch_assignment: max + 1 would drop a trailing one. A genotype with no entry gets a zero context at every layer, so its row is the layer stack applied with attended = 0.

Returns:

  • H_genes_pert: [batch, N, d] - perturbed gene embeddings (EQUIVARIANT!)

  • context: [batch, N, d] - the attended perturbation context, returned separately so downstream heads can condition on it directly (the concat/bilinear/FiLM arms all need c_b, not only h_i + c_b).

Return type:

A tuple (H_genes_pert, context)