torchcell.transforms.coo_regression_to_classification.COOLabelBinningTransform

class COOLabelBinningTransform(dataset: Neo4jCellDataset, label_configs: dict[str, dict[str, Any]], normalizer: COOLabelNormalizationTransform | None = None)

Bases: BaseTransform

Transform binning COO-format continuous labels into class indices.

label_width(label: str) → int

Number of COO entries one value of label expands into.

Categorical and soft labels emit one entry per bin (num_bins); ordinal labels emit one per interior edge (num_bins - 1 thresholds).

forward(data: HeteroData | Batch) → HeteroData | Batch

Transform the data by binning specified labels in COO format.

The type list is rewritten in its own order: a configured label becomes <label>_bin_0 .. <label>_bin_{w-1} (w from label_width) and any other label passes through under its own name. Each COO entry is replaced in place by its w binned entries (configured label) or kept as is (other label), with type indices taken from that rewritten list.

inverse(data: HeteroData | Batch, seed: int = 42) → HeteroData | Batch

Inverse transform to recover continuous values from binned COO format.

Mirrors forward: the <label>_bin_<j> types of each configured label collapse back to <label> at the position of its first bin, and every other type passes through with its entries unchanged and its type index remapped. Output entries follow the rewritten type list; a binned label’s samples are decoded in sorted sample order, one seeded draw each where the decode samples.

get_bin_info(label: str) → dict[str, Any]

Get binning information for a label.