Mercurial > repos > chemteam > mdanalysis_endtoend

comparison extract_rmsd.py @ 1:ce9dc91ff87f draft

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
"planemo upload for repository https://github.com/galaxycomputationalchemistry/galaxy-tools-compchem/ commit 45fe75a3a8ca80f799c85e194429c4c7f38bb5f6"
author chemteam
date Wed, 28 Oct 2020 21:37:06 +0000
parents 78aa3659fcd1
children 795a5996cdc8
comparison
equal deleted inserted replaced
0:78aa3659fcd1 1:ce9dc91ff87f
1 import argparse 1 import argparse
2 import json 2 import json
3 3
4 import MDAnalysis as m 4 import MDAnalysis as m
5 from MDAnalysis.analysis import align, rms 5 from MDAnalysis.analysis import rms
6 from MDAnalysis.analysis.base import AnalysisFromFunction
7 from MDAnalysis.coordinates.memory import MemoryReader
8 6
9 import numpy as np 7 import numpy as np
10 8
11 9
12 def calc_rmsd(str_files, traj_files, ref_str, str_format, traj_format, 10 def calc_rmsd(str_files, traj_files, str_format, traj_format, filepath_out,
13 ref_str_format, filepath_out, group, start, end, step, 11 group, start, end, step):
14 fitting_atoms):
15 """ 12 """
16 the function will cycle through range 0 to no_t and load all files found. 13 the function will cycle through range 0 to no_t and load all files found.
17 14
18 str_files: text file with filepaths for structures, one on each line 15 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 16 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 17 filepath_in: directory where the files are located
22 filepath_out: pickle file where results (3D matrix) should be saved to 18 filepath_out: pickle file where results (3D matrix) should be saved to
23 19
24 group: atoms for which RMSD should be calculated; 20 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 21 use the MDAnalysis selection language
28 22
29 start: first trajectory frame to calculate RMSD 23 start: first trajectory frame to calculate RMSD
30 end: last trajectory frame to calculate RMSD 24 end: last trajectory frame to calculate RMSD
31 step: how frequently frames are sampled between start and end; obviously, 25 step: how frequently frames are sampled between start and end; obviously,
47 41
48 # load files 42 # load files
49 universes = {} 43 universes = {}
50 44
51 for traj in range(no_t): 45 for traj in range(no_t):
52 mobile = m.Universe(str_file_list[traj], traj_file_list[traj], 46 # We no longer align here, users should do this themselves.
53 format=traj_format, topology_format=str_format) 47 universes[traj] = m.Universe(str_file_list[traj], traj_file_list[traj],
54 ref = m.Universe(ref_str, topology_format=ref_str_format) 48 format=traj_format,
55 49 topology_format=str_format)
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 50
69 print("All trajs loaded by MDAnalysis") 51 print("All trajs loaded by MDAnalysis")
70 52
71 # calculate differences 53 # calculate differences
72 for traj1 in range(no_t): 54 for traj1 in range(no_t):
73 print("Calculating differences for traj {}".format(traj1)) 55 print("Calculating differences for traj {}".format(traj1))
74 for traj2 in range(traj1): 56 for traj2 in range(traj1):
75 57 for frame in range(data.shape[2]):
76 u1 = universes[traj1] 58 universes[traj1].trajectory[frame]
77 u2 = universes[traj2] 59 universes[traj2].trajectory[frame]
78 60 A = universes[traj1].select_atoms(group).positions
79 l1 = u1.select_atoms(group) 61 B = universes[traj2].select_atoms(group).positions
80 l2 = u2.select_atoms(group) 62 r = rms.rmsd(A, B)
81 63 data[traj1, traj2, frame] = r
82 rmsd = rms.RMSD(l1, l2) 64 data[traj2, traj1, frame] = r
83
84 rmsd.run()
85
86 data[traj1, traj2] = rmsd.rmsd[:, 2]
87 data[traj2, traj1] = rmsd.rmsd[:, 2]
88 65
89 with open(filepath_out, 'w') as f: 66 with open(filepath_out, 'w') as f:
90 json.dump(data.tolist(), f, indent=4, sort_keys=True) 67 json.dump(data.tolist(), f, indent=4, sort_keys=True)
91 68
92 print("Done!") 69 print("Done!")
97 parser = argparse.ArgumentParser() 74 parser = argparse.ArgumentParser()
98 parser.add_argument('--trajs', required=True, 75 parser.add_argument('--trajs', required=True,
99 help='File containing trajectory filepaths.') 76 help='File containing trajectory filepaths.')
100 parser.add_argument("--strs", 77 parser.add_argument("--strs",
101 help='File containing structure filepaths.') 78 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, 79 parser.add_argument('--traj-format', required=True,
105 help='Trajectory format.') 80 help='Trajectory format.')
106 parser.add_argument("--str-format", help='Structure format.') 81 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', 82 parser.add_argument('-o', '--outfile',
110 help="Path to the output JSON file") 83 help="Path to the output JSON file")
111 parser.add_argument('--group', help="Atoms for which RMSD should be" 84 parser.add_argument('--group', help="Atoms for which RMSD should be"
112 "calculated in MDAnalysis selection language") 85 "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, 86 parser.add_argument('--start', type=int,
116 help="First trajectory frame to calculate RMSD") 87 help="First trajectory frame to calculate RMSD")
117 parser.add_argument('--end', type=int, 88 parser.add_argument('--end', type=int,
118 help="Last trajectory frame to calculate RMSD") 89 help="Last trajectory frame to calculate RMSD")
119 parser.add_argument('--step', type=int, 90 parser.add_argument('--step', type=int,
120 help="Frame sampling frequency for RMSD calculation") 91 help="Frame sampling frequency for RMSD calculation")
121 args = parser.parse_args() 92 args = parser.parse_args()
122 93
123 calc_rmsd(args.strs, args.trajs, args.ref_str, args.str_format, 94 calc_rmsd(args.strs, args.trajs, args.str_format,
124 args.traj_format, args.ref_str_format, args.outfile, 95 args.traj_format, args.outfile,
125 args.group, args.start, args.end, args.step, args.fitting) 96 args.group, args.start, args.end, args.step)
126 97
127 98
128 if __name__ == "__main__": 99 if __name__ == "__main__":
129 main() 100 main()