Generate TRR file (Python)

From Computational Biophysics and Materials Science Group
Jump to: navigation, search

Install MDAnalysis first

# in this file, we have 200000 lines file 'data', each line is 'x y z fx fy fz', 
# this file contains 40 atoms' trajectory, 50000 frames
# another file 50000 lines - 'data2' include pbc size for each frame

# the following 5 parameters you need to change on your need
N = 40    # number of  atoms
FRAME = 50000 # number of frames
OUTPUT = './output.trr'
INPUT1 = './data'
INPUT2 = './data2'
COEFFICIENT = 1
#
import MDAnalysis
from MDAnalysis import Writer
W = Writer(OUTPUT, numatoms=N)
ts = MDAnalysis.coordinates.TRR.Timestep(N)
ts.has_f = True

f = open(INPUT1)
lines = f.readlines()
print('this file has ' + str(len(lines)) + ' lines')
f2 = open(INPUT2)
lines2 = f2.readlines()
print('this file has ' + str(len(lines2)) + ' lines')

for frame in range(0, FRAME):
	ts.lmbda = 0
	ts.frame = frame
	ts.step = frame
	ts.time = frame * COEFFICIENT
	pbc_size = [float(x) for x in lines2[frame]].split()
	ts._unitcell[0] = [pbc_size[0], 0, 0]
	ts._unitcell[1] = [0, pbc_size[1], 0]
	ts._unitcell[2] = [0, 0, pbc_size[2]]
	for n in range(0, N):
		compo = lines[frame * N + n].split()
		ts._pos[n] = [float(x) * 0.1 for x in compo][0:3]
		ts._forces[n] = [float(x) for x in compo][3:6]

	W.write(ts)
	print(frame + 1)