torchcell.models.hetero_cell_bipartite_dango_diff_gi.GeneInteractionDiff

class GeneInteractionDiff(gene_num: int, hidden_channels: int, num_layers: int, gene_multigraph: HeteroData, dropout: float = 0.2, norm: str = 'batch', activation: str = 'relu', gene_encoder_config: dict[str, Any] | None = None, local_predictor_config: dict[str, Any] | None = None, diffusion_config: dict[str, Any] | None = None, decoder_type: str = 'diffusion')[source]

Bases: GeneInteractionDango

Gene interaction model with diffusion-based prediction head.

This model inherits from GeneInteractionDango to reuse the entire encoder pipeline (graph processing, convolutions, aggregation) while replacing only the final prediction head with a diffusion decoder. This ensures consistency in feature extraction between the deterministic and diffusion versions of the model.

Key inherited components: - Gene embeddings and preprocessing - Graph convolution layers - Global and local aggregation - forward_single() for processing graphs

What’s replaced: - The MLP prediction head (global_interaction_predictor) is removed - A diffusion decoder with cross-attention is used instead

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

Forward pass through the model.

Parameters:

  • cell_graph – Full cell graph

  • batch – Batch of perturbed genes

Returns:

Phenotype predictions [batch_size, 1] representations: Dictionary of intermediate representations

Return type:

predictions

compute_diffusion_loss(y_true: Tensor, z_c: Tensor, t_mode: str = 'random') → Tensor[source]

Compute the training loss based on decoder type.

Parameters:

  • y_true – Ground truth phenotype values [batch_size, 1]

  • z_c – Combined latent embeddings for conditioning [batch_size, hidden_dim * 2]

  • t_mode – Timestep sampling mode (“zero”, “partial”, “full”) - only used for diffusion

Returns:

Loss value

sample(cell_graph: HeteroData, batch: HeteroData, num_samples: int | None = None) → Tensor[source]

Sample phenotype predictions using the diffusion model.

Note: This method performs a fresh forward pass through the encoder to get new embeddings, ensuring we don’t reuse stale representations from training.

Parameters:

  • cell_graph – Full cell graph

  • batch – Batch of perturbed genes

  • num_samples – Number of samples to generate

Returns:

Sampled phenotype predictions [batch_size, 1]

property num_parameters: dict[str, int]

Get parameter counts for each component.