| 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 | """ |