torchcell.metabolism.flux_layer.FluxLayer

class FluxLayer(gem: GemTensors, model_gene_ids: list[str], config: FluxLayerConfig | None = None, compartments: CompartmentParameters | None = None, kcat_per_s: Tensor | None = None, molecular_weight_kda: Tensor | None = None, null_space: Tensor | None = None, constitutive_genes: list[str] | None = None)

Bases: Module

Gene tokens in, a feasible flux vector and its feasibility residuals out.

The chain, and every arrow in it is differentiable:

\[\gamma_g \longrightarrow c_u \longrightarrow c_j \longrightarrow [\,v^\ell_j,\ v^u_j\,] \longrightarrow v_j \longrightarrow \{C_{\mathrm{bal}},C_{\mathrm{th}},C_{\mathrm{cap}}, C_{\mathrm{bud}},C_{\mathrm{diss}}\}\]

A deletion sets \(\gamma_g=0\) hard, from the perturbation rather than from a learned gate, because whether a deleted gene is absent is not something the model should be free to disagree with. A complex takes a softmin over its subunits (the scarcest limits); isozymes sum (capacities add). A reaction whose availability falls to zero has \(v^\ell_j=v^u_j=0\) and therefore carries exactly zero flux – the deletion bites through the parameterization, not through a penalty that could be traded off against the data term.

gene_availability(h_genes: Tensor, perturbation_indices: Tensor, perturbation_indices_batch: Tensor) → Tensor

Per-sample \(\gamma_g\in[0,1]\) over the GEM’s genes. [B, n_gem_genes].

Deleted genes are set to zero hard, from the perturbation. Everything else is a learned sigmoid gate, which is what leaves room for dosage, hypomorphic alleles and over-expression later without changing this signature.

Constitutive genes are pinned to one. A heterologous cassette gene has no gene token, so the token-universe mask would otherwise zero it, and zeroing it switches off the pathway it was added to provide in every sample.

reaction_availability(gamma: Tensor) → Tensor

Fold \(\gamma_g\) up the GPR: softmin over complexes, sum over isozymes.

\[c_u=-\tfrac1\beta\log\!\sum_{g\in C_u}e^{-\beta\gamma_g},\qquad c_j=\min\Big(1,\ \sum_{u:\,\rho(u)=j}c_u\Big)\]

Reactions with no GPR keep availability 1: an unannotated reaction is not one that no gene catalyzes, it is one nobody has assigned a gene to, and setting it to zero would let a missing annotation delete a reaction.

dynamic_box(c_j: Tensor) → tuple[Tensor, Tensor]

The genotype-dependent flux box \([v^\ell_j,v^u_j]\).

Availability scales the box multiplicatively, so \(c_j=0\) collapses it to the single point \(0\). Enzyme capacity enters here too rather than as a penalty:

\[\bar v^u_j=\min\Big(v^u_j,\ c_j\sum_{u:\rho(u)=j}k_{\mathrm{cat},u}\bar E\Big)\]

Putting capacity in the box rather than the objective is what makes it exact at no cost, and it is why there is no C_box term in the loss.

reaction_potential(context: Tensor) → tuple[Tensor, Tensor, dict[str, Tensor]]

Per-reaction \(\Delta_r G_j\) and the mask of reactions it is valid on.

In ANCHORED mode this is Thermo-Flux Eq. (4) term by term. In FREE mode it is \(\sum_i S_{ij}\mu_i\) with \(\mu\) unconstrained. In OFF mode the mask is empty and the caller skips the term.

turnover(v: Tensor) → Tensor

\(\omega_i(v)=\tfrac12\sum_j|S_{ij}v_j|\), the throughput of metabolite i.

Doubles as the mass-balance normalizer and as the precursor-pool readout. A residual of 0.01 is negligible on a metabolite carrying flux 10 and nonsense on one carrying 0.01, which is why the balance penalty is a ratio rather than a square.

forward(h_genes: Tensor, context: Tensor, perturbation_indices: Tensor, perturbation_indices_batch: Tensor) → dict[str, Tensor]

Predict a flux vector and every feasibility residual for it.

Parameters:

• h_genes – [B, N, d] perturbed gene token representations.

• context – [B, d] whole-cell context (the perturbed-gene pool).

• perturbation_indices – Flat [total_pert] gene-token indices.

• perturbation_indices_batch – [total_pert] sample assignment.

Returns:

A dict with v [B, r], every constraint residual as a scalar, and the per-sample feasibility diagnostics. Residuals are returned rather than summed so the caller owns the weighting and every term can be logged separately.

constraint_loss(out: dict[str, Tensor]) → Tensor

Weighted sum of the physics terms. Every term is already dimensionless.

Normalization to dimensionless happens inside forward() rather than here, because a term whose scale depends on the phenotype’s units cannot be tuned against another one: that is the failure the 019 joint runs hit, where an un-normalized morphology MSE swamped expression.

coverage_report() → dict[str, Any]

What fraction of the layer rests on data rather than on a default.

This is not optional reporting. A capacity constraint built on a default kcat for every reaction is a uniform rescaling of the box, not an enzyme constraint, and the two are indistinguishable from the loss curve alone.