2_get_us_windows.py

#!/usr/bin/python

import os, sys, numpy, glob, decimal
from collections import OrderedDict
from subprocess import call
import scipy.spatial.distance as distance

#### set GROMACS path ################

gromacs = "/home/p/pkasson/nobackup/"

######################################

'''
python script to get the umbrella sampling windows from z coordinate files

This script accompanies the following publications:
Simulation-guided engineering of antibiotics for improved bacterial uptake, Ricardo Ferreir
a and Peter Kasson, doi:10.1101/2020.10.08.330332.
Antibiotic Uptake Across Gram-Negative Outer Membranes: Better Predictions Towards Better Antibiotics, Ricardo Ferreira and Peter Kasson, doi:10.1021/acsinfecdis.9b00201.

With this script we intent to:
	- get the z coordinate from the pulling run (gmx traj)
	- split into X bins with width defined in keyword 'delta'
	- calls "build_umbrella.sh" script to generate the apropriate gro files from the pulling trajectory
	- no further input is needed (all generated by the script)
	
### version 0.0.4 ###
	- added automatic determination of endpoint (based on distances from startpoint)
	- bin width is now automatically generated from the pulled path
	
### version 0.0.3 ###
	- added function to automatically determine bin width (from ligand 3D coordinates)
	  (needed because smaller ligands often needs smaller bins due to poor histogram overlap)
	- add mdp options to python script
	- add command to run GROMACS
	- gromacs path set within the script
	- script can now be run from main directory

### version 0.0.2 ###
	- changed pull directory name (./PULL)
	- generates input files for US and WHAM
	- bin width is now defined by the user

### version 0.0.1 ###
In this first version:
	- TPR and XTC located in ../PULL_VEC directory
	- it assumes that the ligand has the #21 position in the index file
	- it assumes a 0.350 nm bin width by default (changed only by editingt the file)
'''

#### initialize variables ############
start = 0.0
delta = 0.000
z_list = OrderedDict()
windows = [0.0000]
us_list = 1
points = []
distances = []
times = []


#### get name from ligand ############
ligand = open('ligand.gro', 'r')
lig_lines = ligand.readlines()
lig_name = lig_lines[3][5:8]
ligand.close()

############# DEPRECATED #############
# for lines in lig_lines[2:-2]:
#	newpoints = ()
#	newpoints = newpoints + (float(lines[22:28]), float(lines[30:36]), float(lines[38:44]))
#	points.append(newpoints)
# ligand.close()
#
# dist = distance.pdist(points)					# 3D distances
# if numpy.median(dist) > 0.60:				
# 	 delta = 0.350						# large (long) molecules
#elif numpy.median(dist) < 0.50:			
#	 delta = 0.330						# small (compact) molecules
#else:									
#	 delta = 0.340						# normal bin width
######################################

#### create subdirectory for US ######
if not os.path.exists('US'):
    os.makedirs('US')
os.chdir('./US')

#### generate the z_coordinate.xvg ###
ligand_file = open('ligand.txt', 'w')
ligand_file.write('21\n)')
ligand_file.write('q\n')
ligand_file.close()
call(gromacs + "/bin/gmx traj -quiet -s ../PULL/*.tpr -f ../PULL/*.xtc -nox -noy -ox z_coordinate.xvg -pbc -com -nojump < ligand.txt", shell=True)

#### specify how to start ############
coord_file = open('z_coordinate.xvg', 'r')
coord_lines = coord_file.readlines()[24:]
for l in coord_lines:
	time, position = l.split('\t')
	z_list.update({int(time):float(position)})
coord_file.close()
	
#### get min and max from distance ###
call(gromacs + "/bin/gmx distance -s ../PULL/*.tpr -f ../PULL/*.xtc -selrpos mol_com -seltype res_com -n ../index.ndx -oav distance_com.xvg -select 'com of group 38 plus com of group 21'", shell=True)
distance_file = open("distance_com.xvg", "r")
distance_lines = distance_file.readlines()
for d in distance_lines:
	if '#' not in d:
		if '@' not in d:
			time, position = d.strip('\n').split()
			distances.append(float(position))
			times.append(float(time))

initial_position = distances[50]  # PK mod
initial_time = times[50]
final_time = 0.0
final_position = 0.0

for b in distances:
	if b < initial_position:
		final_position = b
		distance_to_time = distances.index(b)

final_time = times[distance_to_time]
distance_file.close()

#### determine bin width from dist ###
path = float(z_list[0]) - float(z_list[final_time])
delta = round((path / 26),5)

#### define start value ##############
start = z_list[0]
bin_file = open('output_bins.txt', 'w')
bin_file.write('Time (ps)\tPosition (nm)\n')
bin_file.write('0.0000\t\t' + str(z_list[0]) + '\n')
os.system('clear')

for i in z_list:
	if i >= int(final_time):
		break
	if z_list[i] < (start - delta):
		windows.append(int(i))
		bin_file.write(str(i) + '\t\t' + str(z_list[i]) + '\n')
		start = z_list[i]
bin_file.close()

#### generate umbrella windows #######
system_file = open('system.txt', 'w')
system_file.write('0\n)')
system_file.write('q\n')
system_file.close()
for w in windows:
	call(gromacs + "/bin/gmx trjconv -quiet -s ../PULL/*.tpr -f ../PULL/*.xtc -dump " + str(w) + " -pbc mol -nice 0 -o US_" + str(us_list) + ".gro < system.txt", shell=True)
	us_list += 1
	
#### delete unwanted files ###########
os.system('rm ligand.txt system.txt')
	
#### generate input files ############
mdp = open('us_wham_40ns.mdp', 'w')
mdp.write('define = \n')
mdp.write('integrator = md \n')
mdp.write('nsteps = 20000000 \n')
mdp.write('dt = 0.002 \n')
mdp.write('nstxout = 500000 \n')
mdp.write('nstvout = 500000 \n')
mdp.write('nstenergy = 500000 \n')
mdp.write('nstlog = 500000 \n')
mdp.write('nstxout-compressed = 500000 \n')
mdp.write('continuation = yes \n')
mdp.write('constraint_algorithm = lincs \n')
mdp.write('constraints = all-bonds \n')
mdp.write('lincs_iter = 1 \n')
mdp.write('lincs_order = 4 \n')
mdp.write('ns_type = grid \n')
mdp.write('rlist = 1.2 \n')
mdp.write('rcoulomb = 1.2 \n')
mdp.write('rvdw = 1.2 \n')
mdp.write('vdwtype = Cut-off \n')
mdp.write('vdw-modifier = Force-switch \n')
mdp.write('rvdw_switch = 1.0 \n')
mdp.write('coulombtype = PME \n')
mdp.write('pme_order = 4 \n')
mdp.write('fourierspacing = 0.16 \n')
mdp.write('cutoff-scheme = Verlet \n')
mdp.write('nstlist = 40 \n')
mdp.write('tcoupl = Nose-Hoover \n')
mdp.write('tc-grps = Protein_OM ' + str(lig_name) + ' Water_and_ions \n')
mdp.write('tau_t = 1.0 1.0 1.0 \n')
mdp.write('ref_t = 310 310 310 \n')
mdp.write('nsttcouple = 1 \n')
mdp.write('nhchainlength = 1 \n')
mdp.write('pcoupl = Parrinello-Rahman \n')
mdp.write('pcoupltype = semiisotropic \n')
mdp.write('tau_p = 5.0 \n')
mdp.write('ref_p = 1.0 1.0 \n')
mdp.write('compressibility = 4.5e-5 0 \n')
mdp.write('pbc = xyz \n')
mdp.write('DispCorr = EnerPres \n')
mdp.write('gen_vel = no \n')
mdp.write('nstcomm = 1 \n')
mdp.write('nstcalcenergy = 1 \n')
mdp.write('comm-mode = Linear \n')
mdp.write('comm-grps = Protein_OM ' + str(lig_name) + ' Water_and_ions \n')
mdp.write('pull = yes \n')
mdp.write('pull_ngroups = 2 \n')
mdp.write('pull_ncoords = 1 \n')
mdp.write('pull_group1_name = ' + str(lig_name) + '\n')
mdp.write('pull_group2_name = Protein \n')
mdp.write('pull_nstxout = 10 \n')
mdp.write('pull_nstfout = 10 \n')
mdp.write('pull_coord1_type = umbrella \n')
mdp.write('pull_coord1_geometry = direction-periodic \n')
mdp.write('pull_coord1_vec = 0 0 -1\n')
mdp.write('pull_pbc_ref_prev_step_com = yes\n')
mdp.write('pull_group2_pbcatom = 3895\n')
mdp.write('pull_coord1_groups = 1 2 \n')
mdp.write('pull_coord1_dim = N N Y \n')
mdp.write('pull_coord1_rate = 0 \n')
mdp.write('pull_coord1_k = 1000 \n')
mdp.write('pull_coord1_start = yes \n')
mdp.close()

#### run jobs ########################
grofiles = sorted(glob.glob('*.gro'))
if not os.path.exists('WHAM'):
    os.makedirs('WHAM')
for fl in grofiles:
	call(gromacs + "/bin/gmx grompp -f us_wham_40ns.mdp -c " + str(fl) + " -p ../topol_new.top -n ../index.ndx -o ./WHAM/" + str(fl[:-4]) + ".tpr -maxwarn 1", shell=True)	# generate TPR files
	os.system("rm mdout.mdp")
os.chdir('./WHAM')

#### optionally run GROMACS job ######
start = raw_input("Do you want to start GROMACS mdrun (y/n)? ")
if start == "y":
	for fl in grofiles:
		call(gromacs + "/bin/gmx -quiet mdrun -s " + str(fl[:-4]) + ".tpr -nb gpu -deffnm " + str(fl[:-4]) + " -px pullx_" + str(fl[:-4]) + ".xvg -pf pullf_" + str(fl[:-4]) + ".xvg", shell=True)

#### quit the script #################
quit()