Ensure ExpressCoordinatesInFrame perfectly matches Algorithm 29 of the AF3 supplement (#25)

* Ensure `ExpressCoordinatesInFrame` perfectly matches Algorithm 29 of the AF3 supplement

* Update test_af3.py

Author: amorehead

alphafold3_pytorch/alphafold3.py

@@ -2293,22 +2293,26 @@ def forward(
         elif frame.ndim == 3:
             frame = rearrange(frame, 'b fr fc -> b 1 fr fc')
 
-        # Extract frame points
-        a, b, c = frame.unbind(dim = -1)
-
-        # Compute unit vectors of the frame
-        e1 = F.normalize(a - b, dim = -1, eps = self.eps)
-        e2 = F.normalize(c - b, dim = -1, eps = self.eps)
-        e3 = torch.cross(e1, e2, dim = -1)
-
-        # Express coordinates in the frame basis
-        v = coords - b
-
-        transformed_coords = torch.stack([
-            einsum(v, e1, '... i, ... i -> ...'),
-            einsum(v, e2, '... i, ... i -> ...'),
-            einsum(v, e3, '... i, ... i -> ...')
-        ], dim = -1)
+        # Extract frame atoms
+        a, b, c = frame.unbind(dim=-1)
+        w1 = F.normalize(a - b, dim=-1, eps=self.eps)
+        w2 = F.normalize(c - b, dim=-1, eps=self.eps)
+
+        # Build orthonormal basis
+        e1 = F.normalize(w1 + w2, dim=-1, eps=self.eps)
+        e2 = F.normalize(w2 - w1, dim=-1, eps=self.eps)
+        e3 = torch.cross(e1, e2, dim=-1)
+
+        # Project onto frame basis
+        d = coords - b
+        transformed_coords = torch.stack(
+            [
+                einsum(d, e1, '... i, ... i -> ...'),
+                einsum(d, e2, '... i, ... i -> ...'),
+                einsum(d, e3, '... i, ... i -> ...'),
+            ],
+            dim=-1,
+        )
 
         return transformed_coords
 

tests/test_af3.py

@@ -127,14 +127,15 @@ def test_express_coordinates_in_frame():
 
 def test_compute_alignment_error():
     pred_coords = torch.randn(2, 100, 3)
-    true_coords = torch.randn(2, 100, 3)
     pred_frames = torch.randn(2, 100, 3, 3)
-    true_frames = torch.randn(2, 100, 3, 3)
+
+    # `pred_coords` should match itself in frame basis
 
     error_fn = ComputeAlignmentError()
-    alignment_errors = error_fn(pred_coords, true_coords, pred_frames, true_frames)
+    alignment_errors = error_fn(pred_coords, pred_coords, pred_frames, pred_frames)
 
     assert alignment_errors.shape == (2, 100)
+    assert (alignment_errors.mean(-1) < 1e-3).all()
 
 def test_centre_random_augmentation():
     coords = torch.randn(2, 100, 3)