Get the protein indexer. :param rec_features: Protein features. :param edge_cutoff: Edge cutoff. :return: Protein indexer.
(rec_features: Dict[str, Any], edge_cutoff: int = 50)
| 989 | |
| 990 | |
| 991 | def get_protein_indexer(rec_features: Dict[str, Any], edge_cutoff: int = 50) -> Dict[str, Any]: |
| 992 | """Get the protein indexer. |
| 993 | |
| 994 | :param rec_features: Protein features. |
| 995 | :param edge_cutoff: Edge cutoff. |
| 996 | :return: Protein indexer. |
| 997 | """ |
| 998 | # Using a large cutoff here; dynamically remove edges along diffusion |
| 999 | res_xyzs = rec_features["res_atom_positions"] |
| 1000 | n_res = len(res_xyzs) |
| 1001 | res_atom_masks = rec_features["res_atom_mask"] |
| 1002 | ca_xyzs = res_xyzs[:, 1, :] |
| 1003 | distances = np.linalg.norm(ca_xyzs[:, np.newaxis, :] - ca_xyzs[np.newaxis, :, :], axis=2) |
| 1004 | edge_mask = distances < edge_cutoff |
| 1005 | # Mask out residues where the backbone is not resolved |
| 1006 | res_mask = np.all(~res_atom_masks[:, :3], axis=1) |
| 1007 | edge_mask[res_mask, :] = 0 |
| 1008 | edge_mask[:, res_mask] = 0 |
| 1009 | res_ids = np.broadcast_to(np.arange(n_res), (n_res, n_res)) |
| 1010 | src_nid, dst_nid = res_ids[edge_mask], res_ids.T[edge_mask] |
| 1011 | |
| 1012 | indexer = { |
| 1013 | "gather_idx_a_chainid": rec_features["res_chain_id"], |
| 1014 | "gather_idx_a_structid": np.zeros((n_res,), dtype=np.int_), |
| 1015 | "gather_idx_ab_a": src_nid, |
| 1016 | "gather_idx_ab_b": dst_nid, |
| 1017 | } |
| 1018 | return indexer |
| 1019 | |
| 1020 | |
| 1021 | def process_protein( |