passes own tests

mp_nerf/kb_proteins.py

@@ -721,12 +721,11 @@ def make_torsion_mask(aa, fill=False):
     """ Gives the dihedral of the bond originating each atom. """
     mask = np.zeros(14)
     # backbone
-    if aa == "_":
-        return aa
+
     # sidechain
     for i, torsion in enumerate(SC_BUILD_INFO[aa]['torsion-vals']):
         if fill: 
-            mask[4+i] = MP2SC_INFO[aa][ SC_BUILD_INFO[aa]["atom-names"] ]["bond_dihedral"]
+            mask[4+i] = MP2SC_INFO[aa][ SC_BUILD_INFO[aa]["atom-names"][i] ]["bond_dihedral"]
         else: 
             # https://github.com/jonathanking/sidechainnet/blob/master/sidechainnet/structure/StructureBuilder.py#L372
             # 999 is an anotation -- change later 

tests/tes_main.py

@@ -1,6 +1,11 @@
 import numpy as np
 import torch
 
+from mp_nerf import *
+from mp_nerf.utils import *
+from mp_nerf.kb_proteins import *
+from mp_nerf.proteins import *
+
 def test_nerf_and_dihedral():
     # create points
     a = torch.tensor([1,2,3]).float()
@@ -26,5 +31,5 @@ def test_nerf_and_dihedral():
     # reconstruct
     # doesnt work because the scn angle was not measured correctly
     # so the method corrects that incorrection
-    assert (nerf_torch(a, b, c, l, theta, chi - np.pi) - torch.tensor([1,0,6])).sum().abs() < 0.1
-    assert get_dihedral_torch(a, b, c, d).item() == chi
+    assert (mp_nerf_torch(a, b, c, l, theta, chi - np.pi) - torch.tensor([1,0,6])).sum().abs() < 0.1
+    assert get_dihedral(a, b, c, d).item() == chi