File size: 1,700 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 | from dataclasses import dataclass
from typing import Any, Dict, Optional
import torch
@dataclass
class PopulationPerturbationBatch:
"""Batched population-level perturbation data.
Shapes
------
source_cells : [B, Ns, G]
target_cells : [B, Nt, G]
perturbation : [B, G] — multi-hot gene target indicator
source_mask : [B, Ns] — 1 = real cell, 0 = padding
target_mask : [B, Nt] — 1 = real cell, 0 = padding
"""
source_cells : torch.Tensor
target_cells : torch.Tensor
perturbation : torch.Tensor
source_mask : torch.Tensor
target_mask : torch.Tensor
context : Optional[Dict[str, Any]] = None
metadata : Optional[Dict[str, Any]] = None
# cell_line_id: integer id shared by all items in the batch (Strategy B).
# None means unconditioned (single-cell-line or no conditioning).
cell_line_id : Optional[int] = None
# prior_score: [B, G] graph-neighborhood soft prior for inverse design.
# None = no prior (standard training).
prior_score : Optional[torch.Tensor] = None
def to(self, device: torch.device) -> "PopulationPerturbationBatch":
return PopulationPerturbationBatch(
source_cells=self.source_cells.to(device),
target_cells=self.target_cells.to(device),
perturbation=self.perturbation.to(device),
source_mask=self.source_mask.to(device),
target_mask=self.target_mask.to(device),
context=self.context,
metadata=self.metadata,
cell_line_id=self.cell_line_id,
prior_score=self.prior_score.to(device) if self.prior_score is not None else None,
)
|