Transpose SVD matrix V in WeightedRigidAlign to get correct alignments (#22)

* Transpose SVD matrix `V` in `WeightedRigidAlign` to get correct alignment

* Update test_af3.py

Author: amorehead

alphafold3_pytorch/alphafold3.py

@@ -2198,15 +2198,15 @@ def forward(
         U, _, V = torch.svd(cov_matrix)
 
         # Compute the rotation matrix
-        rot_matrix = einsum(U, V, 'b i j, b j k -> b i k')
+        rot_matrix = einsum(U, V, 'b i j, b k j -> b i k')
 
         # Ensure proper rotation matrix with determinant 1
         det = torch.det(rot_matrix)
         det_mask = det < 0
         V_fixed = V.clone()
         V_fixed[det_mask, :, -1] *= -1
 
-        rot_matrix[det_mask] = einsum(U[det_mask], V_fixed[det_mask], 'b i j, b j k -> b i k')
+        rot_matrix[det_mask] = einsum(U[det_mask], V_fixed[det_mask], 'b i j, b k j -> b i k')
 
         # Apply the rotation and translation
         aligned_coords = einsum(pred_coords_centered, rot_matrix, 'b n i, b i j -> b n j') + true_centroid

tests/test_af3.py

@@ -55,13 +55,15 @@ def test_smooth_lddt_loss():
 
 def test_weighted_rigid_align():
     pred_coords = torch.randn(2, 100, 3)
-    true_coords = torch.randn(2, 100, 3)
     weights = torch.rand(2, 100)
 
     align_fn = WeightedRigidAlign()
-    aligned_coords = align_fn(pred_coords, true_coords, weights)
+    aligned_coords = align_fn(pred_coords, pred_coords, weights)
+
+    # `pred_coords` should match itself without any change after alignment
 
-    assert aligned_coords.shape == pred_coords.shape
+    rmsd = torch.sqrt(((pred_coords - aligned_coords) ** 2).sum(dim=-1).mean(dim=-1))
+    assert (rmsd < 1e-5).all()
 
 def test_weighted_rigid_align_with_mask():
     pred_coords = torch.randn(2, 100, 3)