| 138 | |
| 139 | |
| 140 | def load_protein(data_dict): |
| 141 | # Load pickled protein structure |
| 142 | atom_mask = torch.tensor(data_dict['atom_mask']).bool() |
| 143 | atom_type = atom_type_mapping[None, :] |
| 144 | atom_type = atom_type.expand_as(atom_mask)[atom_mask] |
| 145 | atom_name = atom_name_mapping[None, :] |
| 146 | atom_name = atom_name.expand_as(atom_mask)[atom_mask] |
| 147 | node_position = torch.tensor(data_dict['atom_positions'])[atom_mask] |
| 148 | residue_type = torch.tensor(data_dict['aatype']) |
| 149 | residue_type = residue_type_mapping[residue_type] |
| 150 | residue_number = torch.tensor(data_dict['residue_index']) |
| 151 | b_factor = torch.tensor(data_dict['b_factors'])[atom_mask] |
| 152 | chain_id = torch.tensor(data_dict['chain_index']) |
| 153 | num_residue = residue_type.shape[0] |
| 154 | num_atom = atom_name.shape[0] |
| 155 | |
| 156 | atom2residue = torch.arange(num_residue)[:, None] |
| 157 | atom2residue = atom2residue.expand_as(atom_mask)[atom_mask] |
| 158 | |
| 159 | edge_list = torch.zeros((1, 3), dtype=torch.long) |
| 160 | bond_type = torch.zeros((1,), dtype=torch.long) |
| 161 | |
| 162 | residue_feature = F.one_hot(residue_type, len(residue_constants.restypes_with_x)) |
| 163 | atom_feature = torch.cat([ |
| 164 | F.one_hot(atom_name, residue_constants.atom_type_num), |
| 165 | residue_feature[atom2residue] |
| 166 | ], dim=-1) |
| 167 | |
| 168 | protein = data.Protein(edge_list=edge_list, atom_type=atom_type, bond_type=bond_type, |
| 169 | residue_type=residue_type, atom_name=atom_name, atom2residue=atom2residue, |
| 170 | residue_feature=residue_feature, atom_feature=atom_feature, bond_feature=None, |
| 171 | residue_number=residue_number, b_factor=b_factor, chain_id=chain_id, |
| 172 | node_position=node_position, num_node=num_atom, num_residue=num_residue, |
| 173 | ) |
| 174 | return protein |
| 175 | |
| 176 | |
| 177 | def load_surface(surf_dict): |