PDGrapher/src/gidflow/data/collate.py

from typing import Dict, List, Optional

import torch

from .batch import PopulationPerturbationBatch


def population_collate_fn(
    batch: List[Dict],
    max_source_cells: Optional[int] = None,
    max_target_cells: Optional[int] = None,
) -> PopulationPerturbationBatch:
    """Collate a list of per-condition dicts into a PopulationPerturbationBatch.

    Each dict must have keys:
        source_cells : Tensor[Ns, G]
        target_cells : Tensor[Nt, G]
        perturbation : Tensor[G]

    Conditions can have different numbers of cells; they are zero-padded to the
    maximum in the batch (or max_source_cells / max_target_cells if given).

    Parameters
    ----------
    batch : list of dicts
    max_source_cells : if not None, truncate / pad source cells to this length.
    max_target_cells : if not None, truncate / pad target cells to this length.

    Returns
    -------
    PopulationPerturbationBatch with tensors of shape:
        source_cells : [B, Ns_max, G]
        target_cells : [B, Nt_max, G]
        perturbation : [B, G]
        source_mask  : [B, Ns_max]
        target_mask  : [B, Nt_max]
    """
    B = len(batch)
    G = batch[0]["perturbation"].shape[0]

    source_list = [item["source_cells"] for item in batch]
    target_list = [item["target_cells"] for item in batch]
    pert_list   = [item["perturbation"] for item in batch]
    prior_list  = [item.get("prior_score") for item in batch]
    has_prior   = any(p is not None for p in prior_list)

    # Determine padded lengths
    ns_actual = [s.shape[0] for s in source_list]
    nt_actual = [t.shape[0] for t in target_list]

    Ns = max_source_cells if max_source_cells is not None else max(ns_actual)
    Nt = max_target_cells if max_target_cells is not None else max(nt_actual)

    source_padded = torch.zeros(B, Ns, G)
    target_padded = torch.zeros(B, Nt, G)
    source_mask   = torch.zeros(B, Ns, dtype=torch.bool)
    target_mask   = torch.zeros(B, Nt, dtype=torch.bool)

    for i, (src, tgt) in enumerate(zip(source_list, target_list)):
        ns = min(src.shape[0], Ns)
        nt = min(tgt.shape[0], Nt)
        source_padded[i, :ns] = src[:ns]
        target_padded[i, :nt] = tgt[:nt]
        source_mask[i, :ns]   = True
        target_mask[i, :nt]   = True

    perturbation = torch.stack(pert_list, dim=0)   # [B, G]

    prior_score = None
    if has_prior:
        prior_score = torch.stack([p if p is not None else torch.zeros(G) for p in prior_list], dim=0)

    # Collect metadata if present
    metadata_list = [item.get("metadata") for item in batch]
    metadata = metadata_list if any(m is not None for m in metadata_list) else None

    # Extract cell_line_id if all items share one (Strategy B multi-cell training)
    cl_ids = [item.get("cell_line_id") for item in batch]
    if all(c is not None for c in cl_ids):
        # Assert same cell line within batch (required for shared GCN forward)
        assert len(set(cl_ids)) == 1, (
            f"Mixed cell lines in a single batch: {set(cl_ids)}. "
            "Use CellLineBatchSampler to keep batches cell-line-homogeneous."
        )
        cell_line_id = cl_ids[0]
    else:
        cell_line_id = None

    return PopulationPerturbationBatch(
        source_cells=source_padded,
        target_cells=target_padded,
        perturbation=perturbation,
        source_mask=source_mask,
        target_mask=target_mask,
        metadata=metadata,
        cell_line_id=cell_line_id,
        prior_score=prior_score,
    )