Process protein data. :param af_protein: FDProtein object. :param bounding_box: Bounding box to filter atoms. :param no_indexer: If True, the indexer is not added to the output. :param sample_name: Name of the sample. :param prefix_key: Prefix key. :param sequences_to_embedd
(
af_protein: FDProtein,
bounding_box: Optional[np.ndarray] = None,
no_indexer: bool = True,
sample_name: str = "",
sequences_to_embeddings: Optional[Dict[str, np.ndarray]] = None,
plddt: Optional[Iterable[float]] = None,
chain_id: Optional[str] = None,
)
| 1019 | |
| 1020 | |
| 1021 | def process_protein( |
| 1022 | af_protein: FDProtein, |
| 1023 | bounding_box: Optional[np.ndarray] = None, |
| 1024 | no_indexer: bool = True, |
| 1025 | sample_name: str = "", |
| 1026 | sequences_to_embeddings: Optional[Dict[str, np.ndarray]] = None, |
| 1027 | plddt: Optional[Iterable[float]] = None, |
| 1028 | chain_id: Optional[str] = None, |
| 1029 | ) -> Dict[str, Any]: |
| 1030 | """Process protein data. |
| 1031 | |
| 1032 | :param af_protein: FDProtein object. |
| 1033 | :param bounding_box: Bounding box to filter atoms. |
| 1034 | :param no_indexer: If True, the indexer is not added to the output. |
| 1035 | :param sample_name: Name of the sample. |
| 1036 | :param prefix_key: Prefix key. |
| 1037 | :param sequences_to_embeddings: Optional dictionary mapping sequences to embeddings. |
| 1038 | :param plddt: Optional pLDDT values. |
| 1039 | :param chain_id: Optional chain ID for parsing of LM embeddings. |
| 1040 | :return: Processed protein data. |
| 1041 | """ |
| 1042 | lm_embeddings = None if sequences_to_embeddings is None else [] |
| 1043 | if sequences_to_embeddings is not None: |
| 1044 | if chain_id is not None and len(af_protein.letter_sequences) == 1: |
| 1045 | chain_seq = af_protein.letter_sequences[0][1] |
| 1046 | chain_mask = af_protein.letter_sequences[0][2] |
| 1047 | chain_seq_masked = "".join(np.array(list(chain_seq))[chain_mask]) |
| 1048 | lm_embeddings.append(sequences_to_embeddings[chain_seq_masked + f":{chain_id}"]) |
| 1049 | else: |
| 1050 | for i, (_, chain_seq, chain_mask) in enumerate(af_protein.letter_sequences): |
| 1051 | chain_seq_masked = "".join(np.array(list(chain_seq))[chain_mask]) |
| 1052 | if i in sequences_to_embeddings: |
| 1053 | lm_embeddings.append(sequences_to_embeddings[i]) |
| 1054 | elif chain_seq_masked + f":{i}" in sequences_to_embeddings: |
| 1055 | lm_embeddings.append(sequences_to_embeddings[chain_seq_masked + f":{i}"]) |
| 1056 | else: |
| 1057 | raise ValueError( |
| 1058 | f"Sequence {chain_seq_masked}:{i} not found in the provided embeddings." |
| 1059 | ) |
| 1060 | lm_embeddings = np.concatenate(lm_embeddings, axis=0) |
| 1061 | assert len(lm_embeddings) == len( |
| 1062 | af_protein.aatype |
| 1063 | ), f"LM sequence length must match OpenFold-parsed sequence length: {len(lm_embeddings)} != {len(af_protein.aatype)}" |
| 1064 | if bounding_box: |
| 1065 | raise NotImplementedError |
| 1066 | ca_pos = af_protein.atom_positions[:, 1] |
| 1067 | ca_in_box = np.all( |
| 1068 | (ca_pos > bounding_box[0]) & (ca_pos < bounding_box[1]), |
| 1069 | axis=1, |
| 1070 | ) |
| 1071 | af_protein.atom_positions = af_protein.atom_positions[ca_in_box] |
| 1072 | af_protein.aatype = af_protein.aatype[ca_in_box] |
| 1073 | af_protein.atomtypes = af_protein.atomtypes[ca_in_box] |
| 1074 | af_protein.atom_mask = af_protein.atom_mask[ca_in_box] |
| 1075 | af_protein.chain_index = af_protein.chain_index[ca_in_box] |
| 1076 | af_protein.b_factors = af_protein.b_factors[ca_in_box] |
| 1077 | chain_seqs = [ |
| 1078 | (sample_name + seq_data[0], seq_data[1]) for seq_data in af_protein.letter_sequences |
no test coverage detected