Mercurial > repos > chemteam > mdanalysis_distance
diff distance_multiple.py @ 5:d540ea77b909 draft
"planemo upload for repository https://github.com/galaxycomputationalchemistry/galaxy-tools-compchem/ commit 12823af06b7926cc56aaf9b59cfea9f16d342b8c"
author | chemteam |
---|---|
date | Thu, 06 Feb 2020 19:43:03 -0500 |
parents | |
children |
line wrap: on
line diff
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/distance_multiple.py Thu Feb 06 19:43:03 2020 -0500 @@ -0,0 +1,56 @@ +import argparse +import sys + +import MDAnalysis as mda +from MDAnalysis.analysis import distances + +import numpy as np + + +def parse_command_line(argv): + parser = argparse.ArgumentParser() + parser.add_argument('--itraj', help='input traj') + parser.add_argument('--istr', help='input str') + parser.add_argument('--itrajext', help='input traj ext') + parser.add_argument('--istrext', help='input str ext') + parser.add_argument('--list1', help='list 2') + parser.add_argument('--list2', help='list 2') + parser.add_argument('--output', help='output') + parser.add_argument('--header', dest='header', action='store_true') + return parser.parse_args() + + +args = parse_command_line(sys.argv) + +u = mda.Universe(args.istr, args.itraj, + topology_format=args.istrext, format=args.itrajext) + +list1 = np.loadtxt(args.list1, dtype=str, delimiter="\t", ndmin=1) +list2 = np.loadtxt(args.list2, dtype=str, delimiter="\t", ndmin=1) + +sel1 = [u.select_atoms(selection) for selection in list1] +sel2 = [u.select_atoms(selection) for selection in list2] + +d = np.empty((u.trajectory.n_frames, list1.shape[0], list2.shape[0]),) + +for ts in u.trajectory: + c_o_m1 = np.array([selection.center_of_mass() for selection in sel1]) + c_o_m2 = np.array([selection.center_of_mass() for selection in sel2]) + distances.distance_array(c_o_m1, c_o_m2, result=d[ts.frame]) + +d = np.hstack(( + np.array(np.reshape(np.arange( + 0, d.shape[0]), (d.shape[0], 1)), dtype=int), # add column w frame + np.reshape(d, (d.shape[0], d.shape[1] * d.shape[2])) +)) + +if args.header: + header = 'Frame\t' + '\t'.join( + ['-'.join(pair) for pair in zip( + sum([[n, ] * len(list2) for n in list1], []), + list(list2) * len(list1),)]).replace(' ', '_') +else: + header = '' + +np.savetxt(args.output, d, header=header, comments='', + fmt=['%d'] + ['%f'] * (d.shape[1] - 1), delimiter='\t')