forked from mosdef-hub/gmso
-
Notifications
You must be signed in to change notification settings - Fork 0
/
top.py
633 lines (570 loc) · 24.4 KB
/
top.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
"""Write a GROMACS topology (.TOP) file."""
import datetime
import warnings
import unyt as u
from gmso.core.dihedral import Dihedral
from gmso.core.element import element_by_atom_type
from gmso.core.improper import Improper
from gmso.core.views import PotentialFilters
from gmso.exceptions import GMSOError
from gmso.external import to_networkx
from gmso.formats.formats_registry import saves_as
from gmso.lib.potential_templates import PotentialTemplateLibrary
from gmso.parameterization.molecule_utils import (
molecule_angles,
molecule_bonds,
molecule_dihedrals,
molecule_impropers,
)
from gmso.utils.compatibility import check_compatibility
from gmso.utils.connectivity import generate_pairs_lists
@saves_as(".top")
def write_top(top, filename, top_vars=None):
"""Write a gmso.core.Topology object to a GROMACS topology (.TOP) file.
Parameters
----------
top : gmso.Topology
A typed Topology Object
filename : str
Path of the output file
Notes
-----
See https://manual.gromacs.org/current/reference-manual/topologies/topology-file-formats.html for
a full description of the top file format. This method is a work in progress and do not currently
support the full GROMACS specs.
"""
pot_types = _validate_compatibility(top)
top_vars = _get_top_vars(top, top_vars)
# Sanity checks
msg = "System not fully typed"
for site in top.sites:
assert site.atom_type, msg
for connection in top.connections:
assert connection.connection_type, msg
with open(filename, "w") as out_file:
out_file.write(
"; File {} written by GMSO at {}\n\n".format(
top.name if top.name is not None else "",
str(datetime.datetime.now()),
)
)
out_file.write(
"[ defaults ]\n"
"; nbfunc\t"
"comb-rule\t"
"gen-pairs\t"
"fudgeLJ\t\t"
"fudgeQQ\n"
)
out_file.write(
"{0}\t\t"
"{1}\t\t"
"{2}\t\t"
"{3}\t\t"
"{4}\n\n".format(
top_vars["nbfunc"],
top_vars["comb-rule"],
top_vars["gen-pairs"],
top_vars["fudgeLJ"],
top_vars["fudgeQQ"],
)
)
out_file.write(
"[ atomtypes ]\n"
"; name\t"
"at.num\t\t"
"mass\t"
"charge\t\t"
"ptype\t"
"sigma\t"
"epsilon\n"
)
for atom_type in top.atom_types(PotentialFilters.UNIQUE_NAME_CLASS):
out_file.write(
"{0:12s}"
"{1:4s}"
"{2:12.5f}"
"{3:12.5f}\t"
"{4:4s}"
"{5:12.5f}"
"{6:12.5f}\n".format(
atom_type.name,
str(_lookup_atomic_number(atom_type)),
atom_type.mass.in_units(u.amu).value,
atom_type.charge.in_units(u.elementary_charge).value,
"A",
atom_type.parameters["sigma"].in_units(u.nanometer).value,
atom_type.parameters["epsilon"]
.in_units(u.Unit("kJ/mol"))
.value,
)
)
# Define unique molecule by name only
unique_molecules = _get_unique_molecules(top)
# Section headers
headers = {
"bonds": "\n[ bonds ]\n; ai\taj\tfunct\tb0\t\tkb\n",
"bond_restraints": "\n[ bonds ] ;Harmonic potential restraint\n"
"; ai\taj\tfunct\tb0\t\tkb\n",
"pairs": "\n[ pairs ]\n; ai\taj\tfunct\n",
"angles": "\n[ angles ]\n" "; ai\taj\tak\tfunct\tphi_0\t\tk0\n",
"angle_restraints": (
"\n[ angle_restraints ]\n"
"; ai\taj\tai\tak\tfunct\ttheta_eq\tk\tmultiplicity\n"
),
"dihedrals": {
"RyckaertBellemansTorsionPotential": "\n[ dihedrals ]\n"
"; ai\taj\tak\tal\tfunct\tc0\t\tc1\t\tc2\t\tc3\t\tc4\t\tc5\n",
"PeriodicTorsionPotential": "\n[ dihedrals ]\n"
"; ai\taj\tak\tal\tfunct\tphi\tk_phi\tmulitplicity\n",
},
"dihedral_restraints": "\n[ dihedral_restraints ]\n"
"#ifdef DIHRES\n"
"; ai\taj\tak\tal\tfunct\ttheta_eq\tdelta_theta\t\tkd\n",
}
for tag in unique_molecules:
"""Write out nrexcl for each unique molecule."""
out_file.write("\n[ moleculetype ]\n" "; name\tnrexcl\n")
# TODO: Lookup and join nrexcl from each molecule object
out_file.write("{0}\t" "{1}\n\n".format(tag, top_vars["nrexcl"]))
"""Write out atoms for each unique molecule."""
out_file.write(
"[ atoms ]\n"
"; nr\ttype\tresnr\tresidue\t\tatom\tcgnr\tcharge\tmass\n"
)
# Each unique molecule need to be reindexed (restarting from 0)
# The shifted_idx_map is needed to make sure all the atom index used in
# latter connection sections are acurate
shifted_idx_map = dict()
for idx, site in enumerate(unique_molecules[tag]["sites"]):
shifted_idx_map[top.get_index(site)] = idx
out_file.write(
"{0:8s}"
"{1:12s}"
"{2:8s}"
"{3:12s}"
"{4:8s}"
"{5:4s}"
"{6:12.5f}"
"{7:12.5f}\n".format(
str(idx + 1),
site.atom_type.name,
str(site.molecule.number if site.molecule else 1),
tag,
site.atom_type.tags["element"],
"1", # TODO: care about charge groups
site.charge.in_units(u.elementary_charge).value,
site.atom_type.mass.in_units(u.amu).value,
)
)
# Special treatment for water, may ned to consider a better way to tag rigid water
# Built using this https://github.com/gromacs/gromacs/blob/main/share/top/oplsaa.ff/spce.itp as reference
if "water" in tag.lower() and all(
"rigid" in site.label for site in unique_molecules[tag]["sites"]
):
sites_list = unique_molecules[tag]["sites"]
water_sites = {
"O": [
site
for site in sites_list
if site.element.symbol == "O"
],
"H": [
site
for site in sites_list
if site.element.symbol == "H"
],
}
ow_idx = shifted_idx_map[top.get_index(water_sites["O"][0])] + 1
doh = np.linalg.norm(
water_sites["O"][0].position.to(u.nm)
- water_sites["H"][0].position.to(u.nm)
)
dhh = np.linalg.norm(
water_sites["H"][0].position.to(u.nm)
- water_sites["H"][1].position.to(u.nm)
)
# Write settles
out_file.write(
"\n[ settles ] ;Water specific constraint algorithm\n"
"; OW_idx\tfunct\tdoh\tdhh\n"
)
out_file.write(
"{0:4s}{1:4s}{2:15.5f}{3:15.5f}\n".format(
str(ow_idx), "1", doh, dhh
)
)
# Write exclusion
out_file.write(
"\n[ exclusions ] ;Exclude all interactions between water's atoms\n"
"1\t2\t3\n"
"2\t1\t3\n"
"3\t1\t2\n"
)
# Skip the rest of the loop
continue
for conn_group in [
"pairs",
"bonds",
"bond_restraints",
"angles",
"angle_restraints",
"dihedrals",
"dihedral_restraints",
"impropers",
]:
if unique_molecules[tag][conn_group]:
if conn_group == "pairs":
out_file.write(headers[conn_group])
for conn in unique_molecules[tag][conn_group]:
out_file.write(
_write_pairs(top, conn, shifted_idx_map)
)
elif conn_group in ["dihedrals", "impropers"]:
proper_groups = {
"RyckaertBellemansTorsionPotential": list(),
"PeriodicTorsionPotential": list(),
}
for dihedral in unique_molecules[tag][conn_group]:
ptype = pot_types[dihedral.connection_type]
proper_groups[ptype].append(dihedral)
# Improper use same header as dihedral periodic header
if proper_groups["RyckaertBellemansTorsionPotential"]:
out_file.write(
headers["dihedrals"][
"RyckaertBellemansTorsionPotential"
]
)
for conn in proper_groups[
"RyckaertBellemansTorsionPotential"
]:
for line in _write_connection(
top,
conn,
pot_types[conn.connection_type],
shifted_idx_map,
):
out_file.write(line)
if proper_groups["PeriodicTorsionPotential"]:
out_file.write(
headers["dihedrals"]["PeriodicTorsionPotential"]
)
for conn in proper_groups[
"PeriodicTorsionPotential"
]:
for line in _write_connection(
top,
conn,
pot_types[conn.connection_type],
shifted_idx_map,
):
out_file.write(line)
elif "restraints" in conn_group:
out_file.write(headers[conn_group])
for conn in unique_molecules[tag][conn_group]:
out_file.write(
_write_restraint(
top,
conn,
conn_group,
shifted_idx_map,
)
)
if conn_group == "dihedral_restraints":
warnings.warn(
"The diehdral_restraints writer is designed to work with"
"`define = DDIHRES` clause in the GROMACS input file (.mdp)"
)
out_file.write("#endif DIHRES\n")
elif unique_molecules[tag][conn_group]:
out_file.write(headers[conn_group])
for conn in unique_molecules[tag][conn_group]:
out_file.write(
_write_connection(
top,
conn,
pot_types[conn.connection_type],
shifted_idx_map,
)
)
out_file.write("\n[ system ]\n" "; name\n" "{0}\n\n".format(top.name))
out_file.write("[ molecules ]\n" "; molecule\tnmols\n")
for tag in unique_molecules:
out_file.write(
"{0}\t{1}\n".format(tag, len(unique_molecules[tag]["subtags"]))
)
def _accepted_potentials():
"""List of accepted potentials that GROMACS can support."""
templates = PotentialTemplateLibrary()
lennard_jones_potential = templates["LennardJonesPotential"]
harmonic_bond_potential = templates["HarmonicBondPotential"]
harmonic_angle_potential = templates["HarmonicAnglePotential"]
periodic_torsion_potential = templates["PeriodicTorsionPotential"]
rb_torsion_potential = templates["RyckaertBellemansTorsionPotential"]
accepted_potentials = (
lennard_jones_potential,
harmonic_bond_potential,
harmonic_angle_potential,
periodic_torsion_potential,
rb_torsion_potential,
)
return accepted_potentials
def _validate_compatibility(top):
"""Check compatability of topology object with GROMACS TOP format."""
pot_types = check_compatibility(top, _accepted_potentials())
return pot_types
def _get_top_vars(top, top_vars):
"""Generate a dictionary of values for the defaults directive."""
combining_rule_to_gmx = {"lorentz": 2, "geometric": 3}
default_top_vars = dict()
default_top_vars["nbfunc"] = 1 # modify this to check for lj or buckingham
default_top_vars["comb-rule"] = combining_rule_to_gmx[top.combining_rule]
default_top_vars["gen-pairs"] = "yes"
default_top_vars["fudgeLJ"] = top.scaling_factors[0][2]
default_top_vars["fudgeQQ"] = top.scaling_factors[1][2]
default_top_vars["nrexcl"] = 3
if isinstance(top_vars, dict):
default_top_vars.update(top_vars)
return default_top_vars
def _get_unique_molecules(top):
unique_molecules = {
tag: {
"subtags": list(),
}
for tag in top.unique_site_labels("molecule", name_only=True)
}
for molecule in top.unique_site_labels("molecule", name_only=False):
unique_molecules[molecule.name]["subtags"].append(molecule)
if len(unique_molecules) == 0:
unique_molecules[top.name] = dict()
unique_molecules[top.name]["subtags"] = [top.name]
unique_molecules[top.name]["sites"] = list(top.sites)
unique_molecules[top.name]["pairs"] = generate_pairs_lists(
top, refer_from_scaling_factor=True
)["pairs14"]
unique_molecules[top.name]["bonds"] = list(top.bonds)
unique_molecules[top.name]["bond_restraints"] = list(
bond for bond in top.bonds if bond.restraint
)
unique_molecules[top.name]["angles"] = list(top.angles)
unique_molecules[top.name]["angle_restraints"] = list(
angle for angle in top.angles if angle.restraint
)
unique_molecules[top.name]["dihedrals"] = list(top.dihedrals)
unique_molecules[top.name]["dihedral_restraints"] = list(
dihedral for dihedral in top.dihedrals if dihedral.restraint
)
unique_molecules[molecule.name]["impropers"] = list(top.impropers)
else:
for tag in unique_molecules:
molecule = unique_molecules[tag]["subtags"][0]
unique_molecules[tag]["sites"] = list(
top.iter_sites(key="molecule", value=molecule)
)
unique_molecules[tag]["pairs"] = generate_pairs_lists(
top, molecule
)["pairs14"]
unique_molecules[tag]["bonds"] = list(molecule_bonds(top, molecule))
unique_molecules[tag]["bond_restraints"] = list(
bond for bond in molecule_bonds(top, molecule) if bond.restraint
)
unique_molecules[tag]["angles"] = list(
molecule_angles(top, molecule)
)
unique_molecules[tag]["angle_restraints"] = list(
angle
for angle in molecule_angles(top, molecule)
if angle.restraint
)
unique_molecules[tag]["dihedrals"] = list(
molecule_dihedrals(top, molecule)
)
unique_molecules[tag]["dihedral_restraints"] = list(
dihedral
for dihedral in molecule_dihedrals(top, molecule)
if dihedral.restraint
)
unique_molecules[tag]["impropers"] = list(
molecule_impropers(top, molecule)
)
return unique_molecules
def _lookup_atomic_number(atom_type):
"""Look up an atomic_number based on atom type information, 0 if non-element type."""
try:
element = element_by_atom_type(atom_type)
return element.atomic_number
except GMSOError:
return 0
def _lookup_element_symbol(atom_type):
"""Look up an atomic_number based on atom type information, 0 if non-element type."""
try:
element = element_by_atom_type(atom_type)
return element.symbol
except GMSOError:
return "X"
def _write_pairs(top, pair, shifted_idx_map):
"""Workder function to write out pairs information."""
pair_idx = [
shifted_idx_map[top.get_index(pair[0])] + 1,
shifted_idx_map[top.get_index(pair[1])] + 1,
]
line = "{0:8s}{1:8s}{2:4s}\n".format(
str(pair_idx[0]),
str(pair_idx[1]),
"1",
)
return line
def _write_connection(top, connection, potential_name, shifted_idx_map):
"""Worker function to write various connection information."""
worker_functions = {
"HarmonicBondPotential": _harmonic_bond_potential_writer,
"HarmonicAnglePotential": _harmonic_angle_potential_writer,
"RyckaertBellemansTorsionPotential": _ryckaert_bellemans_torsion_writer,
"PeriodicTorsionPotential": _periodic_torsion_writer,
}
return worker_functions[potential_name](top, connection, shifted_idx_map)
def _harmonic_bond_potential_writer(top, bond, shifted_idx_map):
"""Write harmonic bond information."""
line = "{0:8s}{1:8s}{2:4s}{3:15.5f}{4:15.5f}\n".format(
str(shifted_idx_map[top.get_index(bond.connection_members[0])] + 1),
str(shifted_idx_map[top.get_index(bond.connection_members[1])] + 1),
"1",
bond.connection_type.parameters["r_eq"].in_units(u.nm).value,
bond.connection_type.parameters["k"]
.in_units(u.Unit("kJ / (mol*nm**2)"))
.value,
)
return line
def _harmonic_angle_potential_writer(top, angle, shifted_idx_map):
"""Write harmonic angle information."""
line = "{0:8s}{1:8s}{2:8s}{3:4s}{4:15.5f}{5:15.5f}\n".format(
str(shifted_idx_map[top.get_index(angle.connection_members[0])] + 1),
str(shifted_idx_map[top.get_index(angle.connection_members[1])] + 1),
str(shifted_idx_map[top.get_index(angle.connection_members[2])] + 1),
"1",
angle.connection_type.parameters["theta_eq"].in_units(u.degree).value,
angle.connection_type.parameters["k"]
.in_units(u.Unit("kJ/(mol*rad**2)"))
.value,
)
return line
def _ryckaert_bellemans_torsion_writer(top, dihedral, shifted_idx_map):
"""Write Ryckaert-Bellemans Torsion information."""
line = "{0:8s}{1:8s}{2:8s}{3:8s}{4:4s}{5:15.5f}{6:15.5f}{7:15.5f}{8:15.5f}{9:15.5f}{10:15.5f}\n".format(
str(shifted_idx_map[top.get_index(dihedral.connection_members[0])] + 1),
str(shifted_idx_map[top.get_index(dihedral.connection_members[1])] + 1),
str(shifted_idx_map[top.get_index(dihedral.connection_members[2])] + 1),
str(shifted_idx_map[top.get_index(dihedral.connection_members[3])] + 1),
"3",
dihedral.connection_type.parameters["c0"]
.in_units(u.Unit("kJ/mol"))
.value,
dihedral.connection_type.parameters["c1"]
.in_units(u.Unit("kJ/mol"))
.value,
dihedral.connection_type.parameters["c2"]
.in_units(u.Unit("kJ/mol"))
.value,
dihedral.connection_type.parameters["c3"]
.in_units(u.Unit("kJ/mol"))
.value,
dihedral.connection_type.parameters["c4"]
.in_units(u.Unit("kJ/mol"))
.value,
dihedral.connection_type.parameters["c5"]
.in_units(u.Unit("kJ/mol"))
.value,
)
return line
def _periodic_torsion_writer(top, dihedral, shifted_idx_map):
"""Write periodic torsion information."""
if isinstance(dihedral, Dihedral):
if dihedral.connection_type.parameters["phi_eq"].size == 1:
# Normal dihedral
layers, funct = 1, "1"
for key, val in dihedral.connection_type.parameters.items():
dihedral.connection_type.parameters[key] = val.reshape(layers)
else:
# Layered/Multiple dihedral
layers, funct = (
dihedral.connection_type.parameters["phi_eq"].size,
"9",
)
elif isinstance(dihedral, Improper):
layers, funct = 1, "4"
else:
raise TypeError(f"Type {type(dihedral)} not supported.")
lines = list()
for i in range(layers):
line = "{0:8s}{1:8s}{2:8s}{3:8s}{4:4s}{5:15.5f}{6:15.5f}{7:4}\n".format(
str(
shifted_idx_map[top.get_index(dihedral.connection_members[0])]
+ 1
),
str(
shifted_idx_map[top.get_index(dihedral.connection_members[1])]
+ 1
),
str(
shifted_idx_map[top.get_index(dihedral.connection_members[2])]
+ 1
),
str(
shifted_idx_map[top.get_index(dihedral.connection_members[3])]
+ 1
),
funct,
dihedral.connection_type.parameters["phi_eq"][i]
.in_units(u.degree)
.value,
dihedral.connection_type.parameters["k"][i]
.in_units(u.Unit("kJ/(mol)"))
.value,
dihedral.connection_type.parameters["n"][i].value,
)
lines.append(line)
return lines
def _write_restraint(top, connection, type, shifted_idx_map):
"""Worker function to write various connection restraint information."""
worker_functions = {
"bond_restraints": _bond_restraint_writer,
"angle_restraints": _angle_restraint_writer,
"dihedral_restraints": _dihedral_restraint_writer,
}
return worker_functions[type](top, connection, shifted_idx_map)
def _bond_restraint_writer(top, bond, shifted_idx_map):
"""Write bond restraint information."""
line = "{0:8s}{1:8s}{2:4s}{3:15.5f}{4:15.5f}\n".format(
str(shifted_idx_map[top.get_index(bond.connection_members[1])] + 1),
str(shifted_idx_map[top.get_index(bond.connection_members[0])] + 1),
"6",
bond.restraint["r_eq"].in_units(u.nm).value,
bond.restraint["k"].in_units(u.Unit("kJ/(mol * nm**2)")).value,
)
return line
def _angle_restraint_writer(top, angle, shifted_idx_map):
"""Write angle restraint information."""
line = "{0:8s}{1:8s}{2:8s}{3:8s}{4:4s}{5:15.5f}{6:15.5f}{7:4}\n".format(
str(shifted_idx_map[top.get_index(angle.connection_members[1])] + 1),
str(shifted_idx_map[top.get_index(angle.connection_members[0])] + 1),
str(shifted_idx_map[top.get_index(angle.connection_members[1])] + 1),
str(shifted_idx_map[top.get_index(angle.connection_members[2])] + 1),
"1",
angle.restraint["theta_eq"].in_units(u.degree).value,
angle.restraint["k"].in_units(u.Unit("kJ/mol")).value,
angle.restraint["n"],
)
return line
def _dihedral_restraint_writer(top, dihedral, shifted_idx_map):
"""Write dihedral restraint information."""
line = "{0:8s}{1:8s}{2:8s}{3:8s}{4:4s}{5:15.5f}{6:15.5f}{7:15.5f}\n".format(
str(shifted_idx_map[top.get_index(dihedral.connection_members[0])] + 1),
str(shifted_idx_map[top.get_index(dihedral.connection_members[1])] + 1),
str(shifted_idx_map[top.get_index(dihedral.connection_members[2])] + 1),
str(shifted_idx_map[top.get_index(dihedral.connection_members[3])] + 1),
"1",
dihedral.restraint["phi_eq"].in_units(u.degree).value,
dihedral.restraint["delta_phi"].in_units(u.degree).value,
dihedral.restraint["k"].in_units(u.Unit("kJ/(mol * rad**2)")).value,
)
return line