dpacman/data_modules/pair.py

#!/usr/bin/env python3
import argparse
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader, Sampler, BatchSampler
import torch.distributed as dist
from lightning import LightningDataModule
from pathlib import Path
from multiprocessing import cpu_count
import random
import pandas as pd
import shelve
from torch.nn.utils.rnn import pad_sequence
from typing import List, Iterable, Sequence
import sys
import rootutils
import logging
import math
from dpacman.utils import pylogger

root = rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
logger = pylogger.RankedLogger(__name__, rank_zero_only=True)

class PreBatchedDistributedBatchSampler(BatchSampler):
    """
    Accepts a precomputed list of batches (list[list[int]]) and shards them across DDP ranks.
    - shuffle_batch_order: shuffle order of batches each epoch (deterministic via set_epoch)
    - drop_last: drop remainder so each rank gets same #steps
    """
    def __init__(self, batches, shuffle_batch_order=False, drop_last=False, seed: int = 0):
        # super expects attributes batch_size, drop_last, sampler – but we don't need them.
        # We only need to subclass BatchSampler to satisfy Lightning's check.
        self.batches = [list(b) for b in batches]
        self.shuffle = shuffle_batch_order
        self.drop_last = drop_last
        self.seed = int(seed)
        self.epoch = 0

        if dist.is_available() and dist.is_initialized():
            self.world_size = dist.get_world_size()
            self.rank = dist.get_rank()
        else:
            self.world_size = 1
            self.rank = 0

    def __iter__(self):
        n_batches = len(self.batches)
        order = list(range(n_batches))

        if self.shuffle:
            g = torch.Generator()
            g.manual_seed(self.seed + self.epoch)
            order = torch.randperm(n_batches, generator=g).tolist()

        # make divisible across ranks
        if self.drop_last:
            total = (len(order) // self.world_size) * self.world_size
            order = order[:total]
        else:
            pad = (-len(order)) % self.world_size
            if pad:
                order = order + order[:pad]

        # shard by rank
        for i in order[self.rank::self.world_size]:
            yield self.batches[i]

    def __len__(self):
        n = len(self.batches)
        if self.drop_last:
            return (n // self.world_size)
        return math.ceil(n / self.world_size)

    # Lightning will call this if present via its epoch hooks
    def set_epoch(self, epoch: int):
        self.epoch = int(epoch)

class PreBatchedSampler(Sampler[List[int]]):
    """
    Yields precomputed batches of indices, e.g. [[3,7,9], [0,1,2], ...].
    Useful when you've already formed batches by length.
    """

    def __init__(
        self,
        batches: Sequence[Sequence[int]],
        shuffle_batch_order: bool = False,
        generator=None,
    ):
        self.batches = [list(b) for b in batches]
        self.shuffle_batch_order = shuffle_batch_order
        self.generator = generator

    def __iter__(self) -> Iterable[List[int]]:
        if self.shuffle_batch_order:
            # local copy we can shuffle without touching the original
            idxs = list(range(len(self.batches)))
            g = self.generator if self.generator is not None else torch.Generator()
            perm = torch.randperm(len(idxs), generator=g).tolist()
            for i in perm:
                yield self.batches[i]
        else:
            for b in self.batches:
                yield b

    def __len__(self) -> int:
        return len(self.batches)


def compute_tr_lengths_from_shelf(
    tr_shelf_path: str, tr_sequences: list[str]
) -> list[int]:
    """
    Opens the TR shelf once and returns length for each sequence.
    2D array -> length = shape[0]; 1D array (pooled) -> length = 1.
    """
    lengths = []
    with shelve.open(tr_shelf_path, flag="r") as db:
        for s in tr_sequences:
            arr = np.asarray(db[str(s)])
            if arr.ndim == 1:
                lengths.append(1)
            else:
                lengths.append(int(arr.shape[0]))
    return lengths


def make_length_batches(
    dataset_records: list[dict],
    tr_shelf_path: str,
    batch_size: int,
    drop_last: bool = False,
) -> list[list[int]]:
    """
    dataset_records: output of PairDataset._load_and_normalize(...), i.e. list of dicts with
        keys: "dna_sequence", "tr_sequence", "scores", ...
    Returns a list of batches, each a list of indices, sorted by (dna_len, tr_len).
    """
    # DNA length comes from label length
    dna_lens = [len(r["scores"]) for r in dataset_records]
    tr_seqs = [r["tr_sequence"] for r in dataset_records]

    # TR length requires a quick shelf lookup (done once here)
    tr_lens = compute_tr_lengths_from_shelf(tr_shelf_path, tr_seqs)

    # sort indices by (dna_len, tr_len)
    idxs = list(range(len(dataset_records)))
    idxs.sort(key=lambda i: (dna_lens[i], tr_lens[i]))

    # chunk into fixed-size batches
    batches = [idxs[i : i + batch_size] for i in range(0, len(idxs), batch_size)]
    if drop_last and len(batches) and len(batches[-1]) < batch_size:
        batches.pop()
    return batches


# ---- dataset ---------------------------------------------------------
class PairDataset(Dataset):
    def __init__(
        self, dataset: pd.DataFrame, norm_value: int = 1333, round_to: int = 4, score_col="scores", target_col="dna_sequence", binder_col="tr_sequence"
    ):
        """
        Args:
            - dataset: a dataset with the needed information: ID, dna_sequence, tr_sequence, scores
            - norm_value: max score, which we'll use to divide all the integer scores in "scores"
            - round_to: how many decimal places for the numerical score values
        """
        self.fake_scores=False
        self.score_col = score_col
        self.target_col = target_col
        self.binder_col = binder_col
        self.norm_value = norm_value
        self.round_to = round_to
        self.dataset = self._load_and_normalize(dataset)

    def _load_and_normalize(self, dataset):
        """
        Labels come in looking like "0,0,0,100,100,133,133,100,100,0,0,"
        This method turns the labels from strings into floats out to 4 decimal places
        """
        if self.score_col not in dataset.columns:
            logger.info(f"Scores not provided. Adding placeholder scores where all positions are considered binding")
            dataset[self.score_col] = dataset[self.target_col].str.len()
            dataset[self.score_col] = dataset[self.score_col].apply(lambda x: ",".join([str(self.norm_value)]*x))
            self.fake_scores=True
        # split string into list of strings
        dataset[self.score_col] = dataset[self.score_col].apply(lambda x: x.split(","))
        dataset["copycol"] = dataset[self.score_col]
        # turn list of strings into list of normalized, rounded floats
        dataset[self.score_col] = dataset[self.score_col].apply(
            lambda x: [round(int(y) / self.norm_value, self.round_to) for y in x]
        )
        # convert to records for ease of loading
        dataset = dataset.to_dict(orient="records")
        return dataset

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, idx):
        item = self.dataset[idx]
        return {**(item if isinstance(item, dict) else {})}


class PairDataModule(LightningDataModule):
    def __init__(
        self,
        train_file: Path | str = "../data_files/splits/train.csv",
        val_file: Path | str = "../data_files/splits/val.csv",
        test_file: Path | str = "../data_files/splits/test.csv",
        tr_shelf_path: (
            Path | str
        ) = "../data_files/processed/embeddings/fimo_hits_only/trs_esm.shelf",
        dna_shelf_path: (
            Path | str
        ) = "../data_files/processed/embeddings/fimo_hits_only/baby_peaks_segmentnt_pernuc_with_onehot.shelf",
        batch_size: int = 1,
        num_workers=8,
        maximize_num_workers=False,
        debug_run: bool = False,
        pin_memory: bool = False,
        shuffle_train_batch_order: bool = True,
        score_col: str = "scores",
        target_col: str = "dna_sequence",
        binder_col: str = "tr_sequence",
        norm_value: int = 1333
    ):
        super().__init__()
        self.save_hyperparameters()
        self.debug_run = debug_run
        self.norm_value = norm_value

        # Initialize the data files
        self.train_data_file = train_file
        self.val_data_file = val_file
        self.test_data_file = test_file
        self.target_col = target_col
        self.binder_col = binder_col
        self.score_col = score_col

        # Initialize hyperparameters like batch size
        self.batch_size = batch_size
        self.num_workers = (
            cpu_count() if maximize_num_workers else min(num_workers, cpu_count())
        )

        # Set up ShelfCollator
        self.collate = ShelfCollator(
            tr_shelf_path=str(tr_shelf_path),
            dna_shelf_path=str(dna_shelf_path),
            tr_key=self.binder_col,
            dna_key=self.target_col,
            dtype=torch.float32,
            pad_value=-1.0,
            debug_run =self.debug_run,
            score_col = self.score_col
        )
        self.drop_last = False  # or True, your choice
        self.shuffle_batch_order = shuffle_train_batch_order  # False keep batches deterministic per epoch; set True if you want to shuffle batch order

        logger.info(f"num_workers={self.num_workers}")
        logger.info("Initialized BinderDecoyDataModule constants")

    def load_file(self, file_path, lim=None):
        """
        Load and unpack an input csv whose columns are binder_path,glm_path,label
        """
        try:
            df = pd.read_csv(file_path)
            if lim is not None:
                df = df[:lim].reset_index(drop=True)
            return df
        except:
            raise Exception(f"{file_path} is not a valid file")

    def setup(self, stage: str | None = None):
        lim = 5 if self.debug_run else None

        # FIT: build train & val (so val exists during training)
        if stage in (None, "fit"):
            if not hasattr(self, "train_dataset"):
                train_df = self.load_file(self.train_data_file, lim=lim)
                self.train_dataset = PairDataset(train_df, norm_value = self.norm_value, score_col = self.score_col, target_col = self.target_col, binder_col = self.binder_col)
                self.train_batches = make_length_batches(
                    dataset_records=self.train_dataset.dataset,
                    tr_shelf_path=str(self.hparams.tr_shelf_path),
                    batch_size=self.batch_size,
                    drop_last=self.drop_last,
                )
                self.train_batch_sampler = PreBatchedDistributedBatchSampler(
                    self.train_batches,
                    shuffle_batch_order=self.shuffle_batch_order,
                    drop_last=self.drop_last,
                    seed=0,
                )

            if not hasattr(self, "val_dataset"):
                val_df = self.load_file(self.val_data_file, lim=lim)
                self.val_dataset = PairDataset(val_df, norm_value = self.norm_value, score_col = self.score_col, target_col = self.target_col, binder_col = self.binder_col)
                self.val_batches = make_length_batches(
                    dataset_records=self.val_dataset.dataset,
                    tr_shelf_path=str(self.hparams.tr_shelf_path),
                    batch_size=self.batch_size,
                    drop_last=False,
                )
                self.val_batch_sampler = PreBatchedDistributedBatchSampler(
                    self.val_batches, shuffle_batch_order=False, drop_last=False, seed=0
                )

        # VALIDATE called standalone: ensure val is built
        if stage in (None, "validate"):
            if not hasattr(self, "val_dataset"):
                val_df = self.load_file(self.val_data_file, lim=lim)
                self.val_dataset = PairDataset(val_df, norm_value = self.norm_value, score_col = self.score_col, target_col = self.target_col, binder_col = self.binder_col)
                self.val_batches = make_length_batches(
                    dataset_records=self.val_dataset.dataset,
                    tr_shelf_path=str(self.hparams.tr_shelf_path),
                    batch_size=self.batch_size,
                    drop_last=False,
                )
                self.val_batch_sampler = PreBatchedDistributedBatchSampler(
                    self.val_batches, shuffle_batch_order=False, drop_last=False, seed=0
                )
                
        # TEST phase
        if stage in (None, "test"):
            if not hasattr(self, "test_dataset"):
                test_df = self.load_file(self.test_data_file, lim=lim)
                self.test_dataset = PairDataset(test_df, norm_value = self.norm_value, score_col = self.score_col, target_col = self.target_col, binder_col = self.binder_col)
                self.test_batches = make_length_batches(
                    dataset_records=self.test_dataset.dataset,
                    tr_shelf_path=str(self.hparams.tr_shelf_path),
                    batch_size=self.batch_size,
                    drop_last=False,
                )
                self.test_batch_sampler = PreBatchedSampler(
                    self.test_batches, shuffle_batch_order=False
                )

    def train_dataloader(self):
        return DataLoader(
            self.train_dataset,
            batch_sampler=self.train_batch_sampler,
            collate_fn=self.collate,
            num_workers=self.num_workers,
            persistent_workers=(self.num_workers > 0),
            pin_memory=self.hparams.pin_memory,
        )

    def val_dataloader(self):
        return DataLoader(
            self.val_dataset,
            batch_sampler=self.val_batch_sampler,
            collate_fn=self.collate,
            num_workers=self.num_workers,
            persistent_workers=(self.num_workers > 0),
            pin_memory=self.hparams.pin_memory,
        )

    def test_dataloader(self):
        return DataLoader(
            self.test_dataset,
            batch_sampler=self.test_batch_sampler,
            collate_fn=self.collate,
            num_workers=self.num_workers,
            persistent_workers=(self.num_workers > 0),
            pin_memory=self.hparams.pin_memory,
        )
    
    def predict_dataloader(self):
        # Same as test
        return DataLoader(
            self.test_dataset,
            batch_sampler=self.test_batch_sampler,
            collate_fn=self.collate,
            num_workers=self.num_workers,
            persistent_workers=(self.num_workers > 0),
            pin_memory=self.hparams.pin_memory,
        )


class ShelfCollator:
    """
    Lazily opens TR (binder) and DNA shelves the first time each worker calls __call__.
    Expects each item to contain keys:
        - "tr_sequence": str  (key for TR shelf)
        - "dna_sequence": str (key for DNA shelf)
        - "scores": list[float] (per-base labels for DNA)
        - optional "ID"
    Returns a dict with:
        - binder_emb:  FloatTensor [B, Lb_max, Db]  (padded)
        - binder_mask: BoolTensor  [B, Lb_max]
        - glm_emb:     FloatTensor [B, Lg_max, Dg]  (padded)
        - glm_mask:    BoolTensor  [B, Lg_max]
        - labels:      FloatTensor [B, Lg_max]      (padded, zeros where masked)
        - ids, tr_sequences, dna_sequences: lists
    """

    def __init__(
        self,
        tr_shelf_path: str,
        dna_shelf_path: str,
        tr_key: str = "tr_sequence",
        dna_key: str = "dna_sequence",
        dtype: torch.dtype = torch.float32,
        pad_value: float = -1.0,
        debug_run: bool = False,
        score_col = "scores"
    ):
        self.tr_path = tr_shelf_path
        self.dna_path = dna_shelf_path
        self.score_col = score_col
        self.tr_key = tr_key
        self.dna_key = dna_key
        self.dtype = dtype
        self.pad_value = pad_value
        self.debug_run = debug_run

        # opened lazily per worker:
        self._tr_db = None
        self._dna_db = None

    def _ensure_open(self):
        if self._tr_db is None:
            self._tr_db = shelve.open(self.tr_path, flag="r")  # read-only
        if self._dna_db is None:
            self._dna_db = shelve.open(self.dna_path, flag="r")

    def __call__(self, batch):
        """
        batch: list[dict] from Dataset.__getitem__
        """
        self._ensure_open()

        ids = [b.get("ID", None) for b in batch]
        tr_seqs = [b[self.tr_key] for b in batch]
        dna_seqs = [b[self.dna_key] for b in batch]
        scores_list = [b[self.score_col] for b in batch]

        # 1) Fetch embeddings lazily from shelves
        binder_list = []
        glm_list = []
        binder_lens = []
        glm_lens = []

        for tr, dna, scores in zip(tr_seqs, dna_seqs, scores_list):
            # ----- binder/TR -----
            tr_arr = np.asarray(self._tr_db[str(tr)])
            # ensure 2D: [Lb, Db] (if pooled 1D, make length=1)
            if tr_arr.ndim == 1:
                tr_arr = tr_arr[None, :]
            binder_list.append(torch.from_numpy(tr_arr).to(self.dtype))
            binder_lens.append(tr_arr.shape[0])

            # ----- DNA / GLM -----
            dna_arr = np.asarray(self._dna_db[str(dna)])
            if dna_arr.ndim == 1:
                dna_arr = dna_arr[None, :]
            glm_list.append(torch.from_numpy(dna_arr).to(self.dtype))
            glm_lens.append(dna_arr.shape[0])

            # sanity: scores length should match dna length
            if len(scores) != dna_arr.shape[0]:
                raise ValueError(
                    f"Length mismatch for DNA seq: shelf length={dna_arr.shape[0]} "
                    f"but scores length={len(scores)}"
                )

        # 2) Pad sequences to batch max length
        binder_emb = pad_sequence(
            binder_list, batch_first=True, padding_value=self.pad_value
        )  # [B, Lb_max, Db]
        glm_emb = pad_sequence(
            glm_list, batch_first=True, padding_value=self.pad_value
        )  # [B, Lg_max, Dg]
        
        binder_lens = torch.as_tensor(binder_lens, dtype=torch.int64)
        glm_lens = torch.as_tensor(glm_lens, dtype=torch.int64)
        
        binder_mask = torch.arange(binder_emb.size(1)).unsqueeze(
            0
        ) < binder_lens.unsqueeze(
            1
        )  # [B, Lb_max]
        glm_mask = torch.arange(glm_emb.size(1)).unsqueeze(0) < glm_lens.unsqueeze(
            1
        )  # [B, Lg_max]

        # True = PAD  (what MHA expects)
        binder_kpm = ~binder_mask
        glm_kpm    = ~glm_mask

        # 3) Collate labels for DNA and pad
        labels_list = [torch.tensor(s, dtype=torch.float32) for s in scores_list]
        labels = pad_sequence(
            labels_list, batch_first=True, padding_value=self.pad_value
        )  # [B, Lg_max]
        
        if self.debug_run:
            max_binder_len = max(binder_lens)
            max_glm_len = max(glm_lens)
            binder_expected_false = sum(max_binder_len-binder_lens).item()
            binder_expected_true = sum(binder_lens)
            binder_expected_total = binder_expected_true + binder_expected_false
            glm_expected_false = sum(max_glm_len-glm_lens).item()
            glm_expected_true = sum(glm_lens).item()
            glm_expected_total = glm_expected_true + glm_expected_false
            labels_neg1 = sum(sum(labels==-1)).item()
            expected_labels_neg1 = glm_expected_false
            
            logger.info(f"  Max binder length: {max_binder_len}, original lengths: {binder_lens}, ultimate dimensions: {binder_emb.shape}")
            logger.info(f"  Binder expect: true/total = {binder_expected_true}/{binder_expected_total}")
            logger.info(f"  Max DNA length: {max_glm_len}, original lengths: {glm_lens}, ultimate dimensions: {glm_emb.shape}")
            logger.info(f"  DNA expect: true/total = {glm_expected_true}/{glm_expected_total}")
            logger.info(f"  Labels expect -1: -1/total = {expected_labels_neg1}/{glm_expected_total}. True: {labels_neg1}/{labels.numel()}")

        return {
            "binder_emb": binder_emb,  # [B, Lb_max, Db]
            "binder_mask": binder_mask,  # [B, Lb_max]
            "binder_kpm": binder_kpm.bool(),     # True = PAD  ← pass to MHA
            "glm_emb": glm_emb,  # [B, Lg_max, Dg]
            "glm_mask": glm_mask,  # [B, Lg_max]
            "glm_kpm": glm_kpm.bool(),     # True = PAD  ← pass to MHA
            "labels": labels,  # [B, Lg_max]
            "ID": ids,
            "tr_sequence": tr_seqs,
            "dna_sequence": dna_seqs,
        }

# ------------------------ Helpers for main method debugging only ------------------------------------------#
def _peek_batches(dl, n_batches: int = 2, tag: str = "train"):
    logger.info(f"\n=== Peek {n_batches} batch(es) from {tag} loader ===")
    for i, batch in enumerate(dl):
        be = batch["binder_emb"]
        bm = batch["binder_mask"]
        ge = batch["glm_emb"]
        gm = batch["glm_mask"]
        y = batch["labels"]
        ids = batch.get("ID", ["<no-id>"] * be.size(0))

        logger.info(f"\n[{tag}] batch {i+1}")
        logger.info(f"  binder_emb: {tuple(be.shape)}  dtype={be.dtype}")
        logger.info(f"  binder_emb: {tuple(bm.shape)}  dtype={bm.dtype}")
        logger.info(f"  binder_mask true count: {bm.sum().item()} / {bm.numel()}")
        logger.info(f"  glm_emb:    {tuple(ge.shape)}  dtype={ge.dtype}")
        logger.info(f"  glm_mask  true count: {gm.sum().item()} / {gm.numel()}")
        logger.info(f"  glm_mask: {tuple(gm.shape)}  dtype={gm.dtype}")
        logger.info(
            f"  labels:     {tuple(y.shape)}  min={y.min().item():.4f} max={y.max().item():.4f}, total -1 = {sum(sum(y==-1)).item()}"
        )
        logger.info(f"  IDs (first 5): {ids[:5]}")
        # should make sure that the number of labels that are -1 equals the number of padding tokens 
        if i + 1 >= n_batches:
            break

def _warn_on_paths(args):
    import os

    for p, label in [
        (args.train_file, "train_file"),
        (args.val_file, "val_file"),
        (args.test_file, "test_file"),
        (args.tr_shelf_path, "tr_shelf_path"),
        (args.dna_shelf_path, "dna_shelf_path"),
    ]:
        if p and not os.path.exists(p):
            logger.info(f"{label} does not exist: {p}")
    if str(args.tr_shelf_path).endswith(".pkl"):
        logger.info(
            "Warning: tr_shelf_path ends with .pkl but ShelfCollator expects a shelve DB "
            "(e.g., `.shelf`). Pass the correct path via --tr_shelf_path."
        )


def main():
    parser = argparse.ArgumentParser(
        description="Peek pre-batched, shelf-backed dataloaders"
    )
    parser.add_argument(
        "--train_file",
        type=str,
        default="../data_files/processed/splits/by_dna/babytrain.csv",
    )
    parser.add_argument(
        "--val_file",
        type=str,
        default="../data_files/processed/splits/by_dna/babyval.csv",
    )
    parser.add_argument(
        "--test_file",
        type=str,
        default="../data_files/processed/splits/by_dna/babytest.csv",
    )
    parser.add_argument(
        "--tr_shelf_path",
        type=str,
        default="../data_files/processed/embeddings/fimo_hits_only/trs_esm.shelf",
    )
    parser.add_argument(
        "--dna_shelf_path",
        type=str,
        default="../data_files/processed/embeddings/fimo_hits_only/peaks_caduceus.shelf",
    )
    parser.add_argument("--batch_size", type=int, default=4)
    parser.add_argument("--num_workers", type=int, default=4)
    parser.add_argument(
        "--debug_run", default=True, action="store_true", help="limit dataset to a few rows"
    )
    parser.add_argument(
        "--n_batches", type=int, default=2, help="how many batches to print per split"
    )
    parser.add_argument("--shuffle_train_batch_order", action="store_true")
    args = parser.parse_args()

    _warn_on_paths(args)

    dm = PairDataModule(
        train_file=args.train_file,
        val_file=args.val_file,
        test_file=args.test_file,
        tr_shelf_path=args.tr_shelf_path,
        dna_shelf_path=args.dna_shelf_path,
        batch_size=args.batch_size,
        num_workers=args.num_workers,
        debug_run=args.debug_run,
        shuffle_train_batch_order=args.shuffle_train_batch_order,
        pin_memory=False,
        score_col="binary_scores",
        norm_value=1
    )

    # ---- Train ----
    dm.setup(stage="fit")
    train_dl = dm.train_dataloader()
    _peek_batches(train_dl, n_batches=args.n_batches, tag="fit")

    # ---- Val ----
    dm.setup(stage="validate")
    val_dl = dm.val_dataloader()
    _peek_batches(val_dl, n_batches=1, tag="val")  # usually enough to sanity-check

    # ---- Test ----
    dm.setup(stage="test")
    test_dl = dm.test_dataloader()
    _peek_batches(test_dl, n_batches=1, tag="test")

    logger.info("\nAll good")


if __name__ == "__main__":
    # (Optional) set a deterministic seed for batch order shuffling
    torch.manual_seed(0)
    random.seed(0)
    np.random.seed(0)

    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s | %(levelname)s | %(name)s:%(lineno)d | %(message)s",
        datefmt="%H:%M:%S",
        handlers=[logging.StreamHandler(sys.stdout)],  # stdout, not stderr
        force=True,  # override any prior config from imported libs
    )

    main()