258 lines
10 KiB
Python
Executable File
258 lines
10 KiB
Python
Executable File
#!/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()
|
|
|
|
|
|
|
|
|