test and implementation of bending potential in RW module

polyply/src/random_walk.py

@@ -248,6 +248,7 @@ def __init__(self,
         self.nrewind = nrewind
         self.placed_nodes = []
         self.max_bend = max_bend
+        self.prev_prob = 1
 
     def _rewind(self, current_step):
         nodes = [node for _, node in self.placed_nodes[-self.nrewind:-1]]
@@ -283,28 +284,38 @@ def checks_milestones(self, current_node, current_position, fudge=0.7):
 
     def bendiness(self, point, node):
         """
-        Perform Monte-Carlo like sampling of bending potential.
+        Perform Monte-Carlo like sampling of bending probability.
         """
-        b_nodes = list(self.molecule.search_tree.predecessors(node))
-
-        if len(b_nodes) == 1:
-            b_node = b_nodes[0]
-        else:
-            return True
-
+        b_node = list(self.molecule.search_tree.predecessors(node))[0]
         c_nodes = list(self.molecule.search_tree.predecessors(b_node))
         if len(c_nodes) == 1:
             c_node = c_nodes[0]
-            prob, test_prob = self.nonbond_matrix.compute_bending_probability(point, self.mol_idx, node, b_node, c_node)
+            # get the atom types aka resnames
+            typea = self.molecule.nodes[node]['resname']
+            typeb = self.molecule.nodes[b_node]['resname']
+            typec = self.molecule.nodes[c_node]['resname']
+            # get the bending constant
+            lp = self.nonbond_matrix.bending_matrix.get((typea, typeb, typec), None)
+            if lp:
+                prob = self.nonbond_matrix.compute_bending_probability(lp,
+                                                                       point,
+                                                                       self.mol_idx,
+                                                                       b_node,
+                                                                       c_node)
+                if self.prev_prob < prob:
+                    self.prev_prob = prob
+                    return True
 
-            if self.prev_prob < prob:
-                self.prev_prob = prob
-                return True
+                # compute test probability from uniform sampling of prob function
+                test_prob = random.uniform(np.exp(lp*1/180)/(180*(np.exp(lp)-1)/lp),
+                                           np.exp(lp*179/180)/(180*(np.exp(lp)-1)/lp))
+                if prob > test_prob:
+                    self.prev_prob = prob
+                    return True
 
-            elif prob > test_prob:
-                return True
-        else:
-            return True
+                return False
+
+        return True
 
     def update_positions(self, vector_bundle, current_node, prev_node):
         """
@@ -384,7 +395,6 @@ def _random_walk(self, meta_molecule):
         count = 0
         path = list(meta_molecule.search_tree.edges)
         step_count = 0
-        self.prev_prob = 1
         while step_count < len(path):
             prev_node, current_node = path[step_count]
             print(step_count)

polyply/tests/test_random_walk.py

@@ -16,6 +16,7 @@
 """
 import math
 import pytest
+import random
 import numpy as np
 from numpy.linalg import norm
 import networkx as nx
@@ -188,6 +189,63 @@ def test_is_overlap(nonbond_matrix, molecule, new_point, result):
     # node 4 is already placed and hence is skipped over
     assert proccessor._is_overlap(new_point, 7, nrexcl=1) == result
 
+@pytest.mark.parametrize('n_coords, new_point, prev_prob, lp, result', (
+    # only 1 previous node -> always True
+    (1,
+     None,
+     1,
+     10,
+     True,
+     ),
+    # no lp is given -> always True
+    (1,
+     None,
+     1,
+     None,
+     True,
+     ),
+    # 180 degrees should be fine even with high prev prob
+    (2,
+     np.array([1., 1., 1.11]),
+     1,
+     10,
+     True
+     ),
+    # small angle and prev prob high
+    (2,
+     np.array([1., 1.47, 0.1]),
+     1.,
+     10,
+     False
+     ),
+    # small angle and prev prob low
+    (2,
+     np.array([1., 1.47, 0.1]),
+     0.,
+     10,
+     True
+     ),
+    # medium angle, prev prob high, unifrom prob low
+    (2,
+     np.array([1., 1.57, 1.74]),
+     1.,
+     10,
+     True
+     ),
+))
+def test_bendiness(nonbond_matrix, molecule, n_coords, new_point, prev_prob, lp, result):
+    # set random seed for reproducability
+    random.seed(1)
+    nb_matrix = add_positions(nonbond_matrix, n_coords)
+    # the bending constant for a series of PEO monomers
+    # the value is the same as in the high test case from
+    # test nb matrix
+    nb_matrix.bending_matrix = {("PEO", "PEO", "PEO"): lp}
+    processor = RandomWalk(mol_idx=0, nonbond_matrix=nb_matrix)
+    processor.molecule = molecule
+    processor.prev_prob = prev_prob
+    assert processor.bendiness(new_point, n_coords) == result
+
 @pytest.mark.parametrize('new_point, restraint, result', (
    # distance restraint true upper_bound
    # ref_node, upper_bound, lower_bound