From 149fe8a72bc2c180bd2ce30567e1ba8be80aff6c Mon Sep 17 00:00:00 2001
From: Kevin Wu <t-kevinwu@microsoft.com>
Date: Wed, 7 Sep 2022 12:21:57 -0700
Subject: [PATCH] Simplify dihedral calculation code and add tests

---
 protdiff/mynerf.py | 21 +++++++++------------
 tests/test_nerf.py | 24 ++++++++++++------------
 2 files changed, 21 insertions(+), 24 deletions(-)

diff --git a/protdiff/mynerf.py b/protdiff/mynerf.py
index 3ea8ba8..4112cd3 100644
--- a/protdiff/mynerf.py
+++ b/protdiff/mynerf.py
@@ -8,14 +8,15 @@ https://benjamin-computer.medium.com/protein-loops-in-tensorflow-a-i-bio-part-2-
 import numpy as np
 
 
-def place_dihedral(a, b, c, bond_angle, bond_length, torsion_angle):
+def place_dihedral(a: np.ndarray, b:np.ndarray, c: np.ndarray, bond_angle:float, bond_length:float, torsion_angle:float) -> np.ndarray:
     """
     Place the point d such that the bond angle, length, and torsion angle are satisfied
+    with the series a, b, c, d.
     """
+    assert a.ndim == b.ndim == c.ndim == 1
+    unit_vec = lambda x: x / np.linalg.norm(x)
     ab = b - a
-    bc = c - b
-    bcn = bc / np.linalg.norm(bc)
-    # numpy is row major
+    bc = unit_vec(c - b)
     d = np.array(
         [
             -bond_length * np.cos(bond_angle),
@@ -23,15 +24,11 @@ def place_dihedral(a, b, c, bond_angle, bond_length, torsion_angle):
             bond_length * np.sin(torsion_angle) * np.sin(bond_angle),
         ]
     )
-    n = np.cross(ab, bcn)
-    n /= np.linalg.norm(n)
-    nbc = np.cross(n, bcn)
-    m = np.array(
-        [[bcn[0], nbc[0], n[0]], [bcn[1], nbc[1], n[1]], [bcn[2], nbc[2], n[2]]]
-    )
+    n = unit_vec(np.cross(ab, bc))
+    nbc = np.cross(n, bc)
+    m = np.stack([bc, nbc, n]).T
     d = m.dot(d)
-    d = d + c
-    return d
+    return d + c
 
 
 if __name__ == "__main__":
diff --git a/tests/test_nerf.py b/tests/test_nerf.py
index a2c4ac5..0369e7f 100644
--- a/tests/test_nerf.py
+++ b/tests/test_nerf.py
@@ -35,20 +35,20 @@ class TestDihedralPlacement(unittest.TestCase):
 
     def test_randomized(self):
         """Simple test using randomized values"""
-        a, b, c, d = self.rng.uniform(low=-5, high=5, size=(4, 3))
-        print(a, b, c, d)
-        calc_d = nerf.place_dihedral(
-            a,
-            b,
-            c,
-            angle_between(d - c, b - c),
-            dist_between(c, d),
-            dihedral(a, b, c, d),
-        )
-        self.assertTrue(np.allclose(d, calc_d), f"Mismatched: {d} != {calc_d}")
+        for _ in range(5):
+            a, b, c, d = self.rng.uniform(low=-5, high=5, size=(4, 3))
+            calc_d = nerf.place_dihedral(
+                a,
+                b,
+                c,
+                angle_between(d - c, b - c),
+                dist_between(c, d),
+                dihedral(a, b, c, d),
+            )
+            self.assertTrue(np.allclose(d, calc_d), f"Mismatched: {d} != {calc_d}")
 
 
-def angle_between(v1, v2):
+def angle_between(v1, v2) -> float:
     """Gets the angle between u and v"""
     # https://stackoverflow.com/questions/2827393/angles-between-two-n-dimensional-vectors-in-python
     unit_vector = lambda vector: vector / np.linalg.norm(vector)