| 19 | |
| 20 | @R.register("datasets.SideChainDataset") |
| 21 | class SideChainDataset(data.ProteinDataset): |
| 22 | processed_file = None |
| 23 | exclude_pdb_files = [] |
| 24 | |
| 25 | def __init__(self, path=None, pdb_files=None, verbose=1, **kwargs): |
| 26 | if path is not None: |
| 27 | logger.info("Loading dataset from folder %s" % path) |
| 28 | path = os.path.expanduser(path) |
| 29 | if not os.path.exists(path): |
| 30 | os.makedirs(path) |
| 31 | self.path = path |
| 32 | pkl_file = os.path.join(path, self.processed_file) |
| 33 | |
| 34 | if os.path.exists(pkl_file): |
| 35 | logger.info("Found existing pickle file %s" % pkl_file |
| 36 | + ". Loading from pickle file (this may take a while)") |
| 37 | self.load_pickle(pkl_file, verbose=verbose, **kwargs) |
| 38 | else: |
| 39 | logger.info("No pickle file found. Loading from pdb files (this may take a while)" |
| 40 | + " and save to pickle file %s" % pkl_file) |
| 41 | pdb_files = sorted(glob.glob(os.path.join(path, "*.pdb"))) |
| 42 | self.load_pdbs(pdb_files, verbose=verbose, **kwargs) |
| 43 | self.save_pickle(pkl_file, verbose=verbose) |
| 44 | elif pdb_files is not None: |
| 45 | logger.info("Loading dataset from pdb files") |
| 46 | pdb_files = [os.path.expanduser(pdb_file) for pdb_file in pdb_files] |
| 47 | pdb_files = [pdb_file for pdb_file in pdb_files if pdb_file.endswith(".pdb")] |
| 48 | self.load_pdbs(pdb_files, verbose=verbose, **kwargs) |
| 49 | |
| 50 | # Filter out proteins with no residues |
| 51 | indexes = [i for i, (protein, pdb_file) in enumerate(zip(self.data, self.pdb_files)) |
| 52 | if (protein.num_residue > 0).all() and os.path.basename(pdb_file) not in self.exclude_pdb_files] |
| 53 | self.data = [self.data[i] for i in indexes] |
| 54 | self.sequences = [self.sequences[i] for i in indexes] |
| 55 | self.pdb_files = [self.pdb_files[i] for i in indexes] |
| 56 | |
| 57 | def load_pdbs(self, pdb_files, transform=None, lazy=False, verbose=0, sanitize=True, removeHs=True, **kwargs): |
| 58 | """ |
| 59 | Load the dataset from pdb files. |
| 60 | |
| 61 | Parameters: |
| 62 | pdb_files (list of str): pdb file names |
| 63 | transform (Callable, optional): protein sequence transformation function |
| 64 | lazy (bool, optional): if lazy mode is used, the proteins are processed in the dataloader. |
| 65 | This may slow down the data loading process, but save a lot of CPU memory and dataset loading time. |
| 66 | verbose (int, optional): output verbose level |
| 67 | **kwargs |
| 68 | """ |
| 69 | num_sample = len(pdb_files) |
| 70 | |
| 71 | self.transform = transform |
| 72 | self.lazy = lazy |
| 73 | self.kwargs = kwargs |
| 74 | self.data = [] |
| 75 | self.pdb_files = [] |
| 76 | self.sequences = [] |
| 77 | |
| 78 | if verbose: |