diff --git a/adaptive_umbrella.py b/adaptive_umbrella.py index 5e852bf..188ae06 100755 --- a/adaptive_umbrella.py +++ b/adaptive_umbrella.py @@ -76,55 +76,49 @@ class UmbrellaRunner(): smaller E_max""" selection = np.where((pmf <= E_max) & (pmf >= 0)) zipped = list(zip(*selection)) + return zipped + def _generate_neighbor_list(self, root): + """ builds a list of all direct neighbors of the root coordinate """ + dimens = len(root) + coords = [] + for i in range(dimens): + if len(coords) == 0: + new_coords = [] + new_coords.append([root[i]-1]) + new_coords.append([root[i]]) + new_coords.append([root[i]+1]) + coords = new_coords + else: + new_coords = [] + for coord in coords: + new_coords.append(coord + [root[i] - 1]) + new_coords.append(coord + [root[i]]) + new_coords.append(coord + [root[i] + 1]) + coords = new_coords + return [tuple(x) for x in coords] + + + def _is_in_pmf(self, frame): + num_dimens = len(self.pmf.shape) + for dimen in range(num_dimens): + if frame[dimen] < 0 or frame[dimen] >= self.pmf.shape[dimen]: + return False + return True + def _get_new_frames(self, pmf, root_frames): """ returns a dict of all frames surrounding the root_frames that have not an assigned energy yet, as well as their corresponding root frame in the format {new_frame1: root_frame1, new_frame2: root_frame2} """ - def generate_neighbor_list(root, coords=[]): - """ recursively builds a list of all direct neighbors of the root coordinate """ - if len(coords) > 0 and len(coords[0]) == len(root): - return [tuple(x) for x in coords] - elif len(coords) == 0: - coords.append([root[0]-1]) - coords.append([root[0]]) - coords.append([root[0]+1]) - return generate_neighbor_list(root, coords) - else: - new_coords = [] - for coord in coords: - dimen = len(coord) - - new_coord = deepcopy(coord) - new_coord.append(root[dimen]-1) - new_coords.append(new_coord) - - new_coord = deepcopy(coord) - new_coord.append(root[dimen]) - new_coords.append(new_coord) - - new_coord = deepcopy(coord) - new_coord.append(root[dimen]+1) - new_coords.append(new_coord) - return generate_neighbor_list(root, new_coords) - - def in_pmf(frame): - num_dimens = len(self.pmf.shape) - for dimen in range(num_dimens): - if frame[dimen] < 0 or frame[dimen] >= self.pmf.shape[dimen]: - return False - return True - - # find all neighboring frames and create a dict that associates them to the root frame with lowest energy new_frames = {} for frame in root_frames: - neighbors = generate_neighbor_list(frame) + neighbors = self._generate_neighbor_list(frame) # remove neighbors that are not inside the pmf - neighbors = [n for n in neighbors if in_pmf(n)] + neighbors = [n for n in neighbors if self._is_in_pmf(n)] # for each neighbor, check if its already in the list and compare root frame energy for n in neighbors: @@ -165,6 +159,11 @@ class UmbrellaRunner(): if(self.max_iterations > 0 and self.num_iterations > self.max_iterations): print("Max iterations reached ({})".format(self.max_iterations)) return + + # stop if the pmf is sampled + if len(np.where(self.pmf < 0)) == 0: + print("Every window of the PMF appears to be sampled.") + return print("~~~~~~~~~~~~~~~ Iteration {}/{} ~~~~~~~~~~~~~~~~".format(self.num_iterations, self.max_iterations)) lambdas = dict([(self._get_lambdas_for_index(x), self._get_lambdas_for_index(y)) for x,y in new_frames.items()]) @@ -186,7 +185,7 @@ class UmbrellaRunner(): print("Max energy increased to {} (max={})".format(self.E, self.E_max)) else: break - + def run(self): # initialize the pmf @@ -394,28 +393,64 @@ class UmbrellaRunnerTest(unittest.TestCase): self.assertEquals(1, len(root_frames)) self.assertEquals((0, 2, 2), root_frames[0]) - def test_get_new_frames(self): + def test_generate_neighbor_list(self): + runner = UmbrellaRunner() + root = (1, 3) + neighbors = runner._generate_neighbor_list(root) + expected_neighbors = [ + (0, 2), + (0, 3), + (0, 4), + (1, 2), + (1, 3), + (1, 4), + (2, 2), + (2, 3), + (2, 4), + ] + self.assertEquals(neighbors, expected_neighbors) + + def test_generate_neighbor_list2(self): + runner = UmbrellaRunner() + root = (3, 1) + neighbors = runner._generate_neighbor_list(root) + expected_neighbors = [ + (2, 0), + (2, 1), + (2, 2), + (3, 0), + (3, 1), + (3, 2), + (4, 0), + (4, 1), + (4, 2), + ] + self.assertEquals(neighbors, expected_neighbors) + pass + + def test_generate_neighbor_list_3d(self): + runner = UmbrellaRunner() + root = (2, 2, 2) + neighbors = runner._generate_neighbor_list(root) + expected_neighbors = [] + for x in [-1, 0, 1]: + for y in [-1, 0, 1]: + for z in [-1, 0, 1]: + expected_neighbors.append((root[0]+x, root[1]+y, root[2]+z)) + self.assertEquals(len(neighbors), len(expected_neighbors)) + + + def test_is_in_pmf(self): runner = UmbrellaRunner() runner.cvs = np.array([ - (-3, 3, 1), - (0, 4, 1) + (-1, 1, 1), + (-1, 1, 1), + (-1, 1, 1) ]) runner.pmf = runner._init_pmf() - runner.pmf[0, 3] = 5 - runner.pmf[0, 2] = 2 - runner.pmf[0, 4] = 2 - root_frames = [(0, 3), (0, 2), (0, 4)] - new_frames = runner._get_new_frames(runner.pmf, root_frames) - - expected_new_frames = { - (0, 1): (0, 2), - (1, 1): (0, 2), - (1, 2): (0, 2), - (1, 3): (0, 2), - (1, 4): (0, 4) - } - self.assertEquals(len(expected_new_frames.keys()), len(new_frames.keys())) - self.assertDictEqual(expected_new_frames, new_frames) + self.assertFalse(runner._is_in_pmf((-1, 0, 0))) + self.assertFalse(runner._is_in_pmf((0, 0, 3))) + self.assertTrue(runner._is_in_pmf((0, 0, 0))) def test_get_new_frames_3d(self): runner = UmbrellaRunner() @@ -439,6 +474,35 @@ class UmbrellaRunnerTest(unittest.TestCase): self.assertEquals(26, len(new_frames.keys())) self.assertDictEqual(expected_new_frames, new_frames) + def test_get_new_frames(self): + runner = UmbrellaRunner() + runner.cvs = np.array([ + (-2, 2, 1), + (-2, 2, 1), + ]) + runner.pmf = np.array([ + [-1, -1, -1, -1, -1], + [-1, 34, 32, 26, -1], + [-1, 40, 42, 51, -1], + [-1, 28, 41, 37, -1], + [-1, -1, -1, -1, -1]]) + root_frames = [(1, 3), (3, 1)] + expected_new_frames = { + (0, 2): (1, 3), + (0, 3): (1, 3), + (0, 4): (1, 3), + (1, 4): (1, 3), + (2, 4): (1, 3), + (2, 0): (3, 1), + (3, 0): (3, 1), + (4, 0): (3, 1), + (4, 1): (3, 1), + (4, 2): (3, 1), + } + new_frames = runner._get_new_frames(runner.pmf, root_frames) + self.assertDictEqual(expected_new_frames, new_frames) + + if __name__ == '__main__': unittest.main() diff --git a/example.py b/example.py index 027d93d..945f071 100644 --- a/example.py +++ b/example.py @@ -13,12 +13,21 @@ class MyUmbrellaRunner(WHAM2DRunner): pmf_to_plot = deepcopy(self.pmf.T) pmf_to_plot[pmf_to_plot < 0] = None plt.figure() - plt.imshow(pmf_to_plot, origin="lower", cmap='jet') + plt.imshow(pmf_to_plot, origin="bottom", cmap='jet') + ticks = [(x,x) for x in [-3, -2, -1, 0, 1, 2, 3]] + tick_positions = [ self._get_index_for_lambdas(x)[0] for x in ticks ] + tick_labels = [ str(x[0]) for x in ticks ] + + plt.xticks(tick_positions, tick_labels) + plt.yticks(tick_positions, tick_labels) + + cb = plt.colorbar(pad=0.1) cb.set_label("kJ/mol") plt.savefig(filename) os.system("cp {} {}".format(filename, "pmf_current.pdf")) + def simulate_frames(self, lambdas, frames): print("{} new simulations:".format(len(lambdas))) counter = 0 @@ -35,17 +44,16 @@ class MyUmbrellaRunner(WHAM2DRunner): # print("Running {}".format(command)) os.system(command) - runner = MyUmbrellaRunner() runner.WHAM_EXEC = "/opt/wham/wham-2d/wham-2d" runner.cvs = np.array([ - (-3.1, 3.1, 0.1), - (-3.1, 3.1, 0.1), + (-3, 3, 0.2), + (-3, 3, 0.2), ]) -runner.cvs_init = (1, -1.4) +runner.cvs_init = (1.4, -1.4) runner.E_min = 10 -runner.E_max = 200 +runner.E_max = 100 runner.E_incr = 10 -runner.max_iterations = 25 +runner.max_iterations = 100 runner.run()