torchcell.models.dango.DangoPreTrain

class DangoPreTrain(gene_num: int, edge_types: list[str], hidden_channels: int = 64)

Bases: Module

GNN pre-training component of DANGO model that learns node embeddings from protein-protein interaction networks in S. cerevisiae.

As described in the paper, this module: 1. Processes PPI networks from the STRING database 2. Uses a 2-layer GNN for each network to reconstruct graph structure 3. Shares an initial embedding layer across all networks 4. Uses the output embeddings for downstream tasks

property lambda_values: dict[str, float]

Paper lambda (weight of the zero entries in the weighted MSE) per network.

Zhang et al. (2020) set lambda = 0.1 for the networks whose zero entries fell by more than 1% from STRING v9.1 to v11.0 (neighborhood, coexpression, experimental) and 1.0 for the rest (fusion, cooccurence, database). The table is defined only for those six STRING v9.1 and v11.0 network names; any other name (for example a string12_0 network) is refused rather than given 1.0. The 005/006 training scripts do not read this table: they pass determine_lambda_values() to DangoLoss.

Raises:

ValueError – if an edge type is not a STRING v9.1 or v11.0 network name.

reset_parameters() → None

Initialize model parameters

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

Forward pass for the DangoPreTrain model

Parameters:

cell_graph – The cell graph containing multiple edge types

Returns:

• ‘embeddings’: Node embeddings for each edge type (H_i^(2))

• ’reconstructions’: Reconstructed adjacency matrix rows for each edge type (FC(H_i^(2)))

Return type:

Dictionary containing