model/dataset/classification_data.py

"""
https://github.com/ProteinDesignLab/protpardelle
License: MIT
Author: Alex Chu

Dataloader from PDB files.
"""
import logging
import os

import hydra
import numpy as np
import math
import torch
import torch.utils
import torch.utils.data
import tree
from omegaconf import DictConfig
import pandas as pd

from openfold_data import data_transforms
import utils.openfold_rigid_utils as rigid_utils

from utils.pdbUtils import read_pkl, parse_chain_feats

from torch.utils.data import DataLoader, Dataset
from pytorch_lightning import LightningDataModule

from sklearn.model_selection import train_test_split


def get_dataloaders(cfg):
    dataset_cfg = cfg.dataset
    loader_cfg = cfg.loader
    
    num_workers = loader_cfg.num_workers
    
    pdb_csv = pd.read_csv(dataset_cfg.csv_path)
    pdb_csv = pdb_csv[pdb_csv.modeled_seq_len <= dataset_cfg.max_num_res]
    pdb_csv = pdb_csv[pdb_csv.modeled_seq_len >= dataset_cfg.min_num_res]
    
    print(pdb_csv["class"].value_counts())
    
    train_data, test_data = train_test_split(pdb_csv, test_size=0.2, shuffle=True)
    
    train = PdbDataset(
        train_data,
        dataset_cfg,
        is_training=True
    )
    test = PdbDataset(
        test_data,
        dataset_cfg,
        is_training=False
    )
    
    train_loader = DataLoader(
        train,
        batch_size=loader_cfg.batch_size,
        num_workers=num_workers,
        prefetch_factor=None if num_workers == 0 else loader_cfg.prefetch_factor,
        pin_memory=False,
        persistent_workers=True if num_workers > 0 else False
    )
    
    val_loader = DataLoader(
            test,
            shuffle=False,
            num_workers=2,
            prefetch_factor=2,
            persistent_workers=True
        )
    
    return train_loader, val_loader


class PdbDataset(Dataset):
    def __init__(
            self,
            dataset,
            dataset_cfg,
            is_training
    ):
        self.pdb_csv = dataset
        self._log = logging.getLogger(__name__)
        self._is_training = is_training
        self._dataset_cfg = dataset_cfg
        self._init_metadata()
        self._rng = np.random.default_rng(seed=self._dataset_cfg.seed)

    @property
    def is_training(self):
        return self._is_training

    @property
    def dataset_cfg(self):
        return self._dataset_cfg

    def _init_metadata(self):
        self.pdb_csv = self.pdb_csv.sort_values('modeled_seq_len', ascending=False)
        self.csv = self.pdb_csv
        self._metadata_dir = os.path.dirname(os.path.abspath(self._dataset_cfg.csv_path))
        self._log.info(
            f'Dataset: {len(self.csv)} examples.'
        )

    def _resolve_processed_path(self, processed_file_path):
        if os.path.exists(processed_file_path):
            return processed_file_path
        if not isinstance(processed_file_path, str):
            return processed_file_path
        # Allow old absolute paths from another machine by trying local dataset dir.
        basename_path = os.path.join(self._metadata_dir, os.path.basename(processed_file_path))
        if os.path.exists(basename_path):
            return basename_path
        rel_path = os.path.join(self._metadata_dir, processed_file_path)
        if os.path.exists(rel_path):
            return rel_path
        return processed_file_path

    def _process_csv_row(self, processed_file_path):
        processed_features = read_pkl(processed_file_path)
        processed_features = parse_chain_feats(processed_features)

        modeled_idx = processed_features['modeled_idx']
        min_idx, max_idx = np.min(modeled_idx), np.max(modeled_idx)
        del processed_features['modeled_idx']
        processed_features = tree.map_structure(
            lambda x: x[min_idx:(max_idx + 1)], processed_features
        )

        chain_features = {
            'aatype': torch.tensor(processed_features['aatype']).long(),
            'all_atom_positions': torch.tensor(processed_features['atom_positions']).double(),
            'all_atom_mask': torch.tensor(processed_features['atom_mask']).double()
        }

        chain_features = data_transforms.atom37_to_frames(chain_features)
        rigids_1 = rigid_utils.Rigid.from_tensor_4x4(chain_features['rigidgroups_gt_frames'])[:, 0]
        rotmats_1 = rigids_1.get_rots().get_rot_mats()
        trans_1 = rigids_1.get_trans()
        res_idx = processed_features['residue_index']

        return {
            'aatype': chain_features['aatype'],
            'res_idx': res_idx - np.min(res_idx) + 1,
            'rotmats_1': rotmats_1,
            'trans_1': trans_1,
            'res_mask': torch.tensor(processed_features['bb_mask']).int(),
        }

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

    def __getitem__(self, idx):
        example_idx = idx
        csv_row = self.csv.iloc[example_idx]
        class_idx = csv_row["class"]
        processed_file_path = self._resolve_processed_path(csv_row['processed_path'])
        chain_features = self._process_csv_row(processed_file_path)
        chain_features['csv_idx'] = torch.ones(1, dtype=torch.long) * idx
        chain_features["class"] = torch.ones(1, dtype=torch.long) * class_idx
        return chain_features


class PdbDataModule(LightningDataModule):
    def __init__(self, data_cfg):
        super().__init__()
        self.data_cfg = data_cfg
        self.loader_cfg = data_cfg.loader
        self.dataset_cfg = data_cfg.dataset
        self.sampler_cfg = data_cfg.sampler
        
        # self.num_workers = self.loader_cfg.num_workers
        self.num_workers = 0
        
        self.train_csv, self.val_csv = self.prepare_csv()
    
    def prepare_csv(self):
        pdb_csv = pd.read_csv(self.dataset_cfg.csv_path)
        # SAMPLE SIZE ??
        pdb_csv = pdb_csv.sample(10000)
        pdb_csv = pdb_csv[pdb_csv.modeled_seq_len <= self.dataset_cfg.max_num_res]
        pdb_csv = pdb_csv[pdb_csv.modeled_seq_len >= self.dataset_cfg.min_num_res]
        
        print(pdb_csv["class"].value_counts())
        
        train_data, test_data = train_test_split(pdb_csv, test_size=0.2, shuffle=True)
        """
        # Val set
        all_lengths = np.sort(test_data.modeled_seq_len.unique())
        length_indices = (len(all_lengths) - 1) * np.linspace(
                0.0, 1.0, self.dataset_cfg.num_eval_lengths)
        length_indices = length_indices.astype(int)
        eval_lengths = all_lengths[length_indices]
        test_data = test_data[test_data.modeled_seq_len.isin(eval_lengths)]
        
        # Fix a random seed to get the same split each time.
        test_data = test_data.groupby('modeled_seq_len').sample(
            self.dataset_cfg.samples_per_eval_length, replace=True, random_state=123)
        test_data = test_data.sort_values('modeled_seq_len', ascending=False)
        """
        return train_data, test_data
        

    def setup(self, stage: str):
        self._train_dataset = PdbDataset(
            self.train_csv,
            dataset_cfg=self.dataset_cfg,
            is_training=True
        )
        
        self._valid_dataset = PdbDataset(
            self.val_csv,
            dataset_cfg=self.dataset_cfg,
            is_training=False
        )

    def train_dataloader(self, rank=None, num_replicas=None):
        return DataLoader(
            self._train_dataset,
            batch_sampler=LengthBatcher(
                sampler_cfg=self.sampler_cfg,
                metadata_csv=self._train_dataset.csv,
                rank=rank,
                num_replicas=num_replicas,
            ),
            num_workers=self.num_workers,
            prefetch_factor=None if self.num_workers == 0 else self.loader_cfg.prefetch_factor,
            pin_memory=False,
            persistent_workers=True if self.num_workers > 0 else False
        )
    
    def val_dataloader(self):
        return DataLoader(
            self._valid_dataset,
            shuffle=False,
            num_workers=2,
            prefetch_factor=2,
            persistent_workers=True
        )


class LengthBatcher:

    def __init__(
            self,
            *,
            sampler_cfg,
            metadata_csv,
            seed=123,
            shuffle=True,
            num_replicas=None,
            rank=None,
    ):
        super().__init__()
        self.sample_order = None
        self._log = logging.getLogger(__name__)
        if num_replicas is None:
            self.num_replicas = 1
            # self.num_replicas = dist.get_world_size()
        else:
            self.num_replicas = num_replicas
        if rank is None:
            self.rank = 0
            #self.rank = dist.get_rank()
        else:
            self.rank = rank

        self._sampler_cfg = sampler_cfg
        self._data_csv = metadata_csv
        # Each replica needs the same number of batches. We set the number
        # of batches to arbitrarily be the number of examples per replica.
        self._num_batches = math.ceil(len(self._data_csv) / self.num_replicas)
        self._data_csv['index'] = list(range(len(self._data_csv)))
        self.seed = seed
        self.shuffle = shuffle
        self.epoch = 0
        self.max_batch_size = self._sampler_cfg.max_batch_size
        self._log.info(f'Created dataloader rank {self.rank + 1} out of {self.num_replicas}')

    def _replica_epoch_batches(self):
        # Make sure all replicas share the same seed on each epoch.
        rng = torch.Generator()
        rng.manual_seed(self.seed + self.epoch)
        if self.shuffle:
            indices = torch.randperm(len(self._data_csv), generator=rng).tolist()
        else:
            indices = list(range(len(self._data_csv)))

        if len(self._data_csv) > self.num_replicas:
            replica_csv = self._data_csv.iloc[
                indices[self.rank::self.num_replicas]
            ]
        else:
            replica_csv = self._data_csv

        # Each batch contains multiple proteins of the same length.
        sample_order = []
        for seq_len, len_df in replica_csv.groupby('modeled_seq_len'):
            max_batch_size = min(
                self.max_batch_size,
                self._sampler_cfg.max_num_res_squared // seq_len ** 2 + 1,
            )
            num_batches = math.ceil(len(len_df) / max_batch_size)
            for i in range(num_batches):
                batch_df = len_df.iloc[i * max_batch_size:(i + 1) * max_batch_size]
                batch_indices = batch_df['index'].tolist()
                sample_order.append(batch_indices)

        # Remove any length bias.
        new_order = torch.randperm(len(sample_order), generator=rng).numpy().tolist()
        return [sample_order[i] for i in new_order]

    def _create_batches(self):
        # Make sure all replicas have the same number of batches Otherwise leads to bugs.
        # See bugs with shuffling https://github.com/Lightning-AI/lightning/issues/10947
        all_batches = []
        num_augments = -1
        while len(all_batches) < self._num_batches:
            all_batches.extend(self._replica_epoch_batches())
            num_augments += 1
            if num_augments > 1000:
                raise ValueError('Exceeded number of augmentations.')
        if len(all_batches) >= self._num_batches:
            all_batches = all_batches[:self._num_batches]
        self.sample_order = all_batches

    def __iter__(self):
        self._create_batches()
        self.epoch += 1
        return iter(self.sample_order)

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




@hydra.main(version_base=None, config_path="./configs", config_name="classifier.yaml")
def my_app(cfg: DictConfig) -> None:
    train, test = get_dataloaders(cfg.data)
    print(next(iter(train)))
    print(test.shape)
    """
    data = PdbClfDataModule(cfg.data)
    data.setup('train')
    train_loader = data.train_dataloader()
    val_loader = data.val_dataloader()
    # data = PdbDataset(dataset_cfg=cfg.data.dataset, is_training=True)
    print(train_loader)
    print(val_loader)
    #print(data[0])
    """


if __name__ == '__main__':
    my_app()