Method get_item

diffpack/dataset.py:99–134

(self, index)

            self.sequences.append(protein.to_sequence() if protein else None)

    def get_item(self, index):
        if getattr(self, "lazy", False):
            protein = data.Protein.from_pdb(self.pdb_files[index], **self.kwargs)
        else:
            protein = self.data[index].clone()
        protein = protein.subgraph(protein.atom_name < 37)

        with protein.atom():
            # Init atom14 index map
            protein.atom14index = rotamer.restype_atom14_index_map[
                protein.residue_type[protein.atom2residue], protein.atom_name
            ]  # [num_atom, 14]

        with protein.residue():
            # Init residue features
            protein.residue_feature = functional.one_hot(protein.residue_type, 21)  # [num_residue, 21]

            # Init residue masks
            chi_mask = get_chi_mask(protein)
            chi_1pi_periodic_mask = torch.tensor(rotamer.chi_pi_periodic)[protein.residue_type]
            chi_2pi_periodic_mask = ~chi_1pi_periodic_mask
            protein.chi_mask = chi_mask
            protein.chi_1pi_periodic_mask = torch.logical_and(chi_mask, chi_1pi_periodic_mask)  # [num_residue, 4]
            protein.chi_2pi_periodic_mask = torch.logical_and(chi_mask, chi_2pi_periodic_mask)  # [num_residue, 4]

            # Init atom37 features
            protein.atom37_mask = torch.zeros(protein.num_residue, len(atom_name_vocab), device=protein.device,
                                              dtype=torch.bool)  # [num_residue, 37]
            protein.atom37_mask[protein.atom2residue, protein.atom_name] = True
            protein.sidechain37_mask = protein.atom37_mask.clone()  # [num_residue, 37]
            protein.sidechain37_mask[:, bb_atom_name] = False
        item = {"graph": protein}

        if self.transform:
            item = self.transform(item)
        return item

    @staticmethod
    def from_pdb_files(pdb_files, verbose=1, **kwargs):