""" 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 torch import math 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 pytorch_lightning import LightningDataModule from torch.utils.data import DataLoader, Dataset from torch.utils.data.distributed import DistributedSampler, dist 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 def setup(self, stage: str): self._train_dataset = PdbDataset( dataset_cfg=self.dataset_cfg, is_training=True, ) self._valid_dataset = PdbDataset( dataset_cfg=self.dataset_cfg, is_training=False ) def train_dataloader(self, rank=None, num_replicas=None): num_workers = self.loader_cfg.num_workers 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=num_workers, prefetch_factor=None if num_workers == 0 else self.loader_cfg.prefetch_factor, pin_memory=False, persistent_workers=True if num_workers > 0 else False ) def val_dataloader(self): return DataLoader( self._valid_dataset, #sampler=DistributedSampler(self._valid_dataset, shuffle=False), shuffle=False, num_workers=2, prefetch_factor=2, persistent_workers=True, ) class PdbDataset(Dataset): def __init__( self, *, dataset_cfg, is_training ): 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): pdb_csv = pd.read_csv(self._dataset_cfg.csv_path) self._metadata_dir = os.path.dirname(os.path.abspath(self._dataset_cfg.csv_path)) self.raw_csv = pdb_csv 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] if self.dataset_cfg.subset is not None: pdb_csv = pdb_csv.iloc[:self.dataset_cfg.subset] pdb_csv = pdb_csv.sort_values('modeled_seq_len', ascending=False) if self.is_training: self.csv = pdb_csv self._log.info( f'Training: {len(self.csv)} examples.' ) else: eval_csv = pdb_csv[pdb_csv.modeled_seq_len <= self.dataset_cfg.min_eval_length] all_lengths = np.sort(eval_csv.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] eval_csv = eval_csv[eval_csv.modeled_seq_len.isin(eval_lengths)] # Fix a random seed to get the same split each time. eval_csv = eval_csv.groupby('modeled_seq_len').sample( self.dataset_cfg.samples_per_eval_length, replace=True, random_state=123) eval_csv = eval_csv.sort_values('modeled_seq_len', ascending=False) self.csv = eval_csv self._log.info( f'Validation: {len(self.csv)} examples with lengths {eval_lengths}' ) 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 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="../", config_name="train_config") def my_app(cfg: DictConfig) -> None: data = PdbDataset(dataset_cfg=cfg.data.dataset, is_training=True) print(len(data)) print(data[0]) if __name__ == '__main__': my_app()