bauer/adaptiveumbrella
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

first running version

Browse Source
This commit is contained in:
Daniel Bauer
2018-06-17 09:33:16 +02:00
parent 4ed7ea9410
commit 65fc1f6c35
2 changed files with 135 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files

160
adaptive_umbrella.py
View File

@@ -76,55 +76,49 @@ class UmbrellaRunner():
smaller E_max"""
selection = np.where((pmf <= E_max) & (pmf >= 0))
zipped = list(zip(*selection))
return zipped
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)
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:
dimen = len(coord)
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]
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):
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} """
# 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:
@@ -166,6 +160,11 @@ class UmbrellaRunner():
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()])
@@ -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()

22
example.py
View File

@@ -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()