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

only sample neighboring windows

Browse Source
This commit is contained in:
Daniel Bauer
2018-06-17 11:54:25 +02:00
parent 0c2cd3fc0a
commit 025f397bc1
3 changed files with 81 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
.gitignore vendored
View File

@@ -1,3 +1,4 @@
__pycache__
venv
.idea
plot

93
adaptiveumbrella/runner.py
View File

@@ -71,10 +71,12 @@ class UmbrellaRunner():
else: # if len differs, theres no index for every dimension
raise ValueError("{} has no index.".format(lambdas))
def _get_root_frames(self, pmf, E_max):
def _get_root_frames(self, pmf, frames, E_max):
""" returns the index of all positions in the pmf where the energy is
smaller E_max"""
selection = np.where((pmf <= E_max) & (pmf >= 0))
# select positions of the pmf where E <= E_max and that have already been sampled (frames > 0)
selection = np.where((pmf <= E_max) & (frames > 0))
zipped = list(zip(*selection))
return zipped
@@ -107,7 +109,7 @@ class UmbrellaRunner():
return False
return True
def _get_new_frames(self, pmf, root_frames):
def _get_new_frames(self, pmf, frames, 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} """
@@ -132,12 +134,11 @@ class UmbrellaRunner():
new_frames[n] = frame
# remove already sampled frames (where energy >= 0)
# remove already sampled frames (frames > 0)
new_frames_list = list(new_frames.keys())
for idx in range(len(new_frames_list)):
new_frame = new_frames_list[idx]
energy = pmf[new_frame]
if energy >= 0:
if frames[new_frame] > 0:
del(new_frames[new_frame])
return new_frames
@@ -146,7 +147,7 @@ class UmbrellaRunner():
def _main(self):
# get the initial simulation and surrounding frames
root_frames = [self._get_index_for_lambdas(self.cvs_init)]
new_frames = self._get_new_frames(self.pmf, root_frames)
new_frames = self._get_new_frames(self.pmf, self.sample_list, root_frames)
self.num_iterations = 0
self.E = self.E_min
@@ -174,11 +175,15 @@ class UmbrellaRunner():
print("Calculating new PMF")
self.pmf = self.calculate_new_pmf()
# update list of sampled windows
for new_frame in new_frames.keys():
self.sample_list[new_frame] = self.num_iterations
self.after_run_hook()
while self.E <= self.E_max:
root_frames = self._get_root_frames(self.pmf, self.E)
new_frames = self._get_new_frames(self.pmf, root_frames)
root_frames = self._get_root_frames(self.pmf, self.sample_list, self.E)
new_frames = self._get_new_frames(self.pmf, self.sample_list, root_frames)
if len(new_frames) == 0:
self.E += self.E_incr
@@ -191,6 +196,9 @@ class UmbrellaRunner():
# initialize the pmf
self.pmf = self._init_pmf()
# a list of sampled frames of the pmf (0=unsampled, 1,2,3.. = sampled )
self.sample_list = np.zeros(self.pmf.shape)
# start the simulation/evaluation loop
self._main()
print("Finished.")
@@ -280,9 +288,14 @@ class UmbrellaRunnerTest(unittest.TestCase):
(0, 4, 1)
])
runner.pmf = runner._init_pmf()
runner.sample_list = np.zeros(runner.pmf.shape)
runner.pmf[0, 3] = 5
runner.pmf[0, 2] = 2
root_frames = runner._get_root_frames(runner.pmf, 3)
runner.pmf[0, 1] = 2
runner.sample_list[0, 3] = 1
runner.sample_list[0, 2] = 1
root_frames = runner._get_root_frames(runner.pmf, runner.sample_list, 3)
self.assertEquals(1, len(root_frames))
self.assertEquals((0, 2), root_frames[0])
@@ -294,9 +307,13 @@ class UmbrellaRunnerTest(unittest.TestCase):
(0, 4, 1)
])
runner.pmf = runner._init_pmf()
runner.sample_list = np.zeros(runner.pmf.shape)
runner.pmf[0, 3, 3] = 5
runner.pmf[0, 2, 2] = 2
root_frames = runner._get_root_frames(runner.pmf, 3)
runner.pmf[0, 1, 2] = 2
runner.sample_list[0, 3, 3] = 5
runner.sample_list[0, 2, 2] = 2
root_frames = runner._get_root_frames(runner.pmf, runner.sample_list, 3)
self.assertEquals(1, len(root_frames))
self.assertEquals((0, 2, 2), root_frames[0])
@@ -359,28 +376,6 @@ class UmbrellaRunnerTest(unittest.TestCase):
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()
runner.cvs = np.array([
(-1, 1, 1),
(-1, 1, 1),
(-1, 1, 1)
])
runner.pmf = runner._init_pmf()
runner.pmf[1, 1, 1] = 5
root_frames = [(1, 1, 1)]
new_frames = runner._get_new_frames(runner.pmf, root_frames)
expected_new_frames = {}
for x in [0, 1, 2]:
for y in [0, 1, 2]:
for z in [0, 1, 2]:
if (x, y, z) != (1, 1, 1):
expected_new_frames[(x, y, z)] = (1, 1, 1)
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([
@@ -393,6 +388,9 @@ class UmbrellaRunnerTest(unittest.TestCase):
[-1, 40, 42, 51, -1],
[-1, 28, 41, 37, -1],
[-1, -1, -1, -1, -1]])
runner.sample_list = np.zeros(runner.pmf.shape)
runner.sample_list[runner.pmf > 0] = 1
runner.sample_list[runner.pmf < 0] = 0
root_frames = [(1, 3), (3, 1)]
expected_new_frames = {
(0, 2): (1, 3),
@@ -406,9 +404,36 @@ class UmbrellaRunnerTest(unittest.TestCase):
(4, 1): (3, 1),
(4, 2): (3, 1),
}
new_frames = runner._get_new_frames(runner.pmf, root_frames)
new_frames = runner._get_new_frames(runner.pmf, runner.sample_list, root_frames)
self.assertDictEqual(expected_new_frames, new_frames)
def test_get_new_frames_differing_pmf(self):
runner = UmbrellaRunner()
runner.cvs = np.array([
(-2, 2, 1),
(-2, 2, 1),
])
runner.pmf = np.array([
[-1, -1, -1, -1, -1],
[-1, -1, -1, -1, -1],
[-1, -1, 30, -1, -1],
[-1, -1, 20, -1, -1],
[-1, -1, -1, -1, -1]])
runner.sample_list = np.zeros(runner.pmf.shape)
runner.sample_list[2, 3] = 1
root_frames = [(2, 3)]
expected_new_frames = {
(1, 2): (2, 3),
(2, 2): (2, 3),
(3, 2): (2, 3),
(1, 3): (2, 3),
(3, 3): (2, 3),
(1, 4): (2, 3),
(2, 4): (2, 3),
(3, 4): (2, 3)
}
new_frames = runner._get_new_frames(runner.pmf, runner.sample_list, root_frames)
self.assertDictEqual(expected_new_frames, new_frames)
if __name__ == '__main__':

31
examples/example.py
View File

@@ -4,6 +4,8 @@ from copy import deepcopy
import matplotlib.pyplot as plt
import numpy as np
import sys
sys.path.append('..')
from adaptiveumbrella.runner import UmbrellaRunner
@@ -105,18 +107,25 @@ class MyUmbrellaRunner(WHAM2DRunner):
print("Writing new pmf to {}".format(filename))
pmf_to_plot = deepcopy(self.pmf.T)
pmf_to_plot[pmf_to_plot < 0] = None
plt.figure()
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 ]
frames_to_plot = deepcopy(self.sample_list.T)
frames_to_plot[frames_to_plot == 0] = None
plt.xticks(tick_positions, tick_labels)
plt.yticks(tick_positions, tick_labels)
cb = plt.colorbar(pad=0.1)
fig, (ax0, ax1) = plt.subplots(ncols=2, sharey=True)
im = ax0.imshow(pmf_to_plot, origin="bottom", cmap='jet')
cb = fig.colorbar(im, ax=ax0)
cb.set_label("kJ/mol")
im2 = ax1.imshow(frames_to_plot, origin="bottom")
cb2 = fig.colorbar(im2, ax=ax1, ticks=np.arange(0, self.max_iterations))
cb.set_label("cycle")
# 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 ]
# ax1.set_xticks(tick_positions, tick_labels)
# ax1.set_yticks(tick_positions, tick_labels)
# ax2.set_xticks(tick_positions, tick_labels)
# ax2.set_yticks(tick_positions, tick_labels)
plt.savefig(filename)
os.system("cp {} {}".format(filename, "tmp/pmf_current.pdf"))
@@ -146,6 +155,6 @@ runner.cvs_init = (1.4, -1.4)
runner.E_min = 10
runner.E_max = 100
runner.E_incr = 10
runner.max_iterations = 30
runner.max_iterations = 1
runner.run()