torchcell.models

Model implementations, one module per architecture. They range from sequence language-model wrappers (Nucleotide Transformer, ESM-2, ProtT5, the fungal up/down-stream transformer) through set and graph models (DeepSet, GCN and GAT encoders, DiffPool and SAGPool variants) to the cell models used in the experiments: DCell, the DANGO family (for example GeneInteractionDango in hetero_cell_bipartite_dango_gi), and the Cell Graph Transformer (CellGraphTransformer, with its equivariant variant CellGraphTransformer). The package namespace re-exports a small subset (torchcell.models.__all__); every class below is listed under the module that defines it.

`cell_graph_transformer`

Cell Graph Transformer with graph-regularized attention heads.

Classes 

`cell_graph_transformer.GraphRegularizedTransformerLayer`	Transformer layer with graph-regularized attention heads.
`cell_graph_transformer.HyperSAGNN`	Hypergraph Self-Attention Graph Neural Network for perturbation sets.
`cell_graph_transformer.PerturbationHead`	Perturbation head with switchable attention mechanisms.
`cell_graph_transformer.CellGraphTransformer`	Cell Graph Transformer model.

Functions 

`cell_graph_transformer.calculate_weight_l2_norm`	Calculate L2 norm of all model weights.
`cell_graph_transformer.compute_smoothness`	Compute smoothness of node features (oversmoothing diagnostic).

`cell_graph_transformer_metabolism`

CGT-Metabolism: the Cell Graph Transformer with production / metabolome readouts.

Classes 

`cell_graph_transformer_metabolism.FluxMetaboliteHead`	Read a measured metabolite level off the predicted flux vector's turnover.
`cell_graph_transformer_metabolism.FluxScalarHead`	Read a scalar production phenotype off the predicted flux vector.
`cell_graph_transformer_metabolism.ProductScalarHead`	Whole-cell readout for a SCALAR production phenotype (one value per strain).
`cell_graph_transformer_metabolism.MetabolomeVectorHead`	Whole-cell readout for a FIXED-COLUMN metabolome vector (Mulleder's 19 AAs).
`cell_graph_transformer_metabolism.CellGraphTransformerMetabolism`	CellGraphTransformer + named production / metabolome readout heads.

`dango.DangoPreTrain`	GNN pre-training component of DANGO model that learns node embeddings from protein-protein interaction networks in S.
`dango.MetaEmbedding`	Meta-embedding module to integrate embeddings from multiple networks.
`dango.HyperSAGNN`	Fully vectorized Hypergraph Self-Attention Graph Neural Network that handles all perturbation sets in a single forward pass.
`dango.Dango`	DANGO model for predicting higher-order genetic interactions

`dcell`

DCell model: a GO-ontology-structured neural network over gene perturbations.

Classes 

`dcell.DCell`	DCell network mirroring the GO ontology as a hierarchy of subsystems.
`dcell.DCellSubsystem`	Individual subsystem module as described in DCell paper.

`dcell_opt`

Optimized DCell model for torch.compile compatibility.

Classes 

`dcell_opt.DCellOpt`	Optimized DCell with tensorized operations for torch.compile.
`dcell_opt.DCellSubsystem`	Individual subsystem module as described in DCell paper.

`diffusion_decoder.SinusoidalTimeEmbedding`	Sinusoidal time embeddings for diffusion timesteps.
`diffusion_decoder.CrossAttention`	Cross-attention module where diffusion state queries graph embeddings.
`diffusion_decoder.DenoisingBlock`	Single denoising block with cross-attention and feedforward.
`diffusion_decoder.DiffusionDecoder`	Diffusion decoder with cross-attention conditioning on graph embeddings.

`equivariant_cell_graph_transformer`

Equivariant Cell Graph Transformer with graph-regularized attention heads.

Classes 

`equivariant_cell_graph_transformer.GraphRegularizedTransformerLayer`	Transformer layer with graph-regularized attention heads.
`equivariant_cell_graph_transformer.HyperSAGNN`	Hypergraph Self-Attention Graph Neural Network for perturbation sets.
`equivariant_cell_graph_transformer.EquivariantPerturbationTransform`	Equivariant perturbation transformation that preserves per-gene structure.
`equivariant_cell_graph_transformer.PerturbationGraphPropagation`	Route the perturbation to graph NEIGHBOURS -- the missing pair-(p, i) term.
`equivariant_cell_graph_transformer.LowRankBilinear`	Explicit rank-R multiplicative interaction between gene identity and perturbation.
`equivariant_cell_graph_transformer.ObservedLabelEncoder`	Encode PARTIALLY OBSERVED labels so the decoder can condition on them.
`equivariant_cell_graph_transformer.CrossGeneMixing`	GEARS-style cross-gene mixing: one shared global state, then per-gene readout.
`equivariant_cell_graph_transformer.ResponseBasisHead`	Low-rank response-basis decoder: the perturbation activates shared programs.
`equivariant_cell_graph_transformer.PerceiverMixing`	POST-perturbation gene-gene mixing through a learned latent bottleneck.
`equivariant_cell_graph_transformer.PerturbationHead`	Perturbation readout head for gene interaction prediction.
`equivariant_cell_graph_transformer.GlobalHead`	Whole-cell readout head for global phenotypes.
`equivariant_cell_graph_transformer.CrossAttnHead`	S3 cross-attention readout head (DETR / Perceiver style) for vector phenotypes.
`equivariant_cell_graph_transformer.PerGeneHead`	Per-gene (equivariant) readout head for gene-resolved phenotypes.
`equivariant_cell_graph_transformer.PerMetaboliteHead`	Per-metabolite readout head via GPR pooling of gene embeddings.
`equivariant_cell_graph_transformer.MaskedMultitaskLoss`	Masked multitask loss over config-selected heads plus graph regularization.
`equivariant_cell_graph_transformer.CellGraphTransformer`	Equivariant Cell Graph Transformer model.

Functions 

`equivariant_cell_graph_transformer.calculate_weight_l2_norm`	Calculate L2 norm of all model weights.
`equivariant_cell_graph_transformer.compute_smoothness`	Compute smoothness of node features (oversmoothing diagnostic).

`hetero_cell_bipartite_dango.GeneInteractionAttention`	Self-attention for digenic and trigenic gene interaction prediction.
`hetero_cell_bipartite_dango.GeneInteractionPredictor`	Predict interaction scores from squared static/dynamic embedding differences.
`hetero_cell_bipartite_dango.AttentionalGraphAggregation`	Attentional pooling that gates and transforms node features before aggregation.
`hetero_cell_bipartite_dango.PreProcessor`	MLP that projects raw input features to the hidden dimension.
`hetero_cell_bipartite_dango.AttentionConvWrapper`	Wrap a graph conv with projection, normalization, activation, and dropout.
`hetero_cell_bipartite_dango.HeteroCellBipartite`	Hetero GNN with DANGO interaction head over gene/reaction/metabolite graphs.

`hetero_cell_bipartite_dango_diff_gi.LinearDecoder`	Simple linear decoder for baseline comparison.
`hetero_cell_bipartite_dango_diff_gi.GeneInteractionDiff`	Gene interaction model with diffusion-based prediction head.

`hetero_cell_bipartite_dango_gi`

Heterogeneous bipartite cell-graph model for DANGO gene-interaction prediction.

Classes 

`hetero_cell_bipartite_dango_gi.SelfAttentionGraphAggregation`	Self-attention aggregation over multiple graph views of the same nodes.
`hetero_cell_bipartite_dango_gi.PairwiseGraphAggregation`	Pairwise interaction aggregation over multiple graph representations.
`hetero_cell_bipartite_dango_gi.AggregationNormNotImplementedError`	An `aggregation_norm` other than None was requested from this module.
`hetero_cell_bipartite_dango_gi.HeteroConvAggregator`	HeteroConv wrapper with configurable aggregation strategies.
`hetero_cell_bipartite_dango_gi.AttentionalGraphAggregation`	Gated attentional pooling of node features into a graph-level representation.
`hetero_cell_bipartite_dango_gi.DangoLikeHyperSAGNN`	Dango-like HyperSAGNN for local gene interaction prediction.
`hetero_cell_bipartite_dango_gi.GeneInteractionPredictor`	Predict gene-interaction scores from squared static/dynamic embedding diffs.
`hetero_cell_bipartite_dango_gi.PreProcessor`	MLP that projects input features into the hidden space before convolutions.
`hetero_cell_bipartite_dango_gi.AttentionConvWrapper`	Wrap a conv with output projection, optional norm, activation, and dropout.
`hetero_cell_bipartite_dango_gi.GeneInteractionDango`	Full DANGO model: gene embeddings, hetero conv encoder, and interaction head.

Functions 

`hetero_cell_bipartite_dango_gi.get_activation`	Return the `act_register` module named `activation`, refusing unknown names.
`hetero_cell_bipartite_dango_gi.require_no_aggregation_norm`	Refuse any `aggregation_norm` other than None (the only built behavior).
`hetero_cell_bipartite_dango_gi.get_norm_layer`	Return a LayerNorm or BatchNorm1d layer for the given norm name.
`hetero_cell_bipartite_dango_gi.create_conv_layer`	Create appropriate conv layer based on encoder type.
`hetero_cell_bipartite_dango_gi.calculate_weight_l2_norm`	Calculate L2 norm of all model weights.
`hetero_cell_bipartite_dango_gi.calculate_rolling_correlation`	Calculate rolling correlation between two series.

`hetero_cell_bipartite_dango_gi_lazy`

Lazy bipartite Dango gene-interaction model using masked, zero-copy message passing.

Classes 

`hetero_cell_bipartite_dango_gi_lazy.SelfAttentionGraphAggregation`	Self-attention mechanism for aggregating multiple graph representations of the same nodes.
`hetero_cell_bipartite_dango_gi_lazy.PairwiseGraphAggregation`	Pairwise interaction mechanism for aggregating multiple graph representations.
`hetero_cell_bipartite_dango_gi_lazy.HeteroConvAggregator`	HeteroConv wrapper with configurable aggregation strategies.
`hetero_cell_bipartite_dango_gi_lazy.AttentionalGraphAggregation`	Attentional aggregation pooling node features into graph-level vectors.
`hetero_cell_bipartite_dango_gi_lazy.DangoLikeHyperSAGNN`	Dango-like HyperSAGNN for local gene interaction prediction.
`hetero_cell_bipartite_dango_gi_lazy.GeneInteractionPredictor`	Predict gene interaction scores from gene embeddings via a HyperSAGNN.
`hetero_cell_bipartite_dango_gi_lazy.PreProcessor`	MLP that preprocesses node features with linear, norm, activation, and dropout layers.
`hetero_cell_bipartite_dango_gi_lazy.AttentionConvWrapper`	Wrapper applying norm, activation, and dropout around a graph convolution.
`hetero_cell_bipartite_dango_gi_lazy.GeneInteractionDango`	Lazy Dango-style model predicting gene interactions from a gene multigraph.

Functions 

`hetero_cell_bipartite_dango_gi_lazy.get_norm_layer`	Get normalization layer from norm_register - NO FALLBACK.
`hetero_cell_bipartite_dango_gi_lazy.create_conv_layer`	Create appropriate conv layer based on encoder type.
`hetero_cell_bipartite_dango_gi_lazy.calculate_weight_l2_norm`	Calculate L2 norm of all model weights.
`hetero_cell_bipartite_dango_gi_lazy.calculate_rolling_correlation`	Calculate rolling correlation between two series.

`hetero_cell_nsa_retry`

Heterogeneous cell model using Node-Set Attention blocks over gene/reaction/metabolite graphs.

Classes 

`hetero_cell_nsa_retry.AttentionalGraphAggregation`	Attentional aggregation for graph-level representations.
`hetero_cell_nsa_retry.PreProcessor`	MLP for preprocessing node features.
`hetero_cell_nsa_retry.HeteroCellNSA`	Heterogeneous Cell Model using Node-Set Attention (NSA) blocks.

`llm.NucleotideModel`	Abstract interface for nucleotide-sequence language models.
`llm.PeptideModel`	Abstract interface for peptide-sequence language models.
`llm.pretrained_LLM`	Container pairing a tokenizer with a pretrained masked-LM model.

`self_attention_deep_set.SelfAttention`	Multi-head self-attention returning outputs and per-graph weights.
`self_attention_deep_set.SelfAttentionDeepSet`	Deep Sets encoder with self-attention between node and set layers.