Load the dataset from pdb files. Parameters: pdb_files (list of str): pdb file names transform (Callable, optional): protein sequence transformation function lazy (bool, optional): if lazy mode is used, the proteins are processed in the dataloade
(self, pdb_files, transform=None, lazy=False, verbose=0, sanitize=True, removeHs=True, **kwargs)
| 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: |
| 79 | pdb_files = tqdm(pdb_files, "Constructing proteins from pdbs") |
| 80 | for i, pdb_file in enumerate(pdb_files): |
| 81 | if not lazy or i == 0: |
| 82 | mol = Chem.MolFromPDBFile(pdb_file, sanitize=sanitize, removeHs=removeHs) |
| 83 | if not mol: |
| 84 | logger.debug("Can't construct molecule from pdb file `%s`. Ignore this sample." % pdb_file) |
| 85 | continue |
| 86 | protein = data.Protein.from_molecule(mol, **kwargs) |
| 87 | if not protein: |
| 88 | logger.debug("Can't construct protein from pdb file `%s`. Ignore this sample." % pdb_file) |
| 89 | continue |
| 90 | else: |
| 91 | protein = None |
| 92 | if hasattr(protein, "residue_feature"): |
| 93 | with protein.residue(): |
| 94 | protein.residue_feature = protein.residue_feature.to_sparse() |
| 95 | self.data.append(protein) |
| 96 | self.pdb_files.append(pdb_file) |
| 97 | self.sequences.append(protein.to_sequence() if protein else None) |
| 98 | |
| 99 | def get_item(self, index): |
| 100 | if getattr(self, "lazy", False): |