511 lines
17 KiB
Python
Executable File
511 lines
17 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
import numpy as np
|
|
import unittest
|
|
from copy import deepcopy
|
|
|
|
#%%
|
|
class UmbrellaRunner():
|
|
"""
|
|
Runner class for adaptive umbrella sampling
|
|
|
|
Attributes:
|
|
cvs (numpy array): A multidimensional array of all collective variables (cvs) in the form [cv_min, cv_max, cv_delta]
|
|
Example: to sample a 2 dimensional grid where cv1 and cv2 go from 0 to 3 with stepsize 0.5, use
|
|
np.array([[0,3,1],[0,3,1]])
|
|
cvs_init (tuple): Starting coordinates for umbrella sampling. Must be a tuple with the same dimension as cvs
|
|
E_min (float, default=0): Starting energy for exporative umbrella sampling
|
|
E_max (float, default=inf): Final energy. Umbrella sampling is stopped if no frames with E < E_max are found
|
|
E_incr (float, default=1): E_min is incremented by this until E_max is reached
|
|
max_iterations (int, default=-1): Max. number of iterations before umbrella sampling stops. -1 for infinite sampling
|
|
|
|
"""
|
|
|
|
def __init__(self):
|
|
self.max_iterations = -1
|
|
self.E_min = 0
|
|
self.E_max = np.inf
|
|
self.E_incr = 1
|
|
|
|
def _get_pmf_shape(self):
|
|
""" returns the shape of the pmf according to the cvs """
|
|
shape = []
|
|
for dimen in self.cvs:
|
|
windows = np.arange(*dimen)
|
|
size = len(windows)
|
|
if windows[-1] % dimen[-1] == 0:
|
|
size += 1
|
|
shape.append(size)
|
|
return shape
|
|
|
|
|
|
def _init_pmf(self):
|
|
""" returns an empty matrix where each dimension equals the number of frames along the corresponding reaction
|
|
coordinate """
|
|
shape = self._get_pmf_shape()
|
|
pmf = np.empty(shape)
|
|
pmf.fill(-1)
|
|
return pmf
|
|
|
|
|
|
def _get_lambdas_for_index(self, idx):
|
|
""" takes a coordinate tuple of the pmf and returns corresponding lambda values """
|
|
|
|
lambdas = self.cvs.T[0] + idx * self.cvs.T[2]
|
|
return tuple(np.round(lambdas, 10))
|
|
|
|
def _get_index_for_lambdas(self, lambdas):
|
|
""" takes a lambda tuple and returns corresponding indexes of the pmf
|
|
TODO: faster implementation required
|
|
"""
|
|
idx = []
|
|
for dimen in range(len(lambdas)):
|
|
cv = self.cvs[dimen]
|
|
r = np.arange(cv[0], cv[1]+cv[2], cv[2])
|
|
|
|
for i in range(len(r)):
|
|
if abs(r[i]-lambdas[dimen]) < 0.00001:
|
|
idx.append(i)
|
|
break
|
|
if len(idx) == len(lambdas):
|
|
return tuple(idx)
|
|
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):
|
|
""" returns the index of all positions in the pmf where the energy is
|
|
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} """
|
|
|
|
# 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 = self._generate_neighbor_list(frame)
|
|
|
|
# remove neighbors that are not inside the pmf
|
|
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:
|
|
try:
|
|
root_energy = pmf[frame]
|
|
old_root = new_frames[n]
|
|
old_root_energy = pmf[old_root]
|
|
if root_energy < old_root_energy:
|
|
new_frames[n] = frame
|
|
except KeyError:
|
|
new_frames[n] = frame
|
|
|
|
|
|
# remove already sampled frames (where energy >= 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:
|
|
del(new_frames[new_frame])
|
|
|
|
return new_frames
|
|
|
|
|
|
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)
|
|
|
|
self.num_iterations = 0
|
|
self.E = self.E_min
|
|
|
|
# outer main loop: increase E and calculate PMF until E > E_max
|
|
while True:
|
|
|
|
# stop if max iterations is reached
|
|
self.num_iterations += 1
|
|
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()])
|
|
|
|
self.pre_run_hook()
|
|
print("Running simulations")
|
|
self.simulate_frames(lambdas, new_frames)
|
|
|
|
print("Calculating new PMF")
|
|
self.pmf = self.calculate_new_pmf()
|
|
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)
|
|
|
|
if len(new_frames) == 0:
|
|
self.E += self.E_incr
|
|
print("Max energy increased to {} (max={})".format(self.E, self.E_max))
|
|
else:
|
|
break
|
|
|
|
|
|
def run(self):
|
|
# initialize the pmf
|
|
self.pmf = self._init_pmf()
|
|
|
|
# start the simulation/evaluation loop
|
|
self._main()
|
|
print("Finished.")
|
|
|
|
def simulate_frames(self, frames, lambdas):
|
|
""" This should be implemented to run the simulation for each frame in ```frames```
|
|
with ```lambdas```. Both variables are dictionaries where new values are keys and root values are values
|
|
"""
|
|
print("New lambda values to simulate: {}".format(lambdas))
|
|
print("(You should implement `simulate_frames` method yourself.)")
|
|
|
|
def calculate_new_pmf(self):
|
|
""" This should be overwritten to calculate the PMF based on the new simulations and
|
|
return a pmf of correct shape/spacing according to the cvs"""
|
|
pass
|
|
|
|
|
|
def after_run_hook(self):
|
|
""" This can be implemented to hook into the simulation cycle after the reevaluation of the pmf """
|
|
pass
|
|
|
|
def pre_run_hook(self):
|
|
""" This can be implemented to hook into the simulation cycle before the simulation runs """
|
|
pass
|
|
|
|
class WHAM2DRunner(UmbrellaRunner):
|
|
""" Umbrella runner implementation that uses wham-2d to perform
|
|
the pmf calculation.
|
|
|
|
Attributes:
|
|
WHAM_EXEC: path to wham executeable
|
|
|
|
"""
|
|
|
|
def __init__(self):
|
|
UmbrellaRunner.__init__(self)
|
|
self.WHAM_EXEC = 'wham-2d'
|
|
|
|
def calculate_new_pmf(self):
|
|
import os
|
|
from shutil import copyfile
|
|
|
|
simulation_dir = "simulations"
|
|
print("Collecting sampling data from simulations folder")
|
|
|
|
# collect COLVARs
|
|
wham_dir = "WHAM/"
|
|
if not os.path.exists(wham_dir):
|
|
os.makedirs(wham_dir)
|
|
|
|
for folder in os.listdir(simulation_dir):
|
|
src = os.path.join(simulation_dir, folder, "COLVAR")
|
|
dst = os.path.join(wham_dir, folder + ".xvg")
|
|
copyfile(src, dst)
|
|
|
|
# create metadata file
|
|
metadata_file = os.path.join(wham_dir, "{}_metadata.dat".format(self.num_iterations))
|
|
fc_x = 100
|
|
fc_y = 100
|
|
with open(metadata_file, 'w') as out:
|
|
for f in os.listdir(simulation_dir):
|
|
prefix, x, y = f.split("_")
|
|
out.write("WHAM/{}.xvg {} {} {} {}\n".format(f, x, y, fc_x, fc_y))
|
|
|
|
# run WHAM2d
|
|
print("Running WHAM-2d")
|
|
wham_output = os.path.join(wham_dir, "{}_freeenergy.dat".format(self.num_iterations))
|
|
periodicity_x = "pi"
|
|
periodicity_y = "pi"
|
|
tolerance = 0.1
|
|
frames_x, frames_y = 1002, 1002
|
|
min_x = self.cvs[0][0]
|
|
max_x = self.cvs[0][1]
|
|
min_y = self.cvs[1][0]
|
|
max_y = self.cvs[1][1]
|
|
|
|
cmd = "{exec} Px={px} {min_x} {max_x} {frames_x} Py={py} {min_y} {max_y} {frames_y} {tol} 298 0 {metafile} {outfile} 0".format(
|
|
exec=self.WHAM_EXEC,
|
|
px=periodicity_x,
|
|
min_x=min_x,
|
|
max_x=max_x,
|
|
frames_x=frames_x,
|
|
py=periodicity_y,
|
|
min_y=min_y,
|
|
max_y=max_y,
|
|
frames_y=frames_y,
|
|
tol=tolerance,
|
|
metafile=metadata_file,
|
|
outfile=wham_output
|
|
)
|
|
print(cmd)
|
|
os.system(cmd)
|
|
|
|
# read wham to new pmf
|
|
return self.read_pmf(wham_output)
|
|
|
|
def read_pmf(self, pmf_path):
|
|
import pandas as pd
|
|
print("Update PMF from WHAM")
|
|
df = pd.read_csv(pmf_path, delim_whitespace=True, names=['x', 'y', 'e', 'pro'], skiprows=1,
|
|
index_col=None)
|
|
df = df.replace([np.inf, -np.inf], np.nan).dropna(subset=['e'], how='all')
|
|
new_pmf = deepcopy(self.pmf)
|
|
for x in range(new_pmf.shape[0]):
|
|
for y in range(new_pmf.shape[1]):
|
|
lambdax, lambday = self._get_lambdas_for_index((x, y))
|
|
x_selection = (df.x - lambdax).abs() < 0.01
|
|
y_selection = (df.y - lambday).abs() < 0.01
|
|
selected_energies = df[(x_selection) & (y_selection)].e
|
|
if len(selected_energies) == 0:
|
|
new_pmf[x, y] = -1
|
|
else:
|
|
new_pmf[x, y] = selected_energies.iloc[0]
|
|
|
|
return new_pmf
|
|
|
|
|
|
|
|
class UmbrellaRunnerTest(unittest.TestCase):
|
|
|
|
def test_init_pmf_3d(self):
|
|
runner = UmbrellaRunner()
|
|
runner.cvs = np.array([
|
|
(-3, 3, 1),
|
|
(-3, 3, 1),
|
|
(-3, 2, 1)
|
|
])
|
|
pmf = runner._init_pmf()
|
|
|
|
expected_shape = (7, 7, 6)
|
|
self.assertEquals(pmf.shape, expected_shape)
|
|
|
|
def test_init_pmf_odd(self):
|
|
runner = UmbrellaRunner()
|
|
runner.cvs = np.array([
|
|
(-3, 3, 1),
|
|
(-3.5, 3, 1)
|
|
])
|
|
pmf = runner._init_pmf()
|
|
expected_shape = (7, 7)
|
|
self.assertEquals(pmf.shape, expected_shape)
|
|
|
|
def test_get_lambdas_for_index(self):
|
|
runner = UmbrellaRunner()
|
|
runner.cvs = np.array([
|
|
(-3, 3, 1),
|
|
(0, 4, 1)
|
|
])
|
|
|
|
lambdas = runner._get_lambdas_for_index((0, 0))
|
|
self.assertAlmostEqual(lambdas, (-3, 0))
|
|
|
|
lambdas = runner._get_lambdas_for_index((3, 2))
|
|
self.assertEquals(lambdas, (0, 2))
|
|
|
|
def test_get_index_for_lambdas(self):
|
|
runner = UmbrellaRunner()
|
|
runner.cvs = np.array([
|
|
(-3, 3, 1),
|
|
(0, 4, 1)
|
|
])
|
|
|
|
index = runner._get_index_for_lambdas((3, 3))
|
|
self.assertEquals(index, (6, 3))
|
|
|
|
def test_get_index_for_lambdas_error(self):
|
|
runner = UmbrellaRunner()
|
|
runner.cvs = np.array([
|
|
(-3, 3, 1),
|
|
(0, 4, 1)
|
|
])
|
|
|
|
with self.assertRaises(ValueError) as error:
|
|
runner._get_index_for_lambdas((3, 2.5))
|
|
|
|
def test_get_root_frames(self):
|
|
runner = UmbrellaRunner()
|
|
runner.cvs = np.array([
|
|
(-3, 3, 1),
|
|
(0, 4, 1)
|
|
])
|
|
runner.pmf = runner._init_pmf()
|
|
runner.pmf[0, 3] = 5
|
|
runner.pmf[0, 2] = 2
|
|
root_frames = runner._get_root_frames(runner.pmf, 3)
|
|
self.assertEquals(1, len(root_frames))
|
|
self.assertEquals((0, 2), root_frames[0])
|
|
|
|
def test_get_root_frames_3d(self):
|
|
runner = UmbrellaRunner()
|
|
runner.cvs = np.array([
|
|
(-3, 3, 1),
|
|
(0, 4, 1),
|
|
(0, 4, 1)
|
|
])
|
|
runner.pmf = runner._init_pmf()
|
|
runner.pmf[0, 3, 3] = 5
|
|
runner.pmf[0, 2, 2] = 2
|
|
root_frames = runner._get_root_frames(runner.pmf, 3)
|
|
self.assertEquals(1, len(root_frames))
|
|
self.assertEquals((0, 2, 2), root_frames[0])
|
|
|
|
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([
|
|
(-1, 1, 1),
|
|
(-1, 1, 1),
|
|
(-1, 1, 1)
|
|
])
|
|
runner.pmf = runner._init_pmf()
|
|
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()
|
|
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([
|
|
(-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()
|
|
|
|
|