torchcell.models.hetero_cell_bipartite_dango.HeteroCellBipartite

class HeteroCellBipartite(gene_num: int, reaction_num: int, metabolite_num: int, hidden_channels: int, out_channels: int, num_layers: int, dropout: float = 0.1, norm: str = 'layer', activation: str = 'relu', gene_encoder_config: dict[str, Any] | None = None, metabolism_config: dict[str, Any] | None = None, prediction_head_config: dict[str, Any] | None = None, gpr_conv_config: dict[str, Any] | None = None)[source]

Bases: Module

Hetero GNN with DANGO interaction head over gene/reaction/metabolite graphs.

forward_single(data: HeteroData | Batch) → Tensor[source]

Encode one hetero graph and return per-gene node embeddings.

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

Encode reference and perturbed graphs and return predictions plus outputs.

property num_parameters: dict[str, int]

Return trainable parameter counts per component of the model.