torchcell.models.hetero_cell_nsa_retry.HeteroCellNSA

class HeteroCellNSA(gene_num: int, reaction_num: int, metabolite_num: int, hidden_channels: int, out_channels: int, attention_pattern: list[str] = ['M', 'S'], num_heads: int = 8, dropout: float = 0.1, norm: str = 'layer', activation: str = 'relu', prediction_head_config: dict[str, Any] | None = None, graph_names: list[str] | None = None)

Bases: Module

Heterogeneous Cell Model using Node-Set Attention (NSA) blocks.

This model processes heterogeneous biological cell graphs with gene, reaction, and metabolite nodes using attention-based message passing.

forward_single(data: HeteroData | Batch) → Tensor

Process a single graph or batch of graphs through the model.

forward(cell_graph: HeteroData, batch: HeteroData | Batch) → tuple[Tensor, dict[str, Tensor]]

Forward pass comparing reference cell with perturbed cell.

property num_parameters: dict[str, int]

Count the number of parameters in each component of the model.