# Copyright 2024 ByteDance and/or its affiliates. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import random import torch from protenix.metrics.rmsd import rmsd, self_aligned_rmsd from protenix.utils.logger import get_logger from protenix.utils.permutation.chain_permutation.utils import ( apply_transform, get_optimal_transform, num_unique_matches, ) from protenix.utils.permutation.utils import Checker logger = get_logger(__name__) ExtraLabelKeys = [ "pocket_mask", "interested_ligand_mask", "chain_1_mask", "chain_2_mask", "entity_mol_id", "mol_id", "mol_atom_index", "pae_rep_atom_mask", ] def correct_symmetric_chains( pred_dict: dict, label_full_dict: dict, extra_label_keys: list[str] = ExtraLabelKeys, max_num_chains: int = 20, permute_label: bool = True, **kwargs, ): """Inputs Args: pred_dict (dict[str, torch.Tensor]): A dictionary containing: - coordinate: pred_dict["coordinate"] shape = [N_cropped_atom, 3] or [Batch, N_cropped_atom, 3]. - other keys: entity_mol_id, mol_id, mol_atom_index, pae_rep_atom_mask, is_ligand. shape = [N_cropped_atom] label_full_dict (dict[str, torch.Tensor]): A dictionary containing - coordinate: label_full_dict["coordinate"] and label_full_dict["coordinate_mask"] shape = [N_atom, 3] and [N_atom] (for coordinate_mask) - other keys: entity_mol_id, mol_id, mol_atom_index, pae_rep_atom_mask. shape = [N_atom] - extra keys: keys specified by extra_feature_keys. extra_label_keys (list[str]): - Additional features in label_full_dict that should be returned along with the permuted coordinates. max_num_chains (int): if the number of chains is more than this number, than skip permutation to avoid expensive computations. permute_label (bool): if true, permute the groundtruth chains, otherwise premute the prediction chains Return: output_dict: If permute_label=True, this is a dictionary containing - coordinate - coordinate_mask - features specified by extra_label_keys. If permute_label=False, this is a dictionary containing - coordinate. log_dict: statistics. permute_pred_indices / permute_label_indices: If batch_mode, this is a list of LongTensor. Otherwise, this is a LongTensor. The LongTensor gives the indices to permute either prediction or label. """ assert pred_dict["coordinate"].dim() in [2, 3] batch_mode = pred_dict["coordinate"].dim() > 2 if not batch_mode: ( best_match, permute_pred_indices, permute_label_indices, output_dict, log_dict, ) = _correct_symmetric_chains_for_one_sample( pred_dict, label_full_dict, max_num_chains, permute_label, extra_label_keys=extra_label_keys, **kwargs, ) return output_dict, log_dict, permute_pred_indices, permute_label_indices else: assert not permute_label, "Only supports prediction permutations in batch mode." pred_coord = [] log_dict = {} best_matches = [] permute_pred_indices = [] permute_label_indices = [] # Loop over all samples to find best matches one by one for i, pred_coord_i in enumerate(pred_dict["coordinate"]): ( best_match_i, permute_pred_indices_i, permute_label_indices_i, pred_dict_i, log_dict_i, ) = _correct_symmetric_chains_for_one_sample( {**pred_dict, "coordinate": pred_coord_i}, label_full_dict, max_num_chains, permute_label=False, extra_label_keys=[], **kwargs, ) best_matches.append(best_match_i) permute_pred_indices.append(permute_pred_indices_i) permute_label_indices.append(permute_label_indices_i) pred_coord.append(pred_dict_i["coordinate"]) for key, value in log_dict_i.items(): log_dict.setdefault(key, []).append(value) output_dict = {"coordinate": torch.stack(pred_coord, dim=0)} log_dict = {key: sum(value) / len(value) for key, value in log_dict.items()} log_dict["N_unique_perm"] = num_unique_matches(best_matches) return output_dict, log_dict, permute_pred_indices, permute_label_indices def _correct_symmetric_chains_for_one_sample( pred_dict: dict, label_full_dict: dict, max_num_chains: int = 20, permute_label: bool = False, extra_label_keys: list[str] = [], **kwargs, ): """ Correct symmetric chains for a single sample by permuting either the predicted or the ground truth coordinates. """ if not permute_label: """ Permutation will act on the predicted coordinate. In this case, predicted structures and true structure need to have the same number of atoms. """ assert pred_dict["coordinate"].size(-2) == label_full_dict["coordinate"].size( -2 ) with torch.no_grad(): # Do not compute gradient while optimizing the permutation ( best_match, permute_pred_indices, permute_label_indices, log_dict, ) = MultiChainPermutation(**kwargs)( pred_dict=pred_dict, label_full_dict=label_full_dict, max_num_chains=max_num_chains, ) if permute_label: # Permute groundtruth coord and coord mask with indices, along the first dimension. indices = permute_label_indices.tolist() output_dict = { "coordinate": label_full_dict["coordinate"][indices, :], "coordinate_mask": label_full_dict["coordinate_mask"][indices], } # Permute extra label features, along the last dimension. output_dict.update( { k: label_full_dict[k][..., indices] for k in extra_label_keys if k in label_full_dict } ) else: # Permute the predicted coord with permuted_indices indices = permute_pred_indices.tolist() output_dict = { "coordinate": pred_dict["coordinate"][indices, :], } return ( best_match, permute_pred_indices, permute_label_indices, output_dict, log_dict, ) class MultiChainPermutation(object): """Anchor-based heuristic method. Find the best match that maps predicted chains to chains in the true complex. Here the predicted chains could be cropped, which could be fewer and shorter than those in the true complex. """ def __init__( self, use_center_rmsd, find_gt_anchor_first, accept_it_as_it_is, *args, **kwargs ): self.use_center_rmsd = use_center_rmsd self.find_gt_anchor_first = find_gt_anchor_first self.accept_it_as_it_is = accept_it_as_it_is @staticmethod def dict_of_interested_keys( input_dict: dict, keys: list = [ "mol_id", "entity_mol_id", "mol_atom_index", "pae_rep_atom_mask", "coordinate", "coordinate_mask", "is_ligand", ], ): """ Extract a subset of keys from the input dictionary from the list `keys`. """ return {k: input_dict[k] for k in keys if k in input_dict} def process_input( self, pred_dict: dict[str, torch.Tensor], label_full_dict: dict[str, torch.Tensor], max_num_chains: int = 20, ): """Process the input dicts Args: pred_dict (dict[str, torch.Tensor]): A dictionary containing entity_mol_id, mol_id, mol_atom_index, pae_rep_atom_mask, coordinate, is_ligand. All Tensors have shape = [N_cropped_atom] label_full_dict (dict[str, torch.Tensor]): A dictionary containing entity_mol_id, mol_id, mol_atom_index, pae_rep_atom_mask, coordinate, coordinate_mask. All Tensors have shape = [N_atom] max_num_chains (int): if the number of chains is more than this number, than skip permutation to avoid expensive computations. permute_label (bool): if true, permute the groundtruth chains, otherwise premute the prediction chains """ log_dict = {} for key in ["entity_mol_id", "mol_id", "mol_atom_index"]: pred_dict[key] = pred_dict[key].long() label_full_dict[key] = label_full_dict[key].long() # get original unpermuted match pred_mol_id = set(torch.unique(pred_dict["mol_id"]).tolist()) label_mol_id = set(torch.unique(label_full_dict["mol_id"]).tolist()) if pred_mol_id.intersection(label_mol_id) != pred_mol_id: # if the mol_id in predicted structure is not a subset of label structure, # assert they contain the same number of atoms. assert pred_dict["coordinate"].size(-2) == label_full_dict[ "coordinate" ].size(-2) self.unpermuted_match = self.check_pattern_and_create_mapping( pred_dict["mol_id"], label_full_dict["mol_id"] ) else: self.unpermuted_match = { i: i for i in torch.unique(pred_dict["mol_id"]).tolist() } if len(torch.unique(label_full_dict["entity_mol_id"])) == len( torch.unique(label_full_dict["mol_id"]) ): # No permutation is needed has_sym_chain = False return self.unpermuted_match, has_sym_chain else: has_sym_chain = True n_label_chain = len(torch.unique(label_full_dict["mol_id"])) if n_label_chain > 20: logger.warning(f"The label_full_dict contains {n_label_chain} asym chains.") if max_num_chains > 0 and n_label_chain > max_num_chains: logger.warning( f"The label_full_dict contains {n_label_chain} asym chains (max_num_chains: {max_num_chains}). Will skip chain permutation and keep the original chain assignment." ) return self.unpermuted_match, has_sym_chain # parse features to token-level self.label_token_dict, self.label_asym_dict = self._parse_atom_feature_dict( self.dict_of_interested_keys(label_full_dict), rep_atom_mask=label_full_dict["pae_rep_atom_mask"], ) self.pred_token_dict, self.pred_asym_dict = self._parse_atom_feature_dict( self.dict_of_interested_keys(pred_dict), rep_atom_mask=pred_dict["pae_rep_atom_mask"], ) # get mapping between entity_id and asym_id self.label_token_dict.update( self._get_entity_asym_mapping( self.label_token_dict["entity_mol_id"], self.label_token_dict["mol_id"] ) ) self.pred_token_dict.update( self._get_entity_asym_mapping( self.pred_token_dict["entity_mol_id"], self.pred_token_dict["mol_id"] ) ) return None, has_sym_chain @staticmethod def check_pattern_and_create_mapping(mol_id1: torch.Tensor, mol_id2: torch.Tensor): """ Check if the patterns between two mol_id tensors match and create a mapping between them. Args: mol_id1 (torch.Tensor): A tensor of mol IDs from the first set. mol_id2 (torch.Tensor): A tensor of mol IDs from the second set. Returns: dict: A dictionary mapping mol IDs from mol_id1 to mol_id2. """ if mol_id1.shape != mol_id2.shape: raise ValueError("mol_id1 and mol_id2 must have the same shape") pattern_mapping = {} for id1, id2 in zip(mol_id1.tolist(), mol_id2.tolist()): if id1 in pattern_mapping: if pattern_mapping[id1] != id2: raise ValueError( f"Inconsistent pattern: {id1} mapped to different values in mol_id2" ) else: if id2 in pattern_mapping.values(): raise ValueError( f"Value {id2} in mol_id2 already mapped to another value" ) pattern_mapping[id1] = id2 return pattern_mapping def _parse_atom_feature_dict( self, atom_features: dict, rep_atom_mask: torch.Tensor ): """ Parse the atom feature dictionary and convert it to token features and per-asym token features. Args: atom_features (dict): A dictionary containing atom features. rep_atom_mask (torch.Tensor): The rep atom mask. Returns: tuple: A tuple containing: - token_dict (dict): A dictionary containing the token features corresponding to the rep atoms. - asym_token_dict (dict): A dictionary where keys are asym IDs and values are dictionaries of features corresponding to each asym ID. """ # Atom features --> Token features token_dict = self._convert_to_token_dict( atom_dict=atom_features, rep_atom_mask=rep_atom_mask.bool(), ) # Token features --> per asym token features asym_token_dict = self._convert_to_per_asym_feature_dict( asym_id=token_dict["mol_id"], feature_dict={ "coordinate": token_dict.get("coordinate"), "coordinate_mask": token_dict.get("coordinate_mask", None), "mol_atom_index": token_dict.get("mol_atom_index"), }, ) return token_dict, asym_token_dict @staticmethod def _convert_to_token_dict( atom_dict: dict[str, torch.Tensor], rep_atom_mask: torch.Tensor ) -> dict[str, torch.Tensor]: """ Convert the atom feature dictionary to a token feature dictionary based on the rep atom mask. Args: atom_dict (dict[str, torch.Tensor]): A dictionary containing atom features. rep_atom_mask (torch.Tensor): The rep atom mask. Returns: dict[str, torch.Tensor]: A dictionary containing the token features corresponding to the rep atoms. """ rep_atom_mask = rep_atom_mask.bool() return {k: v[rep_atom_mask] for k, v in atom_dict.items() if v is not None} @staticmethod def _convert_to_per_asym_feature_dict(asym_id: torch.Tensor, feature_dict: dict): """ Convert the feature dictionary to a dictionary where keys are asym IDs and values are dictionaries of features corresponding to each asym ID. Args: asym_id (torch.Tensor): A tensor of asym IDs. feature_dict (dict): A dictionary containing features for all atoms. Returns: dict: A dictionary where keys are asym IDs and values are dictionaries of features corresponding to each asym. """ out = {} for aid in torch.unique(asym_id): mask = asym_id == aid out[aid.item()] = { k: v[mask] for k, v in feature_dict.items() if v is not None } return out @staticmethod def _get_entity_asym_mapping( entity_id: torch.Tensor, asym_id: torch.Tensor ) -> tuple[dict]: """ Generate mappings between entity IDs and asym IDs. Args: entity_id (torch.Tensor): A tensor of entity IDs. asym_id (torch.Tensor): A tensor of asym IDs. Returns: tuple[dict]: A tuple containing two dictionaries: - entity_to_asym: A dictionary mapping entity IDs to their corresponding asym IDs. - asym_to_entity: A dictionary mapping asym IDs to their corresponding entity IDs. """ entity_to_asym = {} asym_to_entity = {} for ein in torch.unique(entity_id): ein = ein.item() asyms = torch.unique(asym_id[entity_id == ein]) entity_to_asym[ein] = asyms asym_to_entity.update({a.item(): ein for a in asyms}) return {"entity_to_asym": entity_to_asym, "asym_to_entity": asym_to_entity} def find_anchor_asym_chain_in_predictions(self) -> tuple[int]: """ Find anchor chains in the prediction. Ref: AlphaFold3 SI Chapter 4.2. -> AlphaFold Multimer Chapter 7.3.1 In the alignment phase, we pick a pair of anchor asyms to align, one in the ground truth and one in the prediction. The ground truth anchor asym a_gt is chosen to be the least ambiguous possible, for example in an A3B2 complex an arbitrary B asym is chosen. In the event of a tie e.g. A2B2 stoichiometry, the longest asym is chosen, with the hope that in general the longer asyms are likely to have higher confident predictions. The prediction anchor asym is chosen from the set {a^pred_m} of all prediction asyms with the same sequence as the ground truth anchor asym. Return: anchor_pred_asym_id (int): selected asym chain. """ # Do not consider asym with fewer than 4 tokens in Prediction asym_to_asym_length = { asym_id: len(asym_dict["coordinate"]) for asym_id, asym_dict in self.pred_asym_dict.items() } valid_asyms = [asym_id for asym_id, l in asym_to_asym_length.items() if l >= 4] # Do not consider entities with fewer than 4 resolved tokens in GT valid_entities = [] for ent, asyms in self.label_token_dict["entity_to_asym"].items(): if any( self.label_asym_dict[asym.item()]["coordinate_mask"].sum().item() >= 4 for asym in asyms ): valid_entities.append(ent) valid_entity_asym = [ (ent, asym.item()) for ent in valid_entities for asym in self.pred_token_dict["entity_to_asym"][ent] if asym.item() in valid_asyms ] candidate_entities = set(ent for ent, _ in valid_entity_asym) # Find polymer chains in the prediction pred_polymer_entity_id = [] for ent_id in candidate_entities: mask = self.pred_token_dict["entity_mol_id"] == ent_id is_ligand = self.pred_token_dict["is_ligand"][mask] if ( torch.sum(is_ligand) <= is_ligand.shape[0] / 2 and is_ligand.shape[0] >= 12 # do not prioritize asym with too few tokens ): pred_polymer_entity_id.append(ent_id) # Prioritize polymer if len(pred_polymer_entity_id) > 0: candidate_entities = pred_polymer_entity_id # Choose entities with fewest asyms in GT entity_to_asym_count = { k: len(self.label_token_dict["entity_to_asym"][k]) for k in candidate_entities } min_asym_count = min(list(entity_to_asym_count.values())) candidate_entities = [ ent for ent, count in entity_to_asym_count.items() if count == min_asym_count ] # Choose longest asyms in Prediction candidate_asyms = [ asym_id for ent, asym_id in valid_entity_asym if ent in candidate_entities ] max_asym_length = max( asym_to_asym_length[asym_id] for asym_id in candidate_asyms ) candidate_asyms = [ asym_id for asym_id in candidate_asyms if asym_to_asym_length[asym_id] == max_asym_length ] # If multiple asym chains remain, return a random one. anchor_pred_asym_id = random.choice(candidate_asyms) return anchor_pred_asym_id @staticmethod def _select_atoms_by_mol_atom_index(input_dict: dict, mol_atom_index: torch.Tensor): """ Select atoms from the input dictionary based on the specified mol_atom_index. Args: input_dict (dict): Input dict. mol_atom_index (torch.Tensor): A tensor of atom indices. Returns: dict: A dictionary containing the selected atom features. """ mask = torch.isin(input_dict["mol_atom_index"], mol_atom_index) out_dict = {k: v[mask] for k, v in input_dict.items()} assert (out_dict["mol_atom_index"] == mol_atom_index).all() return out_dict def compute_best_match_heuristic(self): """ Compute the best chain permutation between prediction and groundtruth. Returns: dict[int, int]: A dictionary mapping pred chain IDs to those of the groundtruth. """ # Find anchor asym chain in predictions anchor_pred_asym_id = self.find_anchor_asym_chain_in_predictions() anchor_entity_id = self.pred_token_dict["asym_to_entity"][anchor_pred_asym_id] if self.find_gt_anchor_first: # Randomly sample a groundtruth asym chain using this entity id anchor_gt_asym_id = self.label_token_dict["entity_to_asym"][ anchor_entity_id ].tolist() anchor_gt_asym_id = random.choice(anchor_gt_asym_id) # The candidate anchors to be matched are from prediction candidate_anchors = self.pred_token_dict["entity_to_asym"][anchor_entity_id] else: # The candidate anchors to be matched are from groundtruth candidate_anchors = self.label_token_dict["entity_to_asym"][ anchor_entity_id ] # Find best match best_rmsd = torch.inf best_match = None for anchor_k in candidate_anchors: anchor_k = anchor_k.item() if self.find_gt_anchor_first: gt_anchor, pred_anchor = anchor_gt_asym_id, anchor_k else: gt_anchor, pred_anchor = anchor_k, anchor_pred_asym_id # Find atoms in GT chain to match atoms in predicted chain (which could be cropped) gt_anchor_dict = MultiChainPermutation._select_atoms_by_mol_atom_index( self.label_asym_dict[gt_anchor], mol_atom_index=self.pred_asym_dict[pred_anchor]["mol_atom_index"], ) # Align GT Anchor to Pred Anchor mask = gt_anchor_dict["coordinate_mask"].bool() # use GT coordinate_mask if not mask.any(): continue rot, trans = get_optimal_transform( gt_anchor_dict["coordinate"][mask], self.pred_asym_dict[pred_anchor]["coordinate"][mask], ) # Transform all GT coordinates according to the aligment results aligned_coordinate = apply_transform( self.label_token_dict["coordinate"], rot, trans ) for asym_id in self.label_asym_dict: self.label_asym_dict[asym_id]["aligned_coordinate"] = ( aligned_coordinate[self.label_token_dict["mol_id"] == asym_id] ) # Greedily matches all remaining chains matched_asym = {pred_anchor: gt_anchor} to_be_matched = [k for k in self.pred_asym_dict if k != pred_anchor] candidate_gt_asym_id = [k for k in self.label_asym_dict if k != gt_anchor] # Sort the remaining chains by their length, so that longer chain chooses its match first. to_be_matched = sorted( to_be_matched, key=lambda k: -self.pred_asym_dict[k]["coordinate"].size(-2), ) while len(to_be_matched) > 0: cur_pred_asym_id = to_be_matched.pop(0) cur_entity_id = self.pred_token_dict["asym_to_entity"][cur_pred_asym_id] cur_gt_asym_ids = self.label_token_dict["entity_to_asym"][ cur_entity_id ].tolist() matched_gt_asym_id, _ = self.match_pred_asym_to_gt_asym( cur_pred_asym_id, [asym for asym in cur_gt_asym_ids if asym in candidate_gt_asym_id], ) matched_asym[cur_pred_asym_id] = matched_gt_asym_id candidate_gt_asym_id.remove(matched_gt_asym_id) assert len(matched_asym) == len(self.pred_asym_dict) # Calculate RMSD total_rmsd = self.calculate_rmsd(matched_asym) if total_rmsd < best_rmsd: best_rmsd = total_rmsd best_match = matched_asym assert best_match is not None return best_match def calculate_rmsd(self, asym_match: dict): """ Calculate the RMSD given a match. """ return sum(self._calculate_rmsd(a, b) for a, b in asym_match.items()) / len( asym_match ) def _calculate_rmsd(self, pred_asym_id: int, gt_asym_id: int): """ Calculate the RMSD between the predicted and ground truth chains, either using the average of the representative atoms or all of them. Args: pred_asym_id (int): The ID of the predicted asymmetric chain. gt_asym_id (int): The ID of the ground truth asymmetric chain. Returns: float: The calculated RMSD. """ pred_asym_dict = self.pred_asym_dict[pred_asym_id] label_asym_dict = MultiChainPermutation._select_atoms_by_mol_atom_index( self.label_asym_dict[gt_asym_id], pred_asym_dict["mol_atom_index"] ) mask = label_asym_dict["coordinate_mask"].bool() if not mask.any(): return 0.0 elif self.use_center_rmsd: return rmsd( pred_asym_dict["coordinate"][mask].mean(dim=-2, keepdim=True), label_asym_dict["aligned_coordinate"][mask].mean(dim=-2, keepdim=True), ).item() else: return rmsd( pred_asym_dict["coordinate"][mask], label_asym_dict["aligned_coordinate"][mask], ).item() def match_pred_asym_to_gt_asym(self, pred_asym_id: int, gt_asym_ids: list): """ Match a predicted chain to the groundtruth chain based on the average of the representative atoms. Args: pred_asym_id (int): The ID of the predicted asymmetric chain. gt_asym_ids (list[int]): A list or tensor of ground truth asymmetric chain IDs. Returns: tuple: A tuple containing: - best_gt_asym_id (int): The ID of the best matched ground truth asymmetric chain. - best_error (float): The distance error between the centers of mass of the best matched chains. """ pred_asym_dict = self.pred_asym_dict[pred_asym_id] best_error = torch.inf best_gt_asym_id = None unresolved_gt_asym_id = [] for gt_asym_id in gt_asym_ids: if isinstance(gt_asym_id, torch.Tensor): gt_asym_id = gt_asym_id.item() # Select cropped atoms by comparing to mol_atom_index in prediction label_asym_dict = MultiChainPermutation._select_atoms_by_mol_atom_index( self.label_asym_dict[gt_asym_id], pred_asym_dict["mol_atom_index"] ) mask = label_asym_dict["coordinate_mask"].bool() if not mask.any(): # Skip unresolved ones unresolved_gt_asym_id.append(gt_asym_id) continue gt_center = label_asym_dict["aligned_coordinate"][mask].mean(dim=0) pred_center = pred_asym_dict["coordinate"][mask].mean(dim=0) delta = torch.norm(gt_center - pred_center) if delta < best_error: best_error = delta best_gt_asym_id = gt_asym_id if best_gt_asym_id is None: # If only unresolved ones remains, return the first one assert len(unresolved_gt_asym_id) > 0 best_gt_asym_id, best_error = gt_asym_ids[0], 0 return best_gt_asym_id, best_error @staticmethod def build_permuted_indice( pred_dict: dict, label_full_dict: dict, best_match: dict[int, int] ): """ Build permutation indices from the pred-gt chain mapping. Args: pred_dict (dict): A dictionary containing the predicted coordinates. label_full_dict (dict): A dictionary containing the true coordinates and their masks. best_match (dict[int, int]): {pred_mol_id: gt_mol_id} best match between pred asym chains and gt asym chains Returns: indices (torch.Tensor): Permutation indices. """ # Get the number of predicted (cropped) atoms N_pred_atom = pred_dict["mol_id"].size(0) N_label_atom = label_full_dict["mol_id"].size(0) indices = pred_dict["mol_id"].new_zeros(size=(N_pred_atom,)) full_indices = torch.arange(N_label_atom, device=indices.device) for pred_asym_id, gt_asym_id in best_match.items(): # Create a mask for the predicted asym_id mask = pred_dict["mol_id"] == pred_asym_id mol_atom_index = pred_dict["mol_atom_index"][mask] # Creat a mask for the matched gt asym_id gt_mask = label_full_dict["mol_id"] == gt_asym_id # Extract indices according to 'mol_atom_index' gt_asym_dict = MultiChainPermutation._select_atoms_by_mol_atom_index( { "mol_atom_index": label_full_dict["mol_atom_index"][gt_mask], "indices": full_indices[gt_mask], }, mol_atom_index, ) indices[mask] = gt_asym_dict["indices"].clone() assert len(torch.unique(indices)) == len(indices) return indices @staticmethod def aligned_rmsd( pred_dict: dict, label_full_dict: dict, indices: torch.Tensor, reduce: bool = True, eps: float = 1e-8, ): """ Calculate the global aligned RMSD between predicted and true coordinates. Args: pred_dict (dict): A dictionary containing the predicted coordinates. label_full_dict (dict): A dictionary containing the true coordinates and their masks. indices (torch.Tensor): Indices to select from the true coordinates. reduce (bool): If True, reduce the RMSD over the batch dimension. eps (float): A small value to avoid division by zero. Returns: float: The aligned RMSD value. """ with torch.cuda.amp.autocast(enabled=False): aligned_rmsd, _, _, _ = self_aligned_rmsd( pred_pose=pred_dict["coordinate"].to(torch.float32), true_pose=label_full_dict["coordinate"][indices, :].to(torch.float32), atom_mask=label_full_dict["coordinate_mask"][indices], allowing_reflection=False, reduce=reduce, eps=eps, ) return aligned_rmsd.item() def __call__( self, pred_dict: dict[str, torch.Tensor], label_full_dict: dict[str, torch.Tensor], max_num_chains: int = 20, ): """ Call function for the class Args: pred_dict (dict): A dictionary containing the predicted coordinates. label_full_dict (dict): A dictionary containing the groundtruth and its attributes. max_num_chains (int): Maximum number of chains allowed. Returns: tuple: A tuple containing: - best_match (dict[int, int]): The best match between predicted and groundtruth chains. - permute_pred_indices (torch.Tensor or None): Indices to permute the predicted coordinates. - permuted_indices (torch.Tensor): Indices to permute the groundtruth coordinates. - log_dict (dict): A dictionary detailing the permutation information. """ match, has_sym_chain = self.process_input( pred_dict, label_full_dict, max_num_chains ) if match is not None: """ Either the structure does not contain symmetric chains, or there are too many chains so that the algorithm gives up. """ indices = self.build_permuted_indice(pred_dict, label_full_dict, match) pred_indices = torch.argsort(indices) return match, pred_indices, indices, {"has_sym_chain": False} # Core step: get best mol_id match best_match = self.compute_best_match_heuristic() permuted_indices = self.build_permuted_indice( pred_dict, label_full_dict, best_match ) log_dict = { "has_sym_chain": True, "is_permuted": num_unique_matches([best_match, self.unpermuted_match]) > 1, "algo:no_permute": num_unique_matches([best_match, self.unpermuted_match]) == 1, } if log_dict["algo:no_permute"]: # No permutation, return now pred_indices = torch.argsort(permuted_indices) return best_match, pred_indices, permuted_indices, log_dict # Compare rmsd before/after permutation unpermuted_indices = self.build_permuted_indice( pred_dict, label_full_dict, self.unpermuted_match ) permuted_rmsd = self.aligned_rmsd(pred_dict, label_full_dict, permuted_indices) unpermuted_rmsd = self.aligned_rmsd( pred_dict, label_full_dict, unpermuted_indices ) improved_rmsd = unpermuted_rmsd - permuted_rmsd if improved_rmsd >= 1e-12: # Case with better permutation log_dict.update( { "algo:equivalent_permute": False, "algo:worse_permute": False, "algo:better_permute": True, "algo:better_rmsd": improved_rmsd, } ) elif improved_rmsd < 0: # Case with worse permutation log_dict.update( { "algo:equivalent_permute": False, "algo:worse_permute": True, "algo:better_permute": False, "algo:worse_rmsd": -improved_rmsd, } ) elif not log_dict["algo:no_permute"]: # Case with equivalent permutation log_dict.update( { "algo:equivalent_permute": True, "algo:worse_permute": False, "algo:better_permute": False, } ) else: # No permutation log_dict["debug:zero_rmsd"] = improved_rmsd # Revert worse/equivalent permute to original chain assignment if (not self.accept_it_as_it_is) and ( log_dict["algo:equivalent_permute"] or log_dict["algo:worse_permute"] ): # Revert to original chain assignment best_match = self.unpermuted_match permuted_indices = unpermuted_indices log_dict["is_permuted"] = False if pred_dict["coordinate"].size(-2) == label_full_dict["coordinate"].size(-2): Checker.is_permutation(permuted_indices) # indices to permute/crop label permute_pred_indices = torch.argsort( permuted_indices ) # Indices to permute pred else: # Hard to `define` permute_pred_indices in this case permute_pred_indices = None return best_match, permute_pred_indices, permuted_indices, log_dict