initial commit
This commit is contained in:
Daniel Bauer
6
.idea/vcs.xml
generated
Normal file
6
.idea/vcs.xml
generated
Normal file
@@ -0,0 +1,6 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<project version="4">
|
||||
<component name="VcsDirectoryMappings">
|
||||
<mapping directory="$PROJECT_DIR$" vcs="Git" />
|
||||
</component>
|
||||
</project>
|
||||
257
adaptive_umbrella.py
Executable file
257
adaptive_umbrella.py
Executable file
@@ -0,0 +1,257 @@
|
||||
#!/usr/bin/env python3
|
||||
import numpy as np
|
||||
from copy import deepcopy
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
|
||||
#%%
|
||||
class UmbrellaRunner(ABC):
|
||||
def init_pmf(self):
|
||||
""" returns a NxM matrix filled with -1 where N and M are the total number
|
||||
of lambda frames along the reaction coordinates as determined by
|
||||
self.lambda_min, self.lambda_max and self.lambda_delta """
|
||||
ranges = []
|
||||
for dimen in range(len(self.lambda_delta)):
|
||||
ranges.append(np.arange(self.lambda_min[dimen], self.lambda_max[dimen]+self.lambda_delta[dimen], self.lambda_delta[dimen]))
|
||||
|
||||
mesh = np.meshgrid(*ranges)
|
||||
pmf = np.zeros(np.dstack(mesh).shape[:-1]) # this is magic
|
||||
pmf[:] = -1
|
||||
return pmf
|
||||
|
||||
def get_lambdas_for_index(self, idx):
|
||||
""" takes a coordinate tuple and returns corresponding lambda values """
|
||||
lambdas = self.lambda_min + idx*self.lambda_delta
|
||||
return np.round(lambdas, 10)
|
||||
|
||||
def get_index_for_lambdas(self, lambdas):
|
||||
""" takes a lambda tuple and returns corresponding indeces of the pmf instance variable"""
|
||||
idx = []
|
||||
for dimen in range(len(lambdas)):
|
||||
r = np.arange(self.lambda_min[dimen], self.lambda_max[dimen]+self.lambda_delta[dimen], self.lambda_delta[dimen])
|
||||
|
||||
for i in range(len(r)):
|
||||
if abs(r[i]-lambdas[dimen]) < 0.00001:
|
||||
idx.append(i)
|
||||
break
|
||||
return idx
|
||||
|
||||
# TODO make this work with more then 2 dimensions
|
||||
def get_root_frames(self, pmf, E_max):
|
||||
""" returns the index of all positions in the pmf where the energy is
|
||||
smaller W_max and greater 0 """
|
||||
selection = np.where((pmf <= E_max) & (pmf >= 0))
|
||||
frames = []
|
||||
for i in range(len(selection[0])):
|
||||
frames.append((selection[0][i], selection[1][i]))
|
||||
return frames
|
||||
|
||||
# TODO make this work with more then 2 dimensions
|
||||
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 their
|
||||
# root frames
|
||||
new_frames = {}
|
||||
for frame in root_frames:
|
||||
for x in [-1, 0, 1]:
|
||||
for y in [-1, 0, 1]:
|
||||
new_frame = list(frame)
|
||||
new_frame[0] += x
|
||||
new_frame[1] += y
|
||||
if new_frame[0] < 0 or new_frame[0]+1 > len(self.pmf[0]) or new_frame[1] < 0 or new_frame[1]+1 > len(self.pmf[1]):
|
||||
continue
|
||||
try:
|
||||
old_root = new_frames[tuple(new_frame)]
|
||||
if old_root is not None:
|
||||
old_start_energy = pmf[old_root[0], old_root[1]]
|
||||
new_start_energy = pmf[frame[0], frame[1]]
|
||||
if old_start_energy > new_start_energy:
|
||||
new_frames[tuple(new_frame)] = frame
|
||||
except KeyError:
|
||||
new_frames[tuple(new_frame)] = frame
|
||||
|
||||
# remove already sampled frames (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.lambda_init)]
|
||||
new_frames = self.get_new_frames(self.pmf, root_frames)
|
||||
|
||||
self.num_iterations = 0
|
||||
|
||||
# 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
|
||||
|
||||
self.E = self.E_min
|
||||
print("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~")
|
||||
print("Iteration: {} (max={})".format(self.num_iterations, self.max_iterations))
|
||||
new_lambdas = [ self.get_lambdas_for_index(x) for x in new_frames.keys() ]
|
||||
|
||||
print("Running simulations")
|
||||
self.simulate_frames(new_frames, new_lambdas)
|
||||
|
||||
print("Calculating new PMF")
|
||||
self.pmf = self.wham()
|
||||
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):
|
||||
self.pmf = self.init_pmf()
|
||||
self.main()
|
||||
print("Umbrella sampling finished.")
|
||||
|
||||
@abstractmethod
|
||||
def simulate_frames(new_frames, new_lambdas):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def wham():
|
||||
pass
|
||||
|
||||
def after_run_hook(self):
|
||||
pass
|
||||
|
||||
if __name__ == "__main__":
|
||||
from copy import deepcopy
|
||||
import subprocess
|
||||
import os
|
||||
import pandas as pd
|
||||
import matplotlib.pyplot as plt
|
||||
class MyUmbrellaRunner(UmbrellaRunner):
|
||||
|
||||
def wham(self):
|
||||
# copy colvar files
|
||||
os.system("mkdir -p WHAM")
|
||||
for f in os.listdir("sim"):
|
||||
with open("sim/{}/COLVAR".format(f), "r") as i:
|
||||
with open("WHAM/{}.xvg".format(f), 'w') as o:
|
||||
for line in i.readlines()[100:]:
|
||||
o.write(line)
|
||||
|
||||
|
||||
# generate metadata.dat
|
||||
x_vals = []
|
||||
y_vals = []
|
||||
metadata_file = "WHAM/{}_metadata.dat".format(self.num_iterations)
|
||||
with open(metadata_file, 'w') as o:
|
||||
for f in os.listdir("sim"):
|
||||
prefix, x, y = f.split("_")
|
||||
x_vals.append(float(x))
|
||||
y_vals.append(float(y))
|
||||
o.write("WHAM/{}.xvg {} {} {} {}\n".format(f, x,y, 100, 100))
|
||||
x_vals = np.array(list(set(x_vals)))
|
||||
y_vals = np.array(list(set(y_vals)))
|
||||
|
||||
# min_x = x_vals.min() - self.lambda_delta[0]
|
||||
# min_y = y_vals.min() - self.lambda_delta[1]
|
||||
# max_x = x_vals.max() + self.lambda_delta[0]
|
||||
# max_y = y_vals.max() + self.lambda_delta[1]
|
||||
min_x, min_y = self.lambda_min
|
||||
max_x, max_y = self.lambda_max
|
||||
frames_x, frames_y = 1002, 1002
|
||||
print("Running WHAM-2d:")
|
||||
wham_output = "WHAM/{}_freeenergy.dat".format(self.num_iterations)
|
||||
cmd = "/opt/wham/wham-2d/wham-2d Px=pi {min_x} {max_x} {frames_x} Py=pi {min_y} {max_y} {frames_y} 0.1 298 0 {metafile} {outfile} 0".format(
|
||||
min_x=min_x,
|
||||
max_x=max_x,
|
||||
frames_x=frames_x,
|
||||
min_y=min_y,
|
||||
max_y=max_y,
|
||||
frames_y=frames_y,
|
||||
metafile=metadata_file,
|
||||
outfile=wham_output
|
||||
)
|
||||
print(cmd)
|
||||
os.system(cmd)
|
||||
|
||||
print("Update pmf from wham")
|
||||
df = pd.read_csv(wham_output, 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
|
||||
|
||||
def simulate_frames(self, new_frames, new_lambdas):
|
||||
print("{} new simulations:".format(len(new_lambdas)))
|
||||
counter = 0
|
||||
|
||||
threads = []
|
||||
for f in new_lambdas:
|
||||
counter += 1
|
||||
if os.path.exists("sim/sim_{}_{}/COLVAR".format(*f)):
|
||||
print("{}) Skipping lambdas={}/{}: COLVAR exists".format(counter, *f))
|
||||
continue
|
||||
|
||||
print("{}) Simulate lambda1={}, lambda2={}".format(counter, *f))
|
||||
command = "bash sim.sh {} {} 2>&1 > run.log".format(*f)
|
||||
# print("Running {}".format(command))
|
||||
os.system(command)
|
||||
|
||||
def after_run_hook(self):
|
||||
filename = "pmf_{}.pdf".format(self.num_iterations)
|
||||
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="lower", cmap='jet')
|
||||
cb = plt.colorbar(pad=0.1)
|
||||
cb.set_label("kJ/mol")
|
||||
plt.savefig(filename)
|
||||
os.system("cp {} {}".format(filename, "pmf_current.pdf"))
|
||||
|
||||
|
||||
|
||||
runner = MyUmbrellaRunner()
|
||||
runner.lambda_max = np.array((3.1, 3.1))
|
||||
runner.lambda_min = -runner.lambda_max
|
||||
runner.lambda_delta = np.array((0.1, 0.1))
|
||||
runner.lambda_init = np.array((1,-1.4))
|
||||
runner.E_min = 5
|
||||
runner.E_max = 100
|
||||
runner.E_incr = 10
|
||||
runner.max_iterations = 10
|
||||
|
||||
runner.run()
|
||||
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user