Merge pull request #396 from bigict/relax

fix: violation key error

Author: chungongyu

profold2/relax/amber_minimize.py

@@ -345,29 +345,36 @@ def to_numpy(loss_dict):
     return loss_dict
 
   # Compute between residue backbone violations of bonds and angles.
-  connection_violations = to_numpy(
-      functional.between_residue_bond_loss(
-          pred_points, points_mask, residue_index, aatypes
+  violations = {}
+  violations.update(
+      to_numpy(
+          functional.between_residue_bond_loss(
+              pred_points, points_mask, residue_index, aatypes
+          )
       )
   )
-  between_residue_violations = to_numpy(
-      functional.between_residue_clash_loss(
-          pred_points, points_mask, residue_index, aatypes
+  violations.update(
+      to_numpy(
+          functional.between_residue_clash_loss(
+              pred_points, points_mask, residue_index, aatypes
+          )
       )
   )
-  within_residue_violations = to_numpy(
-      functional.within_residue_clash_loss(
-          pred_points, points_mask, residue_index, aatypes
+  violations.update(
+      to_numpy(
+          functional.within_residue_clash_loss(
+              pred_points, points_mask, residue_index, aatypes
+          )
       )
   )
 
   # Combine them to a single per-residue violation mask (used later for LDDT).
   per_residue_violation_mask = np.max(
       np.stack(
           [
-              connection_violations['per_residue_violation_mask'],
-              np.max(between_residue_violations['per_atom_clash_mask'], axis=-1),
-              np.max(within_residue_violations['per_atom_clash_mask'], axis=-1)
+              violations['per_residue_violation_mask'],
+              np.max(violations['between_residue_per_atom_clash_mask'], axis=-1),
+              np.max(violations['within_residue_per_atom_clash_mask'], axis=-1)
           ]
       ),
       axis=0