Update alphafold3.py (#219)

Author: amorehead

alphafold3_pytorch/alphafold3.py

@@ -3283,7 +3283,7 @@ def get_least_asym_entity_or_longest_length(
 
             # Calculate entity length
             entity_mask = batch["entity_id"] == entity_id
-            entity_length[int(entity_id)] = entity_mask.sum().item()
+            entity_length[int(entity_id)] = entity_mask.sum(-1).mode().values.item()
 
         min_asym_count = min(entity_asym_count.values())
         least_asym_entities = [
@@ -3313,6 +3313,18 @@ def get_least_asym_entity_or_longest_length(
             if asym_id in input_asym_id
         ]
 
+        # Since the entity ID to asym ID mapping is many-to-many, we need to select only
+        # prediction asym IDs with equal length w.r.t. the sampled ground truth asym ID
+        anchor_gt_asym_id_length = (
+            (batch["asym_id"] == anchor_gt_asym_id).sum(-1).mode().values.item()
+        )
+        anchor_pred_asym_ids = [
+            asym_id
+            for asym_id in anchor_pred_asym_ids
+            if (batch["asym_id"] == asym_id).sum(-1).mode().values.item()
+            == anchor_gt_asym_id_length
+        ]
+
         # Remap `asym_id` values to remove any gaps in the ground truth asym IDs,
         # but leave the prediction asym IDs as is since they are used for masking
         sorted_asym_ids = sorted(all_asym_ids)