torchcell.models.cell_graph_transformer_metabolism.CellGraphTransformerMetabolism

class CellGraphTransformerMetabolism(*args: Any, flux_layer: FluxLayer | None = None, **kwargs: Any)[source]

Bases: CellGraphTransformer

CellGraphTransformer + named production / metabolome readout heads.

Every extra heads_config entry (any key outside INHERITED_HEADS) becomes one head, registered as the submodule f"{name}_head" and emitted into reps["head_outputs"][name]. The spec keys are:

  • kind – "scalar" (ProductScalarHead) or "vector" (MetabolomeVectorHead). Required; there is no default, because guessing it is exactly the mistake that lets a scalar target broadcast across a vector head.

  • output_dim – feature count F (must be 1 for scalar).

  • use_gene_pool / dropout / param_dim – as for the inherited heads.

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

Run the inherited forward, then append the metabolism head outputs.

The encoder, graph-regularization loss and equivariant perturbation operator are the parent’s, invoked unchanged; the heads read the h_CLS / H_genes_pert pair the parent already returns.

*args/**kwargs FORWARD WHATEVER THE PARENT ACCEPTS, and that is deliberate rather than lazy. This override exists only to append heads – it has no opinion on the parent’s signature – but by naming the parent’s arguments explicitly it USED to pin them, so when the masked-label objective added observed_values / observed_mask to the parent and the trainer began passing them unconditionally, every metabolism run died with:

TypeError: CellGraphTransformerMetabolism.forward() got an unexpected
keyword argument 'observed_values'

– at the first training batch, i.e. ~20 minutes into a job, after the dataset and embeddings had loaded. A pass-through signature makes the subclass transparent to future parent arguments, which is the property the “subclass rather than transcribe” decision in this module’s header was chosen for in the first place. (It is *args, **kwargs rather than named-plus-**kwargs because that is the form mypy accepts as a Liskov-compatible override.)

property num_parameters: dict[str, int]

Parameter counts, with one entry per metabolism head and the flux layer.

The parent’s total already counts every parameter of this module (it is sum(p.numel() for p in self.parameters())), so the heads and the flux layer get their own entries here and are not added to it again.