Mercurial > repos > chemteam > mdanalysis_hbonds
comparison extract_rmsd.py @ 4:4c36f5ad2799 draft
"planemo upload for repository https://github.com/galaxycomputationalchemistry/galaxy-tools-compchem/ commit 1b23e024af45cc0999d9142d07de6897d4189ec2"
author | chemteam |
---|---|
date | Mon, 24 Aug 2020 16:53:30 -0400 |
parents | |
children | dfda5e713926 |
comparison
equal
deleted
inserted
replaced
3:5efd0c95f97e | 4:4c36f5ad2799 |
---|---|
1 import argparse | |
2 import json | |
3 | |
4 import MDAnalysis as m | |
5 from MDAnalysis.analysis import align, rms | |
6 from MDAnalysis.analysis.base import AnalysisFromFunction | |
7 from MDAnalysis.coordinates.memory import MemoryReader | |
8 | |
9 import numpy as np | |
10 | |
11 | |
12 def calc_rmsd(str_files, traj_files, ref_str, str_format, traj_format, | |
13 ref_str_format, filepath_out, group, start, end, step, | |
14 fitting_atoms): | |
15 """ | |
16 the function will cycle through range 0 to no_t and load all files found. | |
17 | |
18 str_files: text file with filepaths for structures, one on each line | |
19 traj_files: text file with filepaths for trajectories, one on each line | |
20 ref_str: reference structure for fitting | |
21 filepath_in: directory where the files are located | |
22 filepath_out: pickle file where results (3D matrix) should be saved to | |
23 | |
24 group: atoms for which RMSD should be calculated; | |
25 use the MDAnalysis selection language | |
26 fitting_atoms: atoms used for str alignment prior to RMSD calculation; | |
27 use the MDAnalysis selection language | |
28 | |
29 start: first trajectory frame to calculate RMSD | |
30 end: last trajectory frame to calculate RMSD | |
31 step: how frequently frames are sampled between start and end; obviously, | |
32 the larger the step, the quicker the script finishes | |
33 """ | |
34 | |
35 # open list of files | |
36 with open(str_files) as f1, open(traj_files) as f2: | |
37 str_file_list = f1.read().strip().split('\n') | |
38 traj_file_list = f2.read().strip().split('\n') | |
39 | |
40 if sum(1 for line in f1) != sum(1 for line in f2): | |
41 raise IOError('Number of structure and trajectory files unequal.') | |
42 | |
43 no_t = len(traj_file_list) | |
44 | |
45 data = np.zeros((no_t, no_t, | |
46 int((end - start)/step + ((end - start) % step > 0)))) | |
47 | |
48 # load files | |
49 universes = {} | |
50 | |
51 for traj in range(no_t): | |
52 mobile = m.Universe(str_file_list[traj], traj_file_list[traj], | |
53 format=traj_format, topology_format=str_format) | |
54 ref = m.Universe(ref_str, topology_format=ref_str_format) | |
55 | |
56 mobile.trajectory[-1] # set mobile trajectory to last frame | |
57 ref.trajectory[0] # set reference trajectory to first frame | |
58 | |
59 # perform alignment | |
60 align.AlignTraj(mobile, ref, select=fitting_atoms, | |
61 in_memory=True).run() | |
62 | |
63 grp = mobile.select_atoms(group) | |
64 universes[traj] = m.core.universe.Merge(grp) # create Universe w grp | |
65 coordinates = AnalysisFromFunction(lambda ag: ag.positions.copy(), | |
66 grp).run().results # write to uv | |
67 universes[traj].load_new(coordinates, format=MemoryReader) | |
68 | |
69 print("All trajs loaded by MDAnalysis") | |
70 | |
71 # calculate differences | |
72 for traj1 in range(no_t): | |
73 print("Calculating differences for traj {}".format(traj1)) | |
74 for traj2 in range(traj1): | |
75 | |
76 u1 = universes[traj1] | |
77 u2 = universes[traj2] | |
78 | |
79 l1 = u1.select_atoms(group) | |
80 l2 = u2.select_atoms(group) | |
81 | |
82 rmsd = rms.RMSD(l1, l2) | |
83 | |
84 rmsd.run() | |
85 | |
86 data[traj1, traj2] = rmsd.rmsd[:, 2] | |
87 data[traj2, traj1] = rmsd.rmsd[:, 2] | |
88 | |
89 with open(filepath_out, 'w') as f: | |
90 json.dump(data.tolist(), f, indent=4, sort_keys=True) | |
91 | |
92 print("Done!") | |
93 return | |
94 | |
95 | |
96 def main(): | |
97 parser = argparse.ArgumentParser() | |
98 parser.add_argument('--trajs', required=True, | |
99 help='File containing trajectory filepaths.') | |
100 parser.add_argument("--strs", | |
101 help='File containing structure filepaths.') | |
102 parser.add_argument("--ref-str", | |
103 help='File containing reference structure.') | |
104 parser.add_argument('--traj-format', required=True, | |
105 help='Trajectory format.') | |
106 parser.add_argument("--str-format", help='Structure format.') | |
107 parser.add_argument("--ref-str-format", | |
108 help='Reference structure format.') | |
109 parser.add_argument('-o', '--outfile', | |
110 help="Path to the output JSON file") | |
111 parser.add_argument('--group', help="Atoms for which RMSD should be" | |
112 "calculated in MDAnalysis selection language") | |
113 parser.add_argument('--fitting', help="Fitting atoms for alignment" | |
114 "prior to RMSD calculation") | |
115 parser.add_argument('--start', type=int, | |
116 help="First trajectory frame to calculate RMSD") | |
117 parser.add_argument('--end', type=int, | |
118 help="Last trajectory frame to calculate RMSD") | |
119 parser.add_argument('--step', type=int, | |
120 help="Frame sampling frequency for RMSD calculation") | |
121 args = parser.parse_args() | |
122 | |
123 calc_rmsd(args.strs, args.trajs, args.ref_str, args.str_format, | |
124 args.traj_format, args.ref_str_format, args.outfile, | |
125 args.group, args.start, args.end, args.step, args.fitting) | |
126 | |
127 | |
128 if __name__ == "__main__": | |
129 main() |