torchcell.losses.point_dist_graph_reg.PointDistGraphReg

class PointDistGraphReg(point_estimator: dict[str, Any] | None = None, distribution_loss: dict[str, Any] | None = None, graph_regularization: dict[str, Any] | None = None, buffer: dict[str, Any] | None = None, ddp: dict[str, Any] | None = None, weights: Tensor | None = None, point_loss_type: str | None = None, lambda_point: float | None = None, dist_loss_type: str | None = None, lambda_dist: float | None = None, lambda_graph_reg: float | None = None, dist_bandwidth: float | None = None, wasserstein_blur: float | None = None, wasserstein_p: float | None = None, wasserstein_scaling: float | None = None, use_buffer: bool | None = None, buffer_size: int | None = None, min_samples_for_dist: int | None = None, use_ddp_gather: bool | None = None, gather_interval: int | None = None)

Bases: Module

Modular composite loss for Cell Graph Transformer combining: 1. Point estimator (MSE or LogCosh) - required 2. Distribution loss (Dist or Wasserstein) - optional 3. Graph regularization (from model) - optional

Designed for single-phenotype predictions with graph-regularized attention.

gather_across_gpus(tensor: Tensor) → Tensor

Gather tensors from all GPUs for DDP synchronization.

Ensures all ranks see same data for buffer updates, preventing distribution mismatch across GPUs.

forward(predictions: Tensor, targets: Tensor, representations: dict[str, Any], epoch: int | None = None) → tuple[Tensor, dict[str, Any]]

Forward pass computing composite loss.

Parameters:

• predictions – Model predictions [batch_size, 1]

• targets – Target values [batch_size, 1]

• representations – Dict from model containing “graph_reg_loss”

• epoch – Current training epoch (optional, for adaptive weighting)

Returns:

Tuple of (total_loss, loss_dict)