add ANM calculations node anotation

src/graphpro/annotations.py

@@ -110,6 +110,32 @@ def encode(self, G: Graph) -> torch.tensor:
         ]
         return torch.tensor(values ,dtype=torch.float)
 
+class ANMSlowModes(NodeAnnotation):
+    """ Calculate ANM and select specific slow modes per residue
+    """
+    def __init__(self, attr_name: str = 'anm_slow_mode', modes: int = 3):
+        """ Inits the residue generation
+            Args:
+                attr_name: name of the attribute prefix
+                modes: number of modes.
+        """
+        self.attr_name  = attr_name
+        self.modes = modes
+
+    def generate(self, G: Graph, atom_group: AtomGroup):
+        resids, modes_eigenvec = compute_gnm_slow_modes(atom_group, self.modes)
+        for i, resid  in enumerate(resids):
+             node_id = G.get_node_by_resid(int(resid))
+             for m in range(0, self.modes):
+                 G.node_attr_add(node_id, f"{self.attr_name}_{m}", modes_eigenvec[i,m].round(3))
+
+    def encode(self, G: Graph) -> torch.tensor:
+        values = [
+            [G.node_attr(n)[f"{self.attr_name}_{m}"] for m in range(0, self.modes)] for n in G.nodes()
+        ]
+        return torch.tensor(values ,dtype=torch.float)
+
+
 class Polarity(NodeAnnotation):
     """ Anotates and encode per residue polarity
     """

src/graphpro/util/modes.py

@@ -10,3 +10,13 @@ def compute_gnm_slow_modes(ag: AtomGroup, n_modes: int = 20):
     gnm.calcModes(n_modes)
 
     return ([a.getResnum() for a in calphas], gnm.getEigvecs())
+
+
+def compute_anm_slow_modes(ag: AtomGroup, n_modes: int = 20):
+    pdy_ag = ag.to_prody()
+    calphas = pdy_ag.select('calpha')
+    anm = pdy.ANM()
+    anm.buildHessian(calphas)
+    anm.calcModes(n_modes)
+
+    return ([a.getResnum() for a in calphas], anm.getEigvecs())

test/graphpro/annnotations/modes_test.py

@@ -3,7 +3,7 @@
 
 from graphpro import md_analisys
 from graphpro.graphgen import ContactMap
-from graphpro.annotations import GNMSlowModes
+from graphpro.annotations import GNMSlowModes, ANMSlowModes
 
 from MDAnalysis.tests.datafiles import PDB_small
 
@@ -27,4 +27,12 @@ def test_gnm_encoding():
     data = G.to_data(node_encoders=[GNMSlowModes(modes=3)])
 
     assert data.x.size() == (214, 3)
-    assert data.x.dtype == torch.float
+    assert data.x.dtype == torch.float
+
+
+def test_anm_slow_modes_annotation():
+    G = md_analisys(u1).generate(ContactMap(cutoff=6), [ANMSlowModes(modes=3)])
+    assert len(G.nodes()) == 214
+    assert G.node_attr(0)['anm_slow_mode_0'] == -0.038
+    assert G.node_attr(0)['anm_slow_mode_1'] == -0.047
+    assert G.node_attr(0)['anm_slow_mode_2'] == 0.074