Coverage for biobb_structure_utils/gro_lib/gro.py: 72%
210 statements
« prev ^ index » next coverage.py v7.5.3, created at 2024-06-14 19:03 +0000
1import sys
2import re
5class Gro:
6 """
7 the central class in gropy
8 """
9 # -- constructor(s) --
10 def __init__(self,
11 system_name=None, num_of_atoms=None,
12 residue_id=None, residue_name=None,
13 atom_name=None, atom_id=None,
14 x=None, y=None, z=None,
15 v_x=None, v_y=None, v_z=None,
16 box=None):
17 """
18 wrap the contents in a GROMACS gro file in a class
19 """
20 self.system_name = system_name or 'This is a Gro!'
21 self.num_of_atoms = num_of_atoms or 0
22 self.residue_id = residue_id or []
23 self.residue_name = residue_name or []
24 self.atom_name = atom_name or []
25 self.atom_id = atom_id or []
26 self.x = x or []
27 self.y = y or []
28 self.z = z or []
29 self.v_x = v_x or []
30 self.v_y = v_y or []
31 self.v_z = v_z or []
32 self.box = box or [0.0, 0.0, 0.0]
34 # -- deconstructor --
35 # not mandatory in python
37 # -- file i/o --
38 def read_gro_file(self, file_name):
39 """
40 read a gro file and store information in a Gro object
41 """
42 with open(file_name, 'r') as file_id:
43 for i_line, line in enumerate(file_id):
44 line = line.replace('\n', '')
46 if i_line == 0:
47 self.system_name = line
48 continue
50 if i_line == 1:
51 self.num_of_atoms = int(line)
52 final_line_of_atoms = self.num_of_atoms + 1
53 continue
55 if i_line <= final_line_of_atoms:
56 # store atom information
57 self.residue_id.append(int(line[0:5]))
58 self.residue_name.append(line[5:10].strip()) # remove leading spaces
59 self.atom_name.append(line[10:15].strip()) # remove leading spaces
60 self.atom_id.append(int(line[15:20]))
61 self.x.append(float(line[20:28]))
62 self.y.append(float(line[28:36]))
63 self.z.append(float(line[36:44]))
64 if len(line) > 44:
65 self.v_x.append(float(line[44:52]))
66 self.v_y.append(float(line[52:60]))
67 self.v_z.append(float(line[60:68]))
68 else:
69 self.v_x.append(0.0)
70 self.v_y.append(0.0)
71 self.v_z.append(0.0)
72 else:
73 self.box = line.split()
74 self.box = [float(box_size) for box_size in self.box]
76 def write_gro_file(self, file_name):
77 """
78 write a gro file based on a Gro object
79 """
80 with open(file_name, 'w') as file_id:
81 file_id.write("%s\n" % self.system_name)
82 file_id.write(" %d\n" % self.num_of_atoms)
83 if self.v_x != []:
84 for i in range(self.num_of_atoms):
85 file_id.write("%5d%-5s%5s%5d%8.3f%8.3f%8.3f%8.4f%8.4f%8.4f\n" % (self.residue_id[i], self.residue_name[i],
86 self.atom_name[i], self.atom_id[i], self.x[i],
87 self.y[i], self.z[i], self.v_x[i], self.v_y[i],
88 self.v_z[i]))
89 else:
90 for i in range(self.num_of_atoms):
91 file_id.write("%5d%-5s%5s%5d%8.3f%8.3f%8.3f%8.4f%8.4f%8.4f\n" % (self.residue_id[i], self.residue_name[i],
92 self.atom_name[i], self.atom_id[i], self.x[i],
93 self.y[i], self.z[i], 0.0, 0.0, 0.0))
95 # Writting box coordinates
96 list(map(lambda box_coord: file_id.write("%10.5f" % box_coord), self.box))
97 # for box_coord in self.box:
98 # file_id.write("%10.5f" % box_coord)
99 file_id.write("\n")
101 # -- preservative operations --
102 def rename_atoms(self, old_atom_names, new_atom_names):
103 """
104 rename atoms from the old_atom_names to the new_atom_names
105 """
106 assert (len(old_atom_names) == len(new_atom_names)), "old_atom_names doesn't have the same length as new_atom_names"
107 for i_atom in range(self.num_of_atoms):
108 for i_name in range(len(old_atom_names)):
109 if self.atom_name[i_atom] == old_atom_names[i_name]:
110 self.atom_name[i_atom] = new_atom_names[i_name]
111 break
113 # TODO: may add flexibility to rename atoms with specific residue_names
114 def rename_residues(self, old_residue_names, new_residue_names):
115 """
116 rename residues with an old_residue_name to a new_residue_name
117 """
118 assert (len(old_residue_names) == len(new_residue_names)), "old_residue_names doesn't have the same length as new_residue_names"
119 for i_atom in range(self.num_of_atoms):
120 for i_name in range(len(old_residue_names)):
121 if self.residue_name[i_atom] == old_residue_names[i_name]:
122 self.residue_name[i_atom] = new_residue_names[i_name]
123 break
125 def renumber_atoms(self, renumber_residues=True, renumber_residues_per_chain=True):
126 """
127 renumber residue_id and atom_id starting from 1; the original composition of each resdiue is maintained
128 """
129 atom_mapping = {}
130 residue_mapping = {}
131 residue_count = 0
132 last_residue_id = sys.maxsize * -1
133 chain_number = 1
134 residue_mapping[str(chain_number)] = {}
135 new_chain = False
136 for i_atom in range(self.num_of_atoms):
137 gro_atom_number = str(self.atom_id[i_atom])
138 new_atom_number = i_atom + 1
139 atom_mapping[str(gro_atom_number)] = str(new_atom_number)
140 self.atom_id[i_atom] = new_atom_number
142 # As in the GRO files there is no chain information, the residue number is used as heuristics.
143 if self.residue_id[i_atom] < last_residue_id:
144 new_chain = True
145 chain_number += 1
146 residue_mapping[str(chain_number)] = {}
148 # New residue
149 if self.residue_id[i_atom] != last_residue_id:
150 residue_count += 1
151 # New chain
152 if renumber_residues_per_chain and new_chain:
153 residue_count = 1
154 new_chain = False
156 if not renumber_residues:
157 residue_count = str(self.residue_id[i_atom])
159 residue_mapping[str(chain_number)][str(self.residue_id[i_atom])] = str(residue_count)
160 last_residue_id = self.residue_id[i_atom]
161 self.residue_id[i_atom] = residue_count
163 return residue_mapping, atom_mapping
165 def replace_atom_entry(self, i_atom, another_gro_object, j_atom):
166 """
167 replace the i-th atom of the current gro object with the j-th atom of another gro object
168 """
169 self.residue_id[i_atom] = another_gro_object.residue_id[j_atom]
170 self.residue_name[i_atom] = another_gro_object.residue_name[j_atom]
171 self.atom_name[i_atom] = another_gro_object.atom_name[j_atom]
172 self.atom_id[i_atom] = another_gro_object.atom_id[j_atom]
173 self.x[i_atom] = another_gro_object.x[j_atom]
174 self.y[i_atom] = another_gro_object.y[j_atom]
175 self.z[i_atom] = another_gro_object.z[j_atom]
176 self.v_x[i_atom] = another_gro_object.v_x[j_atom]
177 self.v_y[i_atom] = another_gro_object.v_y[j_atom]
178 self.v_z[i_atom] = another_gro_object.v_z[j_atom]
180 def sort_residues(self, residue_name_list):
181 """
182 sort residues in the provided order, attaching other unspecified residues to the end
183 """
184 sorted_gro = Gro()
185 # copy provided to another
186 for residue_name in residue_name_list:
187 sorted_gro.copy_residues(self, [residue_name])
188 # remove provided
189 self.remove_residues(residue_name_list)
190 # copy unprovided to another
191 # for i_atom in range(self.num_of_atoms):
192 for i_atom in range(self.num_of_atoms):
193 sorted_gro.copy_atom_entry(self, i_atom)
194 # delete unprovided
195 # for i_atom in range(self.num_of_atoms):
196 for i_atom in range(self.num_of_atoms):
197 self.remove_atom_entry(0)
198 # copy all back from another
199 # for i_atom in range(sorted_gro.num_of_atoms):
200 for i_atom in range(sorted_gro.num_of_atoms):
201 self.copy_atom_entry(sorted_gro, i_atom)
203 def sort_residues2(self, residue_name_list):
204 """
205 sort residues in the provided order, attaching other unspecified residues to the end
206 """
207 ini_gro = Gro()
208 sorted_gro = Gro()
209 # copy provided to sorted slice
210 for residue_name in residue_name_list:
211 sorted_gro.copy_residues(self, [residue_name])
212 # remove provided
213 self.remove_residues(residue_name_list)
215 # copy unprovided to initial slice
216 for i_atom in range(self.num_of_atoms):
217 ini_gro.copy_atom_entry(self, i_atom)
218 # delete unprovided
219 for i_atom in range(self.num_of_atoms):
220 self.remove_atom_entry(0)
222 # copy all back from another
223 for i_atom in range(ini_gro.num_of_atoms):
224 self.copy_atom_entry(ini_gro, i_atom)
225 for i_atom in range(sorted_gro.num_of_atoms):
226 self.copy_atom_entry(sorted_gro, i_atom)
228 # -- additive operations --
229 def copy_atom_entry(self, another_gro_object, i_atom):
230 """
231 copy the i-th atom entry from another gro object and append to the end of current gro object
232 """
233 self.num_of_atoms += 1
234 self.residue_id.append(another_gro_object.residue_id[i_atom])
235 self.residue_name.append(another_gro_object.residue_name[i_atom])
236 self.atom_name.append(another_gro_object.atom_name[i_atom])
237 self.atom_id.append(another_gro_object.atom_id[i_atom])
238 self.x.append(another_gro_object.x[i_atom])
239 self.y.append(another_gro_object.y[i_atom])
240 self.z.append(another_gro_object.z[i_atom])
241 self.v_x.append(another_gro_object.v_x[i_atom])
242 self.v_y.append(another_gro_object.v_y[i_atom])
243 self.v_z.append(another_gro_object.v_z[i_atom])
245 def copy_residue_entry(self, another_gro_object, residue_id, residue_name):
246 """
247 copy atoms of the specified residue from another gro object and append to the end of current gro object
248 """
249 for i_atom in range(another_gro_object.num_of_atoms):
250 if another_gro_object.residue_id[i_atom] == residue_id and another_gro_object.residue_name[i_atom] == residue_name:
251 self.copy_atom_entry(another_gro_object, i_atom)
253 def copy_atoms(self, another_gro_object, atom_name_list):
254 """
255 copy atoms with the provided atom names from another gro object and append to the end of current gro object
256 """
257 for i_atom in range(another_gro_object.num_of_atoms):
258 for atom_name in atom_name_list:
259 if another_gro_object.atom_name[i_atom] == atom_name:
260 self.copy_atom_entry(another_gro_object, i_atom)
261 break
263 def copy_residues(self, another_gro_object, residue_name_list):
264 """
265 copy atoms with the provided residue names from another gro object and append to the end of current gro object
266 """
267 for i_atom in range(another_gro_object.num_of_atoms):
268 for residue_name in residue_name_list:
269 if another_gro_object.residue_name[i_atom] == residue_name:
270 self.copy_atom_entry(another_gro_object, i_atom)
271 break
273 # TODO: may add to copy atoms with the providied atom names and residue names
275 # Python doesn't allow the overloading of assignment operator "=";
276 # In python, copying an object is often achived by utilizing the copy and deepcopy functions in the copy module.
278 # -- subtractive operations --
279 def remove_atom_entry(self, i_atom):
280 """
281 remove the i-th atom entry from current gro object
282 """
283 self.num_of_atoms -= 1
284 del self.residue_id[i_atom]
285 del self.residue_name[i_atom]
286 del self.atom_name[i_atom]
287 del self.atom_id[i_atom]
288 del self.x[i_atom]
289 del self.y[i_atom]
290 del self.z[i_atom]
291 del self.v_x[i_atom]
292 del self.v_y[i_atom]
293 del self.v_z[i_atom]
295 def remove_residue_entry(self, residue_id, residue_name):
296 """
297 remove atoms of the specified residue
298 """
299 atom_indice_to_be_removed = []
300 for i_atom in range(self.num_of_atoms):
301 if self.residue_id[i_atom] == residue_id and self.residue_name[i_atom] == residue_name:
302 atom_indice_to_be_removed.append(i_atom) # save indice first; direct removal would shrink the atom list
304 num_of_atoms_to_be_removed = len(atom_indice_to_be_removed)
305 for i_atom in range(num_of_atoms_to_be_removed):
306 self.remove_atom_entry(atom_indice_to_be_removed[i_atom] - i_atom) # shift atom indice to match the shrinkage of atom list
308 def remove_atoms(self, atom_name_list):
309 """
310 remove atoms with the provided atom names
311 """
312 atom_indice_to_be_removed = []
313 for i_atom in range(self.num_of_atoms):
314 for atom_name in atom_name_list:
315 if self.atom_name[i_atom] == atom_name:
316 atom_indice_to_be_removed.append(i_atom)
317 break
318 num_of_atoms_to_be_removed = len(atom_indice_to_be_removed)
319 for i_atom in range(num_of_atoms_to_be_removed):
320 self.remove_atom_entry(atom_indice_to_be_removed[i_atom] - i_atom) # shift atom indice to match the shrinkage of atom list
322 def select_atoms(self, regular_expression_pattern):
323 atom_indice_to_be_removed = []
324 for i_atom in range(self.num_of_atoms):
325 if not re.search(regular_expression_pattern, self.atom_name[i_atom]):
326 atom_indice_to_be_removed.append(i_atom)
328 num_of_atoms_to_be_removed = len(atom_indice_to_be_removed)
329 for i_atom in range(num_of_atoms_to_be_removed):
330 self.remove_atom_entry(atom_indice_to_be_removed[i_atom] - i_atom)
332 def remove_residues(self, residue_name_list):
333 """
334 remove atoms with the provided residue names
335 """
336 atom_indice_to_be_removed = []
337 for i_atom in range(self.num_of_atoms):
338 for residue_name in residue_name_list:
339 if self.residue_name[i_atom] == residue_name:
340 atom_indice_to_be_removed.append(i_atom)
341 break
342 num_of_atoms_to_be_removed = len(atom_indice_to_be_removed)
343 for i_atom in range(num_of_atoms_to_be_removed):
344 self.remove_atom_entry(atom_indice_to_be_removed[i_atom] - i_atom) # shift atom indice to match the shrinkage of atom list
346 # TODO: may add to copy atoms with the providied atom names and residue names