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