From b9c749eae29e3ec6e8def33ab25c5457e41fe886 Mon Sep 17 00:00:00 2001
From: Kevin Wu <t-kevinwu@microsoft.com>
Date: Wed, 7 Sep 2022 11:43:17 -0700
Subject: [PATCH] Initial working dihedral reversal

---
 protdiff/mynerf.py | 40 +++++++++++++++++++++++++++
 tests/test_nerf.py | 67 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 107 insertions(+)
 create mode 100644 protdiff/mynerf.py
 create mode 100644 tests/test_nerf.py

diff --git a/protdiff/mynerf.py b/protdiff/mynerf.py
new file mode 100644
index 0000000..3ea8ba8
--- /dev/null
+++ b/protdiff/mynerf.py
@@ -0,0 +1,40 @@
+"""
+NERF!
+
+References:
+https://benjamin-computer.medium.com/protein-loops-in-tensorflow-a-i-bio-part-2-f1d802ef8300
+"""
+
+import numpy as np
+
+
+def place_dihedral(a, b, c, bond_angle, bond_length, torsion_angle):
+    """
+    Place the point d such that the bond angle, length, and torsion angle are satisfied
+    """
+    ab = b - a
+    bc = c - b
+    bcn = bc / np.linalg.norm(bc)
+    # numpy is row major
+    d = np.array(
+        [
+            -bond_length * np.cos(bond_angle),
+            bond_length * np.cos(torsion_angle) * np.sin(bond_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]]]
+    )
+    d = m.dot(d)
+    d = d + c
+    return d
+
+
+if __name__ == "__main__":
+    rng = np.random.default_rng(seed=1)
+    a, b, c, d = rng.uniform(low=-5, high=5, size=(4, 3))
+    print(a, b, c, d)
diff --git a/tests/test_nerf.py b/tests/test_nerf.py
new file mode 100644
index 0000000..a2c4ac5
--- /dev/null
+++ b/tests/test_nerf.py
@@ -0,0 +1,67 @@
+"""
+Unit tests for NERF conversion of internal coordinates to cartesian coordinates
+"""
+
+import os, sys
+import unittest
+
+import numpy as np
+from sequence_models import pdb_utils
+from biotite.structure import dihedral
+
+SRC_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), "protdiff")
+assert os.path.isdir(SRC_DIR)
+sys.path.append(SRC_DIR)
+
+import mynerf as nerf
+
+
+class TestDihedralPlacement(unittest.TestCase):
+    """
+    Test the dihedral placement
+    """
+
+    def setUp(self) -> None:
+        self.rng = np.random.default_rng(seed=6489)
+
+    def test_simple(self):
+        """Simple test about origin"""
+        a = np.array([1, 0, 0])
+        b = np.array([0, 0, 0])
+        c = np.array([0, 1, 0])
+        d = np.array([0, 1, 1])
+        calc_d = nerf.place_dihedral(a, b, c, np.pi / 2, 1.0, -np.pi / 2)
+        self.assertTrue(np.allclose(d, calc_d), f"Mismatched: {d} != {calc_d}")
+
+    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}")
+
+
+def angle_between(v1, v2):
+    """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)
+    v1_u = unit_vector(v1)
+    v2_u = unit_vector(v2)
+    return np.arccos(np.clip(np.dot(v1_u, v2_u), -1.0, 1.0))
+
+
+def dist_between(a, b):
+    """Distance between a and b"""
+    d = a - b
+    return np.linalg.norm(d, 2)
+
+
+if __name__ == "__main__":
+    unittest.main()