data_utils.py

import os
import h5py

import torch
from torch.utils.data import Dataset, DataLoader
from torch.utils.data.distributed import DistributedSampler
import torch.distributed as dist
from typing import List
import random



class HDF5GenomeDataset(Dataset):
    def __init__(self, hdf5_files: List[str], bp_len: int = 249, shuffle_files: bool = True, seed: int = 42, multihead: bool = False):
        """
        Args:
            hdf5_files: List of paths to HDF5 files
            bp_len: Length to pad/truncate sequences to
            shuffle_files: Whether to shuffle file order
            seed: Random seed for reproducibility
        """
        self.hdf5_files = sorted(hdf5_files)
        self.bp_len = bp_len
        self.seed = seed
        self.multihead = multihead
        
        # Precompute file lengths and index mapping
        self.file_lengths = self._precompute_file_lengths()
        self.total_examples = sum(self.file_lengths)
        self.index_map = self._create_index_map()
        
        # One-hot encoding mapping (A,C,G,T,N)
        self.mapping = torch.tensor([
            [1, 0, 0, 0],  # A
            [0, 1, 0, 0],  # C
            [0, 0, 1, 0],  # G
            [0, 0, 0, 1],  # T
            [0, 0, 0, 0],  # N
        ], dtype=torch.float32)
        
        if shuffle_files:
            rng = random.Random(seed)
            rng.shuffle(self.index_map)
        
        # Cache open file handles for efficiency
        self.file_handles = [None] * len(self.hdf5_files)
        
    def _precompute_file_lengths(self) -> List[int]:
        """Get number of examples in each file"""
        lengths = []
        for file_path in self.hdf5_files:
            with h5py.File(file_path, 'r') as f:
                lengths.append(len(f['sequence']))
        return lengths
    
    def _create_index_map(self) -> List[tuple]:
        """Create (file_idx, local_idx) mapping"""
        return [(file_idx, i) 
                for file_idx, length in enumerate(self.file_lengths) 
                for i in range(length)]
    
    def _get_file_handle(self, file_idx: int):
        """Get or create file handle with lazy loading"""
        if self.file_handles[file_idx] is None:
            self.file_handles[file_idx] = h5py.File(self.hdf5_files[file_idx], 'r')
        return self.file_handles[file_idx]
    
    def _pad_or_truncate(self, seq: str) -> str:
        """Optimized sequence padding/truncation"""
        seq_len = len(seq)
        if seq_len < self.bp_len:
            add_seq = self.bp_len - seq_len
            left_add = random.randint(0, add_seq)
            right_add = add_seq - left_add
            left_pad = 'N' * left_add
            right_pad = 'N' * right_add
            seq = left_pad + seq + right_pad
        else:
            extra_len = seq_len - self.bp_len
            if extra_len > 0:
                add_left = random.randint(0, extra_len//2)
                add_right = extra_len - add_left
                seq = seq[add_left:seq_len - add_right]
        return seq
    
    def _one_hot_encode(self, seq: str) -> torch.Tensor:
        """Vectorized one-hot encoding"""
        indices = torch.zeros(self.bp_len, dtype=torch.long)
        for i, c in enumerate(seq):
            if c == 'A': indices[i] = 0
            elif c == 'C': indices[i] = 1
            elif c == 'G': indices[i] = 2
            elif c == 'T': indices[i] = 3
            else: indices[i] = 4  # N or invalid
        return self.mapping[indices]
    
    def __len__(self):
        return self.total_examples
    
    def __getitem__(self, idx):
        file_idx, local_idx = self.index_map[idx]
        hf = self._get_file_handle(file_idx)
        
        # Get all data in one go
        seq_bytes = hf['sequence'][local_idx]
        seq = seq_bytes.decode('utf-8').upper()
        seq = self._pad_or_truncate(seq)
        
        # Convert to tensors
        one_hot = self._one_hot_encode(seq)

        if self.multihead:
            cdna_b = torch.tensor(hf['cDNA_count_bost_boaec'][local_idx], dtype=torch.float32)
            cdna_h = torch.tensor(hf['cDNA_count_bost_haec'][local_idx], dtype=torch.float32)
            cdna = torch.stack([cdna_b, cdna_h])
        else:
            cdna = torch.tensor(hf['cDNA_count'][local_idx], dtype=torch.float32)
        
        # Get metadata
        chrom = hf['chrom'][local_idx].decode('utf-8')
        start = int(hf['start'][local_idx])
        end = int(hf['end'][local_idx])
        strand = hf['strand'][local_idx].decode('utf-8')
        # type = hf['type'][local_idx].decode('utf-8')
        
        return (
            one_hot,  # shape: (seq_len, 4)
            cdna,
            chrom,
            start,
            end,
            strand,
            # type
        )
    
    def __del__(self):
        """Clean up file handles"""
        for handle in self.file_handles:
            if handle is not None:
                handle.close()

def create_dataloader(
    hdf5_files: List[str],
    batch_size: int,
    is_train: bool = True,
    num_workers: int = min(4, os.cpu_count() // 2),
    bp_len: int = 249,
    seed: int = 42,
    prefetch_factor: int = 2,
    shuffle: bool = None,
    multihead: bool = False
):
    """
    Optimized DataLoader for HDF5 files with distributed training support.
    
    Args:
        hdf5_files: List of HDF5 file paths
        batch_size: Batch size
        is_train: Whether this is for training (affects shuffling/sampling)
        num_workers: Number of parallel workers
        bp_len: Sequence length
        seed: Random seed
        prefetch_factor: How many batches to prefetch
        shuffle: Override for shuffle behavior (None=auto based on is_train)
    """
    if shuffle is None:
        shuffle = is_train
    
    dataset = HDF5GenomeDataset(
        hdf5_files,
        bp_len=bp_len,
        shuffle_files=is_train,
        seed=seed,
        multihead=multihead
    )
    
    # Distributed sampler if needed
    sampler = None
    if dist.is_initialized():
        sampler = DistributedSampler(
            dataset,
            num_replicas=dist.get_world_size(),
            rank=dist.get_rank(),
            shuffle=shuffle,
            seed=seed,
            drop_last=is_train
        )
        shuffle = False
    
    # Optimized DataLoader settings
    loader = DataLoader(
        dataset,
        batch_size=batch_size,
        sampler=sampler,
        shuffle=shuffle,
        num_workers=num_workers,
        pin_memory=True,
        drop_last=is_train,
        persistent_workers=num_workers > 0,
        prefetch_factor=prefetch_factor if num_workers > 0 else None,
        multiprocessing_context='spawn' if num_workers > 0 else None
    )
    
    return loader, sampler