File size: 3,504 Bytes
07fcdfe | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 | 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,
)
|