torchcell.losses

Loss functions. The package exports ListMLELoss, a listwise ranking loss, and DCellLoss, the root plus subsystem loss of the DCell model. The remaining modules (multi_dim_nan_tolerant, distributional, mle_wasserstein, point_dist_graph_reg and others) are imported by module path from the experiment training scripts and are not re-exported. Every class below is listed under the module that defines it.

dango

DANGO loss with epoch-scheduled weighting of reconstruction vs interaction loss.

Classes

dango.DangoLossSched	Abstract base class for DANGO loss schedulers.
dango.PreThenPost	Schedule that uses reconstruction loss only before transition_epoch, then switches to interaction loss only.
dango.LinearUntilUniform	Schedule that linearly transitions from reconstruction-focused to uniform weighting.
dango.LinearUntilFlipped	Schedule that linearly transitions from reconstruction-focused to interaction-focused.
dango.DangoLoss	Combined loss for the DANGO model with dynamic weighting between reconstruction loss and interaction loss.

dcell

Loss for the DCell model: root MSE plus optional subsystem auxiliary losses.

Classes

dcell.DCellLoss

Loss function for the DCell model (Ma et al. 2018).

diffusion_loss

Strict diffusion loss with optional explicit x0 supervision.

Classes

diffusion_loss.DiffusionLoss

Strict diffusion loss for conditional decoders.

distributional

Distributional decoder heads + proper-scoring-rule losses for the multitask CGT.

Classes

distributional.DistHead

Pluggable distributional readout: interpret a head's raw params, score them, pool.

Functions

distributional.gaussian_crps	Closed-form CRPS of a Gaussian \(\mathcal N(\mu,\sigma)\) vs observations.
distributional.laplace_crps	Closed-form CRPS of a Laplace \(\mathrm{Lap}(\mu,b)\) vs observations.
distributional.gaussian_nll	Gaussian negative log-likelihood, WITHOUT the constant term.
distributional.energy_score	Monte-Carlo energy score of a low-rank-plus-diagonal Gaussian, per row.
distributional.pinball	Pinball (quantile / check) loss for a set of predicted quantiles.
distributional.masked_mean	Mean of elem over the entries feature_mask selects.
distributional.pit_values	Probability integral transform -- ONE calibration diagnostic for every mode.
distributional.coverage	Empirical coverage of the central alpha predictive interval, from PIT values.
distributional.pit_ks	Kolmogorov-Smirnov distance between the PIT distribution and Uniform(0,1).
distributional.dist_param_dim	Params-per-feature for a config dist value (the head Linear-width multiplier).
distributional.make_dist_head	Build a DistHead from a config dist value.

isomorphic_cell_loss

Isomorphic cell losses combining MSE, distribution, and SupCR terms.

Classes

isomorphic_cell_loss.ICLoss	Weighted sum of MSE, distribution, and SupCR cell losses.
isomorphic_cell_loss.ICLossStd	IC loss with learnable per-task standard deviations (heteroscedastic).

list_mle

ListMLE listwise ranking loss for ordered prediction tasks.

Classes

list_mle.ListMLELoss

Listwise Maximum Likelihood Estimation ranking loss.

logcosh

Log-cosh regression loss module.

Classes

logcosh.LogCoshLoss

Log-cosh loss: a smooth, outlier-robust alternative to MSE.

mle_dist_supcr

Combined MLE, distribution, and supervised contrastive loss with buffering.

Classes

mle_dist_supcr.AdaptiveWeighting	Manages adaptive weighting for buffer contributions during training.
mle_dist_supcr.TemperatureScheduler	Manages temperature scheduling for SupCR loss.
mle_dist_supcr.BufferedWeightedDistLoss	Enhanced distribution loss with circular buffer for small batch training.
mle_dist_supcr.BufferedWeightedSupCRCell	Enhanced SupCR loss with circular buffer for small batch training.
mle_dist_supcr.MleDistSupCR	Composite loss combining MSE, distribution matching, and SupCR.

mle_wasserstein

Composite MSE + Wasserstein + SupCR loss with buffers and DDP gathering.

Classes

mle_wasserstein.AdaptiveWeighting	Manages adaptive weighting for buffer contributions during training.
mle_wasserstein.TemperatureScheduler	Manages temperature scheduling for SupCR loss.
mle_wasserstein.WeightedWassersteinLoss	Weighted Wasserstein distance loss for multi-dimensional outputs.
mle_wasserstein.BufferedWeightedWassersteinLoss	Enhanced Wasserstein loss with circular buffer for small batch training.
mle_wasserstein.BufferedWeightedSupCRCell	Enhanced SupCR loss with circular buffer for small batch training.
mle_wasserstein.MleWassSupCR	Composite MSE + Wasserstein + SupCR loss.

multi_dim_nan_tolerant

NaN-tolerant multi-dimensional loss functions for cell-state regression.

Classes

multi_dim_nan_tolerant.SupCR	Supervised Contrastive Regression (SupCR) loss over multiple target dimensions.
multi_dim_nan_tolerant.WeightedSupCRCell	SupCR loss with per-dimension weights for cell-state embeddings.
multi_dim_nan_tolerant.NaNTolerantL1Loss	Per-dimension L1 loss that ignores NaN targets.
multi_dim_nan_tolerant.NaNTolerantHuberLoss	Per-dimension Huber loss that ignores NaN targets.
multi_dim_nan_tolerant.NaNTolerantLogCoshLoss	Per-dimension log-cosh loss that ignores NaN targets.
multi_dim_nan_tolerant.WeightedMSELoss	Per-dimension MSE loss with optional normalized dimension weights.
multi_dim_nan_tolerant.NaNTolerantMSELoss	Per-dimension MSE loss that ignores NaN targets.
multi_dim_nan_tolerant.NaNTolerantQuantileLoss	Quantile regression loss that handles NaN values.
multi_dim_nan_tolerant.FastSoftSort	Fast differentiable sorting following the original paper implementation.
multi_dim_nan_tolerant.WeightedDistLoss	Distribution-matching (DistLoss) loss using KDE and differentiable sorting.
multi_dim_nan_tolerant.CombinedRegressionLoss	Dispatch to a NaN-tolerant regression loss and apply per-dimension weights.
multi_dim_nan_tolerant.MultiDimNaNTolerantCELoss	Multi-task cross-entropy loss that ignores NaN targets per task.
multi_dim_nan_tolerant.CombinedCELoss	Cross-entropy loss combining per-task losses with dimension weights.
multi_dim_nan_tolerant.MonotonicParameter	Parameter that ensures values are monotonically increasing via cumulative softplus.
multi_dim_nan_tolerant.MultiDimNaNTolerantOrdinalCELoss	Multi-task ordinal cross-entropy loss with learnable monotonic thresholds.
multi_dim_nan_tolerant.CombinedOrdinalCELoss	Combined ordinal classification loss with dimension-wise weighting.
multi_dim_nan_tolerant.CategoricalEntropyRegLoss	Entropy regularization loss that works with both soft and one-hot categorical targets.
multi_dim_nan_tolerant.MseCategoricalEntropyRegLoss	Combine MSE loss with entropy regularization, tracking per-dimension losses.

Functions

multi_dim_nan_tolerant.isotonic_l2_pav	Solve isotonic regression using the PAV (Pool Adjacent Violators) algorithm.
multi_dim_nan_tolerant.fast_soft_sort	Fast differentiable sorting that follows the original paper's implementation.

point_dist_graph_reg

Composite point + distribution + graph-regularization loss for the cell transformer.

Classes

point_dist_graph_reg.PointDistGraphReg

Modular composite loss for Cell Graph Transformer combining: 1.