Merge pull request #406 from dingquanyu/update-permutation-unittest

Update multi-chain permutation and permutation unittest 

openfold/utils/multi_chain_permutation.py

 import logging
 import random
 import torch
-
+from typing import Tuple, List, Dict
 from openfold.np import residue_constants as rc
 
 logger = logging.getLogger(__name__)
@@ -13,6 +13,17 @@ def compute_rmsd(
     atom_mask: torch.Tensor = None,
     eps: float = 1e-6,
 ) -> torch.Tensor:
+    """
+    Function to calculate RMSD between predicted and ground truth atom position
+
+    Args: 
+        true_atom_pos: a [nres*3] tensor
+        pred_atom_pos: a [nres*3] tensor
+        atom_mask: a [1*nres] tensor
+
+    Return:
+        RMSD value between true and predicted atom positions
+    """
     sq_diff = torch.square(true_atom_pos - pred_atom_pos).sum(dim=-1, keepdim=False)
     if atom_mask is not None:
         sq_diff = torch.masked_select(sq_diff, atom_mask.to(sq_diff.device))
@@ -21,7 +32,7 @@ def compute_rmsd(
     return torch.sqrt(msd + eps)  # prevent sqrt 0
 
 
-def kabsch_rotation(P, Q):
+def kabsch_rotation(P: torch.Tensor, Q: torch.Tensor) -> torch.Tensor:
     """
     Calculate the best rotation that minimises the RMSD between P and Q.
 
@@ -29,11 +40,11 @@ def kabsch_rotation(P, Q):
     https://en.wikipedia.org/wiki/Kabsch_algorithm
 
     Args:
-    P: [N * 3] Nres is the number of atoms and each row corresponds to the atom's x,y,z coordinates
-    Q: [N * 3] the same dimension as P
+        P: [N * 3] Nres is the number of atoms and each row corresponds to the atom's x,y,z coordinates
+        Q: [N * 3] the same dimension as P
 
     return:
-    A 3*3 rotation matrix
+        one 3*3 rotation matrix that best aligns the sorce and target atoms 
     """
     assert P.shape == torch.Size([Q.shape[0], Q.shape[1]])
 
@@ -54,11 +65,20 @@ def get_optimal_transform(
     src_atoms: torch.Tensor,
     tgt_atoms: torch.Tensor,
     mask: torch.Tensor = None,
-):
+) -> Tuple[torch.Tensor, torch.Tensor]:
     """
-    src_atoms: predicted CA positions, shape:[num_res,3]
-    tgt_atoms: ground-truth CA positions, shape:[num_res,3]
-    mask: a vector of boolean values, shape:[num_res]
+    A function that obtain the transformation that optimally align 
+    src_atoms with tgt_atoms 
+
+    Args:
+        src_atoms: predicted CA positions, shape:[num_res,3]
+        tgt_atoms: ground-truth CA positions, shape:[num_res,3]
+        mask: a vector of boolean values, shape:[num_res]
+    
+    Returns:
+    a rotation matrix that record the optimal rotation 
+    that will best align selected anchor prediction to selected anchor truth
+    a matrix records how the atoms should be shifted after applying r i.e. optimal alignment requires 1) rotate 2) shift the positions 
     """
     assert src_atoms.shape == tgt_atoms.shape, (src_atoms.shape, tgt_atoms.shape)
     assert src_atoms.shape[-1] == 3
@@ -88,7 +108,7 @@ def get_optimal_transform(
     return r, x
 
 
-def get_least_asym_entity_or_longest_length(batch, input_asym_id):
+def get_least_asym_entity_or_longest_length(batch: dict, input_asym_id: list) -> Tuple[torch.Tensor, List[torch.Tensor]]:
     """
     First check how many subunit(s) one sequence has. Select the subunit that is less
     common, e.g. if the protein was AABBB then select one of the A as anchor
@@ -97,15 +117,15 @@ def get_least_asym_entity_or_longest_length(batch, input_asym_id):
     then choose one of the corresponding subunits as anchor
 
     Args:
-    batch: in this function batch is the full ground truth features
-    input_asym_id: A list of asym_ids that are in the cropped input features
+        batch: in this function batch is the full ground truth features
+        input_asym_id: A list of asym_ids that are in the cropped input features
 
     Return:
-    anchor_gt_asym_id: Tensor(int) selected ground truth asym_id
-    anchor_pred_asym_ids: list(Tensor(int)) a list of all possible pred anchor candidates
+        anchor_gt_asym_id: Tensor(int) selected ground truth asym_id
+        anchor_pred_asym_ids: list(Tensor(int)) a list of all possible pred anchor candidates
     """
     entity_2_asym_list = get_entity_2_asym_list(batch)
-    unique_entity_ids = torch.unique(batch["entity_id"])
+    unique_entity_ids = [i for i in torch.unique(batch["entity_id"]) if i !=0]# if entity_id is 0, that means this entity_id comes from padding
     entity_asym_count = {}
     entity_length = {}
 
@@ -145,19 +165,38 @@ def get_least_asym_entity_or_longest_length(batch, input_asym_id):
 
 
 def greedy_align(
-    batch,
-    per_asym_residue_index,
-    entity_2_asym_list,
-    pred_ca_pos,
-    pred_ca_mask,
-    true_ca_poses,
-    true_ca_masks,
-):
+    batch: dict,
+    per_asym_residue_index: dict,
+    entity_2_asym_list: dict,
+    pred_ca_pos: torch.Tensor,
+    pred_ca_mask: torch.Tensor,
+    true_ca_poses: list,
+    true_ca_masks: list
+) -> List[Tuple[int, int]]:
     """
     Implement Algorithm 4 in the Supplementary Information of AlphaFold-Multimer paper:
     Evans,R et al., 2022 Protein complex prediction with AlphaFold-Multimer, bioRxiv 2021.10.04.463034; doi: https://doi.org/10.1101/2021.10.04.463034
+    
+    Args:
+        batch: a dictionary of ground truth features 
+        per_asym_residue_index: a dictionary recording which residues belong to which aysm_id 
+        entity_2_asym_list: a dictionary recording which asym_id(s) belong to which entity_id 
+        pred_ca_pos: predicted positions of c-alpha atoms from the results of model.forward()
+        pred_ca_mask: a boolean tensor that masks pred_ca_pos
+        true_ca_poses: a list of tensors, corresponding to the c-alpha positions of the ground truth structure. e.g. If there are 5 chains, this list will have a length of 5 
+        true_ca_masks: a list of tensors, corresponding to the masks of c-alpha positions of the ground truth structure. If there are 5 chains, this list will have a length of 5 
+    
+    Return:
+        A list of tuple(int,int) that provides instructions of how the ground truth chains should be permuated 
+        e.g. if 3 chains in the imput model have the same sequences, an example return would be:
+        [(0,2),(1,1),(2,0)], meaning the 1st chain in the predicted structure should be aligned to the 3rd chain in the ground truth,
+        and the 2nd chain in the predicted structure is ok to stay with the 2nd chain in the ground truth.
+
+    Note: the tuples in the returned list begin with 0 indexing but aym_id begins with 1. The reason why tuples in the return are 0-indexing 
+    is that at the stage of loss calculation, the ground truth atom positions: true_ca_poses, are already split up into a list of matrices.
+    Hence, now this function needs to return tuples that provide the index to select from the list: true_ca_poses, and list index starts from 0.
     """
-    used = [False for _ in range(len(true_ca_poses))]
+    used = [False for _ in range(len(true_ca_poses))] # a list the keeps recording whether a ground truth chain has been used or not
     align = []
     unique_asym_ids = [i for i in torch.unique(batch["asym_id"]) if i != 0]
     for cur_asym_id in unique_asym_ids:
@@ -189,21 +228,38 @@ def greedy_align(
     return align
 
 
-def pad_features(feature_tensor, nres_pad, pad_dim):
-    """Pad input feature tensor"""
+def pad_features(feature_tensor: torch.Tensor, nres_pad: int, pad_dim: int) -> torch.Tensor:
+    """
+    Pad input feature tensor. Padding values will be 0 and put behind the true feature values
+    
+    Args:
+        feature_tensor: A feature tensor 
+        nres_pad: number of residues to add
+        pad_dim: along which dimension of the feature_tensor to pad 
+    
+    Returns:
+        a padded feature tensor
+    """
     pad_shape = list(feature_tensor.shape)
     pad_shape[pad_dim] = nres_pad
     padding_tensor = feature_tensor.new_zeros(pad_shape, device=feature_tensor.device)
     return torch.concat((feature_tensor, padding_tensor), dim=pad_dim)
 
 
-def merge_labels(per_asym_residue_index, labels, align, original_nres):
+def merge_labels(per_asym_residue_index: Dict[int,List[int]], 
+                 labels: List[Dict], align: List[Tuple[int, int]], 
+                 original_nres: int) -> Dict[str, torch.Tensor]:
     """
     Merge ground truth labels according to the permutation results
 
-    labels: list of original ground truth feats
-    align: list of tuples, each entry specify the corresponding label of the asym.
+    Args:
+        per_asym_residue_index: a dictionary recording which residues belong to which aysm_id 
+        labels: list of original ground truth feats e.g. if there're 5 chains, labels will have a length of 5
+        align: list of tuples, each entry specify the corresponding label of the asym.
+        original_nres: int, corresponding to the number of residues specified by crop_size in config.py
 
+    Returns:
+        A new dictionary of permuated ground truth features
     modified based on UniFold:
     https://github.com/dptech-corp/Uni-Fold/blob/b1c89a2cebd4e4ee4c47b4e443f92beeb9138fbb/unifold/losses/chain_align.py#L176C1-L176C1
     """
@@ -230,13 +286,20 @@ def merge_labels(per_asym_residue_index, labels, align, original_nres):
     return outs
 
 
-def split_ground_truth_labels(gt_features):
+def split_ground_truth_labels(gt_features: dict) -> List[Dict]:
     """
     Splits ground truth features according to chains
 
+    Args:
+    gt_features: A dictionary within a the PyTorch DataSet iteration, which returns by the upstream DataLoader.iter() method
+    In the DataLoader pipeline, all tensors belonging to all the ground truth changes are concatenated so it stays the same as monomer data input format/pipeline,
+    thus, this function is needed to 1) detect the number of chains i.e. unique(asym_id) 
+    2) split the concatenated tensors back to individual ones that correspond to individual asym_ids
+
     Returns:
-    a list of feature dictionaries with only necessary ground truth features
-    required to finish multi-chain permutation
+        a list of feature dictionaries with only necessary ground truth features
+        required to finish multi-chain permutation, e.g. it will be a list of 5 elements if there 
+        are 5 chains in total.
     """
     unique_asym_ids, asym_id_counts = torch.unique(gt_features["asym_id"], sorted=True, return_counts=True)
     n_res = gt_features["asym_id"].shape[-1]
@@ -251,7 +314,16 @@ def split_ground_truth_labels(gt_features):
     return labels
 
 
-def get_per_asym_residue_index(features):
+def get_per_asym_residue_index(features: dict) -> Dict[int, torch.Tensor]:
+    """
+    A function that retrieve which residues belong to which asym_id 
+
+    Args:
+        features: a dictionary that contains input features after cropping 
+
+    Returns:
+        A dictionary that records which region of the sequence belongs to which asym_id
+    """
     unique_asym_ids = [i for i in torch.unique(features["asym_id"]) if i != 0]
     per_asym_residue_index = {}
     for cur_asym_id in unique_asym_ids:
@@ -261,34 +333,36 @@ def get_per_asym_residue_index(features):
     return per_asym_residue_index
 
 
-def get_entity_2_asym_list(batch):
+def get_entity_2_asym_list(features: dict) -> Dict[int, list]:
     """
     Generates a dictionary mapping unique entity IDs to lists of unique asymmetry IDs (asym_id) for each entity.
 
     Args:
-        batch (dict): A dictionary containing data batches, including "entity_id" and "asym_id" tensors.
+        features (dict): A dictionary containing data features, including "entity_id" and "asym_id" tensors.
 
     Returns:
         entity_2_asym_list (dict): A dictionary where keys are unique entity IDs, and values are lists of unique asymmetry IDs
         associated with each entity.
     """
     entity_2_asym_list = {}
-    unique_entity_ids = torch.unique(batch["entity_id"])
+    unique_entity_ids = torch.unique(features["entity_id"])
     for cur_ent_id in unique_entity_ids:
-        ent_mask = batch["entity_id"] == cur_ent_id
-        cur_asym_id = torch.unique(batch["asym_id"][ent_mask])
+        ent_mask = features["entity_id"] == cur_ent_id
+        cur_asym_id = torch.unique(features["asym_id"][ent_mask])
         entity_2_asym_list[int(cur_ent_id)] = cur_asym_id
     return entity_2_asym_list
 
 
-def calculate_input_mask(true_ca_masks, anchor_gt_idx, anchor_gt_residue,
-                         asym_mask, pred_ca_mask):
+def calculate_input_mask(true_ca_masks: List[torch.Tensor], anchor_gt_idx: torch.Tensor, 
+                         anchor_gt_residue: torch.Tensor,
+                         asym_mask: torch.Tensor, pred_ca_mask: torch.Tensor) -> torch.Tensor:
     """
     Calculate an input mask for downstream optimal transformation computation
 
     Args:
-        true_ca_masks (Tensor): ca mask from ground truth.
-        anchor_gt_idx (Tensor): The index of selected ground truth anchor.
+        true_ca_masks: list of masks from ground truth chains.
+        anchor_gt_idx (Tensor): a tensor with one integer in it. The index of selected ground truth anchor.
+        anchor_gt_residue:a 1D vector tensor of residue indexes that belongs to the selected ground truth anchor
         asym_mask (Tensor): Boolean tensor indicating which regions are selected predicted anchor.
         pred_ca_mask (Tensor): ca mask from predicted structure.
 
@@ -303,11 +377,38 @@ def calculate_input_mask(true_ca_masks, anchor_gt_idx, anchor_gt_residue,
     return input_mask
 
 
-def calculate_optimal_transform(true_ca_poses,
-                                anchor_gt_idx, anchor_gt_residue,
-                                true_ca_masks, pred_ca_mask,
-                                asym_mask,
-                                pred_ca_pos):
+def calculate_optimal_transform(true_ca_poses: List[torch.Tensor],
+                                anchor_gt_idx: int, anchor_gt_residue: torch.Tensor,
+                                true_ca_masks: List[torch.Tensor], pred_ca_mask: torch.Tensor,
+                                asym_mask: torch.Tensor,
+                                pred_ca_pos: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    
+    """
+    Takes selected anchor ground truth c-alpha positions and 
+    selected predicted anchor c-alpha position then calculate the optimal rotation matrix
+    to align ground-truth anchor and predicted anchor
+
+    Args: 
+        true_ca_poses: a list of tensors, corresponding to the c-alpha positions of the ground truth structure. e.g. If there are 5 chains, this list will have a length of 5
+        anchor_gt_idx (Tensor): a tensor with one integer in it. The index of selected ground truth anchor.
+        anchor_gt_residue:a 1D vector tensor of residue indexes that belongs to the selected ground truth anchor
+        true_ca_masks: list of masks from ground truth chains e.g. it will be length=5 if there are 5 chains in ground truth structure
+        pred_ca_mask: A boolean tensor corresponds to the mask to mask the predicted features
+        asym_mask: A boolean tensor that mask out other elements in a tensor if they do not belong to a this asym_id
+        pred_ca_pos: a [nres*3] tensor of predicted c-alpha atom positions
+
+    Process:
+    1) select an achor chain from ground truth, denoted by anchor_gt_idx, and 
+    an chor chain from the predicted structure. Both anchor_gt and anchor_pred have exactly the same sequence
+    2) obtain the C-alpha positions corresponding to the selected anchor_gt, done be slicing the true_ca_pose according to anchor_gt_residue
+    3) calculate the optimal transformation that can best align the C-alpha atoms of anchor_pred to those of anchor_gt, 
+    done by Kabsch algorithm: source https://en.wikipedia.org/wiki/Kabsch_algorithm
+
+    Returns:
+    a rotation matrix that record the optimal rotation 
+    that will best align selected anchor prediction to selected anchor truth
+    a matrix records how the atoms should be shifted after applying r i.e. optimal alignment requires 1) rotate 2) shift the positions 
+    """
     input_mask = calculate_input_mask(true_ca_masks,
                                       anchor_gt_idx,
                                       anchor_gt_residue,
@@ -326,13 +427,27 @@ def calculate_optimal_transform(true_ca_poses,
     return r, x
 
 
-def compute_permutation_alignment(out, features, ground_truth):
+def compute_permutation_alignment(out: Dict[str,torch.Tensor], 
+                                  features: Dict[str,torch.Tensor], 
+                                  ground_truth: List[Dict[str, torch.Tensor]]) -> Tuple[List[Tuple[int, int]], Dict[int, List[int]]]:
     """
-    A class method that first permutate chains in ground truth first
-    before calculating the loss.
+    A method that permutes chains in ground truth before calculating the loss 
+    because the mapping between the predicted and ground-truth will become arbitrary.
+    The model cannot be assumed to predict chains in the same order as the ground truth. 
+    Thus, this function pick the optimal permutaion of predicted chains that best matches the ground truth,
+    by minimising the RMSD i.e. the best permutation of ground truth chains is selected based on which permutation has the lowest RMSD calculation
 
     Details are described in Section 7.3 in the Supplementary of AlphaFold-Multimer paper:
     https://www.biorxiv.org/content/10.1101/2021.10.04.463034v2
+
+    Args:
+        out: a dictionary of output tensors from model.forward()
+        features: a dictionary of feature tensors that are used as input for model.forward()
+        ground_truth: a list of dictionaries of features corresponding to chains in ground truth structure e.g. it will be a length of 5 if  there are 5 chains in ground truth structure
+
+    Returns:
+        a list of tuple(int,int) that instructs how ground truth chains should be permutated
+        a dictionary recording which residues belong to which aysm_id 
     """
     unique_asym_ids = set(torch.unique(features['asym_id']).tolist())
     unique_asym_ids.discard(0)  # Remove padding asym_id
@@ -397,13 +512,19 @@ def compute_permutation_alignment(out, features, ground_truth):
     return best_align, per_asym_residue_index
 
 
-def multi_chain_permutation_align(out, features, ground_truth):
-    """Compute multi-chain permutation alignment.
+def multi_chain_permutation_align(out: Dict[str, torch.Tensor], 
+                                  features: Dict[str, torch.Tensor], 
+                                  ground_truth: List[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
+    """
+    Compute multi-chain permutation alignment.
 
     Args:
-        out: The output of model.forward()
-        features: Input features
-        ground_truth: Ground truth features
+        out: a dictionary of output tensors from model.forward()
+        features: a dictionary of feature tensors that are used as input for model.forward()
+        ground_truth: a list of dictionaries of features corresponding to chains in ground truth structure e.g. it will be a length of 5 if  there are 5 chains in ground truth structure
+    
+    Returns:
+        features: a dictionary with updated ground truth feature tensors, ready for downstream loss calculations.
     """
 
     labels = split_ground_truth_labels(ground_truth)

tests/test_permutation.py

@@ -48,7 +48,8 @@ class TestPermutation(unittest.TestCase):
         self.chain_a_num_res = 9
         self.chain_b_num_res = 13
         # below create default fake ground truth structures for a hetero-pentamer A2B3
-        self.residue_index = list(range(self.chain_a_num_res)) * 2 + list(range(self.chain_b_num_res)) * 3
+        self.residue_index = list(
+            range(self.chain_a_num_res)) * 2 + list(range(self.chain_b_num_res)) * 3
         self.num_res = self.chain_a_num_res * 2 + self.chain_b_num_res * 3
         self.asym_id = torch.tensor([[1] * self.chain_a_num_res + [2] * self.chain_a_num_res + [
             3] * self.chain_b_num_res + [4] * self.chain_b_num_res + [5] * self.chain_b_num_res], device=device)
@@ -63,19 +64,44 @@ class TestPermutation(unittest.TestCase):
             'entity_id': self.entity_id,
             'seq_length': torch.tensor([57])
         }
-        anchor_gt_asym, anchor_pred_asym = get_least_asym_entity_or_longest_length(batch, batch['asym_id'])
+        anchor_gt_asym, anchor_pred_asym = get_least_asym_entity_or_longest_length(
+            batch, batch['asym_id'])
         anchor_gt_asym = int(anchor_gt_asym)
         anchor_pred_asym = {int(i) for i in anchor_pred_asym}
         expected_anchors = {1, 2}
         expected_non_anchors = {3, 4, 5}
 
-        self.assertIn(anchor_gt_asym, expected_anchors) 
+        self.assertIn(anchor_gt_asym, expected_anchors)
         self.assertNotIn(anchor_gt_asym, expected_non_anchors)
         # Check that predicted anchors are within expected anchor set
         self.assertEqual(anchor_pred_asym, expected_anchors & anchor_pred_asym)
-        self.assertEqual(set(), anchor_pred_asym & expected_non_anchors) 
+        self.assertEqual(set(), anchor_pred_asym & expected_non_anchors)
 
     def test_2_permutation_pentamer(self):
+        """
+        Test the permutation results on a pentamer A2B3, in which protein A has 9 residues
+        and protein B has 13 residues. 
+
+        Expected outputs:
+            Only protein A should be selected as an anchor thus, in the output list, either [(0,1), (1,0)] or [(0,0), (1,1)] are allowed 
+            The 3 chains from protein B should ALWAYS be aligned in a way that predicted b1 to be aligned with ground truth b1, pred b2 to ground truth b2 
+            as shown below:
+
+            predicted structure:              a2 - a1 - b2 - b3 - b1
+            indexes in the predicted list:    0    1    2    3    4
+
+            ground truth structure:           a1 - a2 - b1 - b2 - b3
+            indexes in the ground truth list: 0    1    2    3    4
+
+            then the 2 protein A chains are free to be aligned by either order, thus either [(0,1),(1,0)] or [(0,0),(1,1)] is valid.
+
+            However, the 3 protein B chains should be strictly aligned in the following order: 
+            [(2,3), (3,4), (4,2)], regardless of how protein A chains are aligned. 
+
+            Therefore, the only 2 correct permutations are :
+            [(0, 1), (1, 0), (2, 3), (3, 4), (4, 2)] and 
+            [(0, 0), (1, 1), (2, 3), (3, 4), (4, 2)]
+        """
         batch = {
             'asym_id': self.asym_id,
             'sym_id': self.sym_id,
@@ -85,7 +111,7 @@ class TestPermutation(unittest.TestCase):
         }
         batch['asym_id'] = batch['asym_id'].reshape(1, self.num_res)
         batch["residue_index"] = torch.tensor([self.residue_index])
-        # create fake ground truth atom positions        
+        # create fake ground truth atom positions
         chain_a1_pos = torch.randint(15, (self.chain_a_num_res, 3 * 37),
                                      dtype=torch.float).reshape(1, self.chain_a_num_res, 37, 3)
         chain_a2_pos = torch.matmul(chain_a1_pos, self.rotation_matrix_x) + 10
@@ -93,16 +119,22 @@ class TestPermutation(unittest.TestCase):
         chain_b1_pos = torch.randint(low=15, high=30, size=(self.chain_b_num_res, 3 * 37),
                                      dtype=torch.float).reshape(1, self.chain_b_num_res, 37, 3)
         chain_b2_pos = torch.matmul(chain_b1_pos, self.rotation_matrix_y) + 10
-        chain_b3_pos = torch.matmul(torch.matmul(chain_b1_pos, self.rotation_matrix_z), self.rotation_matrix_x) + 30
-        # Below permutate predicted chain positions 
-        pred_atom_position = torch.cat((chain_a2_pos, chain_a1_pos, chain_b2_pos, chain_b3_pos, chain_b1_pos), dim=1)
+        chain_b3_pos = torch.matmul(torch.matmul(
+            chain_b1_pos, self.rotation_matrix_z), self.rotation_matrix_x) + 30
+        # Below permutate predicted chain positions
+        # here the b2 chain from the ground truth is deliberately put in b1 chain's position, and predicted b3 chain to b2's position
+        # and predicted b1 chain to b3's position
+        pred_atom_position = torch.cat(
+            (chain_a2_pos, chain_a1_pos, chain_b2_pos, chain_b3_pos, chain_b1_pos), dim=1)
+
         pred_atom_mask = torch.ones((1, self.num_res, 37))
         out = {
             'final_atom_positions': pred_atom_position,
             'final_atom_mask': pred_atom_mask
         }
 
-        true_atom_position = torch.cat((chain_a1_pos, chain_a2_pos, chain_b1_pos, chain_b2_pos, chain_b3_pos), dim=1)
+        true_atom_position = torch.cat(
+            (chain_a1_pos, chain_a2_pos, chain_b1_pos, chain_b2_pos, chain_b3_pos), dim=1)
         true_atom_mask = torch.cat((torch.ones((1, self.chain_a_num_res, 37)),
                                     torch.ones((1, self.chain_a_num_res, 37)),
                                     torch.ones((1, self.chain_b_num_res, 37)),
@@ -111,27 +143,47 @@ class TestPermutation(unittest.TestCase):
         batch['all_atom_positions'] = true_atom_position
         batch['all_atom_mask'] = true_atom_mask
 
-        aligns, _ = compute_permutation_alignment(out, batch,
-                                                  batch)
-        print(f"##### aligns is {aligns}")
-        possible_outcome = [[(0, 1), (1, 0), (2, 3), (3, 4), (4, 2)], [(0, 0), (1, 1), (2, 3), (3, 4), (4, 2)]]
-        wrong_outcome = [[(0, 1), (1, 0), (2, 4), (3, 2), (4, 3)], [(0, 0), (1, 1), (2, 2), (3, 3), (4, 4)]]
-        self.assertIn(aligns, possible_outcome)
-        self.assertNotIn(aligns, wrong_outcome)
+        aligns, per_asym_residue_index = compute_permutation_alignment(out, batch,
+                                                                       batch)
+
+        expected_asym_residue_index = {
+            1: torch.tensor(list(range(self.chain_a_num_res))),
+            2: torch.tensor(list(range(self.chain_a_num_res))),
+            3: torch.tensor(list(range(self.chain_b_num_res))),
+            4: torch.tensor(list(range(self.chain_b_num_res))),
+            5: torch.tensor(list(range(self.chain_b_num_res)))
+        }
+        chain_a_permutated_chain_b_permutated = [
+            (0, 1), (1, 0), (2, 3), (3, 4), (4, 2)]
+        chain_a_not_permutated_chain_b_permutated = [
+            (0, 0), (1, 1), (2, 3), (3, 4), (4, 2)]
+        chain_a_permutated_chain_b_not_permuated = [
+            (0, 1), (1, 0), (2, 2), (3, 3), (4, 4)]
+        chain_a_not_permutated_chain_b_not_permuated = [
+            (0, 0), (1, 1), (2, 2), (3, 3), (4, 4)]
+        
+        # test on the permutation alignments 
+        self.assertIn(aligns, [chain_a_permutated_chain_b_permutated,
+                      chain_a_not_permutated_chain_b_permutated])
+        self.assertNotIn(aligns, [chain_a_permutated_chain_b_not_permuated,
+                         chain_a_not_permutated_chain_b_not_permuated])
+
+        # test on the per_aysm_residue_index 
+        for k, v in expected_asym_residue_index.items():
+            self.assertTrue(torch.equal(v, per_asym_residue_index[k]))
 
-    @unittest.skip("Test needs to be fixed post-refactor")
     def test_3_merge_labels(self):
         nres_pad = 325 - 57  # suppose the cropping size is 325
         batch = {
-            'asym_id': pad_features(self.asym_id, nres_pad, pad_dim=1),
-            'sym_id': pad_features(self.sym_id, nres_pad, pad_dim=1),
-            'entity_id': pad_features(self.entity_id, nres_pad, pad_dim=1),
-            'aatype': torch.randint(21, size=(1, 325)),
+            'asym_id': self.asym_id,
+            'sym_id': self.sym_id,
+            'entity_id': self.entity_id,
+            'aatype': torch.randint(21, size=(1, 57)),
             'seq_length': torch.tensor([57])
         }
-        batch['asym_id'] = batch['asym_id'].reshape(1, 325)
-        batch["residue_index"] = pad_features(torch.tensor(self.residue_index).reshape(1, 57), nres_pad, pad_dim=1)
-        # create fake ground truth atom positions        
+        batch['asym_id'] = batch['asym_id'].reshape(1, 57)
+        batch["residue_index"] = torch.tensor([self.residue_index])
+        # create fake ground truth atom positions
         chain_a1_pos = torch.randint(15, (self.chain_a_num_res, 3 * 37),
                                      dtype=torch.float).reshape(1, self.chain_a_num_res, 37, 3)
         chain_a2_pos = torch.matmul(chain_a1_pos, self.rotation_matrix_x) + 10
@@ -139,39 +191,64 @@ class TestPermutation(unittest.TestCase):
         chain_b1_pos = torch.randint(low=15, high=30, size=(self.chain_b_num_res, 3 * 37),
                                      dtype=torch.float).reshape(1, self.chain_b_num_res, 37, 3)
         chain_b2_pos = torch.matmul(chain_b1_pos, self.rotation_matrix_y) + 10
-        chain_b3_pos = torch.matmul(torch.matmul(chain_b1_pos, self.rotation_matrix_z), self.rotation_matrix_x) + 30
-        # Below permutate predicted chain positions 
-        pred_atom_position = torch.cat((chain_a2_pos, chain_a1_pos, chain_b2_pos, chain_b3_pos, chain_b1_pos), dim=1)
+        chain_b3_pos = torch.matmul(torch.matmul(
+            chain_b1_pos, self.rotation_matrix_z), self.rotation_matrix_x) + 30
+        # Below permutate predicted chain positions
+        pred_atom_position = torch.cat(
+            (chain_a2_pos, chain_a1_pos, chain_b2_pos, chain_b3_pos, chain_b1_pos), dim=1)
         pred_atom_mask = torch.ones((1, self.num_res, 37))
-        pred_atom_position = pad_features(pred_atom_position, nres_pad, pad_dim=1)
+        pred_atom_position = pad_features(
+            pred_atom_position, nres_pad, pad_dim=1)
         pred_atom_mask = pad_features(pred_atom_mask, nres_pad, pad_dim=1)
         out = {
             'final_atom_positions': pred_atom_position,
             'final_atom_mask': pred_atom_mask
         }
-        true_atom_position = torch.cat((chain_a1_pos, chain_a2_pos, chain_b1_pos, chain_b2_pos, chain_b3_pos), dim=1)
+        true_atom_position = torch.cat(
+            (chain_a1_pos, chain_a2_pos, chain_b1_pos, chain_b2_pos, chain_b3_pos), dim=1)
         true_atom_mask = torch.cat((torch.ones((1, self.chain_a_num_res, 37)),
                                     torch.ones((1, self.chain_a_num_res, 37)),
                                     torch.ones((1, self.chain_b_num_res, 37)),
                                     torch.ones((1, self.chain_b_num_res, 37)),
                                     torch.ones((1, self.chain_b_num_res, 37))), dim=1)
-        batch['all_atom_positions'] = pad_features(true_atom_position, nres_pad, pad_dim=1)
-        batch['all_atom_mask'] = pad_features(true_atom_mask, nres_pad=nres_pad, pad_dim=1)
 
-        # tensor_to_cuda = lambda t: t.to('cuda')
-        # ground_truth = tensor_tree_map(tensor_to_cuda,ground_truth)
+        batch['all_atom_positions'] = true_atom_position
+        batch['all_atom_mask'] = true_atom_mask
+
+        # Below create a fake_input_features
+        fake_input_features = {
+            'asym_id': pad_features(self.asym_id, nres_pad, pad_dim=1),
+            'sym_id': pad_features(self.sym_id, nres_pad, pad_dim=1),
+            'entity_id': pad_features(self.entity_id, nres_pad, pad_dim=1),
+            'aatype': torch.randint(21, size=(1, 325)),
+            'seq_length': torch.tensor([57])
+        }
+        fake_input_features['asym_id'] = fake_input_features['asym_id'].reshape(
+            1, 325)
+        fake_input_features["residue_index"] = pad_features(
+            torch.tensor(self.residue_index).reshape(1, 57), nres_pad, pad_dim=1)
+        fake_input_features['all_atom_positions'] = pad_features(
+            true_atom_position, nres_pad, pad_dim=1)
+        fake_input_features['all_atom_mask'] = pad_features(
+            true_atom_mask, nres_pad=nres_pad, pad_dim=1)
+
+        # NOTE
+        # batch: simulates ground_truth features
+        # fake_input_features: simulates the data that are going be used as input for model.forward(fake_input_features)
+        # out: simulates the output of model.forward(fake_input_features)
         aligns, per_asym_residue_index = compute_permutation_alignment(out,
-                                                                       batch,
+                                                                       fake_input_features,
                                                                        batch)
-        print(f"##### aligns is {aligns}")
         labels = split_ground_truth_labels(batch)
 
         labels = merge_labels(per_asym_residue_index, labels, aligns,
                               original_nres=batch['aatype'].shape[-1])
 
-        self.assertTrue(torch.equal(labels['residue_index'], batch['residue_index']))
+        self.assertTrue(torch.equal(
+            labels['residue_index'], batch['residue_index']))
 
         expected_permutated_gt_pos = torch.cat((chain_a2_pos, chain_a1_pos, chain_b2_pos, chain_b3_pos, chain_b1_pos),
                                                dim=1)
-        expected_permutated_gt_pos = pad_features(expected_permutated_gt_pos, nres_pad, pad_dim=1)
-        self.assertTrue(torch.equal(labels['all_atom_positions'], expected_permutated_gt_pos))
+
+        self.assertTrue(torch.equal(
+            labels['all_atom_positions'], expected_permutated_gt_pos))