dataset.py

# dataset.py
# Dataset loader for protein fragment distance maps

import os
import torch
import numpy as np
import pandas as pd
from torch.utils.data import Dataset, DataLoader
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt

class ProteinFragmentDataset(Dataset):
    """
    Dataset for protein fragment distance maps.
    """
    def __init__(self, root_dir, split='train', transform=None, visualize_samples=False):
        """
        Args:
            root_dir (str): Directory with the dataset.
            split (str): 'train', 'val', or 'test'.
            transform (callable, optional): Optional transform to be applied on a sample.
            visualize_samples (bool): Whether to visualize some samples for debugging.
        """
        self.root_dir = root_dir
        self.split_dir = os.path.join(root_dir, split)
        self.transform = transform
        self.visualize_samples = visualize_samples
        
        # Load metadata
        metadata_file = os.path.join(root_dir, 'metadata.csv')
        self.metadata = pd.read_csv(metadata_file)
        self.metadata = self.metadata[self.metadata['split'] == split]
        
        # Reset index
        self.metadata = self.metadata.reset_index(drop=True)
        
        # Verify that files exist
        self.valid_indices = []
        for idx, row in self.metadata.iterrows():
            if os.path.exists(row['filename']):
                self.valid_indices.append(idx)
        
        print(f"Found {len(self.valid_indices)} valid samples for {split} split")
        
        # Visualize a few samples if requested
        if visualize_samples:
            self.visualize_random_samples(5)
    
    def __len__(self):
        return len(self.valid_indices)
    
    def __getitem__(self, idx):
        """
        Args:
            idx (int): Index
            
        Returns:
            dict: A dictionary containing:
                'distance_map': The original distance map
                'masked_map': The masked distance map (input)
                'mask': The binary mask (1 = visible, 0 = masked)
        """
        # Get metadata for the sample
        idx = self.valid_indices[idx]
        row = self.metadata.iloc[idx]
        
        # Load the .npz file
        data = np.load(row['filename'])
        
        # Extract distance map and mask
        distance_map = data['distance_map'].astype(np.float32)
        masked_map = data['masked_map'].astype(np.float32) if 'masked_map' in data else distance_map.copy()
        mask = data['mask'].astype(np.float32) if 'mask' in data else np.ones_like(distance_map, dtype=np.float32)
        
        # Add channel dimension if missing
        if len(distance_map.shape) == 2:
            distance_map = distance_map[np.newaxis, :, :]
            masked_map = masked_map[np.newaxis, :, :]
            mask = mask[np.newaxis, :, :]
        
        # Convert to torch tensors
        distance_map = torch.from_numpy(distance_map)
        masked_map = torch.from_numpy(masked_map)
        mask = torch.from_numpy(mask)
        
        # Apply transforms if specified
        if self.transform:
            distance_map = self.transform(distance_map)
            masked_map = self.transform(masked_map)
        
        # Get sequence if available
        sequence = row['sequence'] if 'sequence' in row else ""
        
        # Create sample dictionary
        sample = {
            'distance_map': distance_map,
            'masked_map': masked_map,
            'mask': mask,
            'pdb_id': row['pdb_id'],
            'chain_id': row['chain_id'],
            'sequence': sequence
        }
        
        return sample
    
    def visualize_random_samples(self, num_samples=5):
        """
        Visualize random samples from the dataset for debugging.
        """
        if len(self.valid_indices) == 0:
            print("No valid samples to visualize")
            return
            
        indices = np.random.choice(len(self.valid_indices), size=min(num_samples, len(self.valid_indices)), replace=False)
        
        fig, axes = plt.subplots(num_samples, 2, figsize=(10, 2.5 * num_samples))
        if num_samples == 1:
            axes = axes.reshape(1, -1)
            
        for i, idx in enumerate(indices):
            sample = self[idx]
            distance_map = sample['distance_map'][0].numpy()  # Remove channel dim
            masked_map = sample['masked_map'][0].numpy()
            
            axes[i, 0].imshow(distance_map, cmap='viridis')
            axes[i, 0].set_title(f"Original - {sample['pdb_id']} Chain {sample['chain_id']}")
            
            axes[i, 1].imshow(masked_map, cmap='viridis')
            axes[i, 1].set_title("Masked Input")
            
        plt.tight_layout()
        plt.savefig(os.path.join(self.root_dir, f"{self.split_dir.split('/')[-1]}_samples.png"))
        plt.close()
        
def get_dataloaders(root_dir, batch_size=32, num_workers=4, visualize_samples=True):
    """
    Create dataloaders for training, validation, and testing.
    
    Args:
        root_dir (str): Directory with the dataset.
        batch_size (int): Batch size.
        num_workers (int): Number of workers for data loading.
        visualize_samples (bool): Whether to visualize some samples for debugging.
        
    Returns:
        dict: A dictionary containing dataloaders for 'train', 'val', and 'test' splits.
    """
    # Create datasets
    train_dataset = ProteinFragmentDataset(root_dir, split='train', visualize_samples=visualize_samples)
    val_dataset = ProteinFragmentDataset(root_dir, split='val', visualize_samples=visualize_samples)
    test_dataset = ProteinFragmentDataset(root_dir, split='test', visualize_samples=visualize_samples)
    
    # Create dataloaders
    train_loader = DataLoader(
        train_dataset, 
        batch_size=batch_size, 
        shuffle=True, 
        num_workers=num_workers,
        pin_memory=True
    )
    
    val_loader = DataLoader(
        val_dataset, 
        batch_size=batch_size, 
        shuffle=False, 
        num_workers=num_workers,
        pin_memory=True
    )
    
    test_loader = DataLoader(
        test_dataset, 
        batch_size=batch_size, 
        shuffle=False, 
        num_workers=num_workers,
        pin_memory=True
    )
    
    return {
        'train': train_loader,
        'val': val_loader,
        'test': test_loader
    }