forked from leeping/crank-legacy
-
Notifications
You must be signed in to change notification settings - Fork 3
/
molecule.py
4073 lines (3751 loc) · 181 KB
/
molecule.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
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
#======================================================================#
#| |#
#| Chemical file format conversion module |#
#| |#
#| Lee-Ping Wang ([email protected]) |#
#| Last updated July 5, 2016 |#
#| |#
#| This is free software released under version 2 of the GNU GPL, |#
#| please use or redistribute as you see fit under the terms of |#
#| this license. (http://www.gnu.org/licenses/gpl-2.0.html) |#
#| |#
#| This program is distributed in the hope that it will be useful, |#
#| but without any warranty; without even the implied warranty of |#
#| merchantability or fitness for a particular purpose. See the |#
#| GNU General Public License for more details. |#
#| |#
#| Feedback and suggestions are encouraged. |#
#| |#
#| What this is for: |#
#| Converting a molecule between file formats |#
#| Loading and processing of trajectories |#
#| (list of geometries for the same set of atoms) |#
#| Concatenating or slicing trajectories |#
#| Combining molecule metadata (charge, Q-Chem rem variables) |#
#| |#
#| Supported file formats: |#
#| See the __init__ method in the Molecule class. |#
#| |#
#| Note to self / developers: |#
#| Please make this file as standalone as possible |#
#| (i.e. don't introduce dependencies). If we load an external |#
#| library to parse a file, do so with 'try / except' so that |#
#| the module is still usable even if certain parts are missing. |#
#| It's better to be like a Millennium Falcon. :P |#
#| |#
#| Please make sure this file is up-to-date in |#
#| both the 'nanoreactor' and 'forcebalance' modules |#
#| |#
#| At present, when I perform operations like adding two objects, |#
#| the sum is created from deep copies of data members in the |#
#| originals. This is because copying by reference is confusing; |#
#| suppose if I do B += A and modify something in B; it should not |#
#| change in A. |#
#| |#
#| A consequence of this is that data members should not be too |#
#| complicated; they should be things like lists or dicts, and NOT |#
#| contain references to themselves. |#
#| |#
#| To-do list: Handling of comments is still not very good. |#
#| Comments from previous files should be 'passed on' better. |#
#| |#
#| Contents of this file: |#
#| 0) Names of data variables |#
#| 1) Imports |#
#| 2) Subroutines |#
#| 3) Molecule class |#
#| a) Class customizations (add, getitem) |#
#| b) Instantiation |#
#| c) Core functionality (read, write) |#
#| d) Reading functions |#
#| e) Writing functions |#
#| f) Extra stuff |#
#| 4) "main" function (if executed) |#
#| |#
#| Required: Python 2.7, Numpy 1.6 |#
#| Optional: Mol2, PDB, DCD readers |#
#| (can be found in ForceBalance) |#
#| NetworkX package (for topologies) |#
#| |#
#| Thanks: Todd Dolinsky, Yong Huang, |#
#| Kyle Beauchamp (PDB) |#
#| John Stone (DCD Plugin) |#
#| Pierre Tuffery (Mol2 Plugin) |#
#| #python IRC chat on FreeNode |#
#| |#
#| Contributors: Leah Isseroff Bendavid |#
#| Yudong Qiu |#
#| |#
#| Instructions: |#
#| |#
#| To import: |#
#| from molecule import Molecule |#
#| To create a Molecule object: |#
#| MyMol = Molecule(fnm) |#
#| To convert to a new file format: |#
#| MyMol.write('newfnm.format') |#
#| To concatenate geometries: |#
#| MyMol += MyMolB |#
#| |#
#======================================================================#
#=========================================#
#| DECLARE VARIABLE NAMES HERE |#
#| |#
#| Any member variable in the Molecule |#
#| class must be declared here otherwise |#
#| the Molecule class won't recognize it |#
#=========================================#
#| Data attributes in FrameVariableNames |#
#| must be a list along the frame axis, |#
#| and they must have the same length. |#
#=========================================#
# xyzs = List of arrays of atomic xyz coordinates
# comms = List of comment strings
# boxes = List of 3-element or 9-element arrays for periodic boxes
# qm_grads = List of arrays of gradients (i.e. negative of the atomistic forces) from QM calculations
# qm_espxyzs = List of arrays of xyz coordinates for ESP evaluation
# qm_espvals = List of arrays of ESP values
# qm_zpe = Zero point energy, kcal/mol (from a qchem freq calculation)
# qm_entropy = Entropy contribution at STP, cal/mol.K (from a qchem freq calculation)
# qm_enthalpy= Enthalpic contribution at STP, excluding electronic energy and ZPE, kcal/mol (from a qchem freq calculation)
FrameVariableNames = set(['xyzs', 'comms', 'boxes', 'qm_hessians', 'qm_grads', 'qm_energies', 'qm_interaction',
'qm_espxyzs', 'qm_espvals', 'qm_extchgs', 'qm_mulliken_charges', 'qm_mulliken_spins',
'qm_zpe', 'qm_entropy', 'qm_enthalpy'])
#=========================================#
#| Data attributes in AtomVariableNames |#
#| must be a list along the atom axis, |#
#| and they must have the same length. |#
#=========================================#
# elem = List of elements
# partial_charge = List of atomic partial charges
# atomname = List of atom names (can come from MM coordinate file)
# atomtype = List of atom types (can come from MM force field)
# tinkersuf = String that comes after the XYZ coordinates in TINKER .xyz or .arc files
# resid = Residue IDs (can come from MM coordinate file)
# resname = Residue names
# terminal = List of true/false denoting whether this atom is followed by a terminal group.
AtomVariableNames = set(['elem', 'partial_charge', 'atomname', 'atomtype', 'tinkersuf', 'resid', 'resname', 'qcsuf', 'qm_ghost', 'chain', 'altloc', 'icode', 'terminal'])
#=========================================#
#| This can be any data attribute we |#
#| want but it's usually some property |#
#| of the molecule not along the frame |#
#| atom axis. |#
#=========================================#
# bonds = A list of 2-tuples representing bonds. Carefully pruned when atom subselection is done.
# fnm = The file name that the class was built from
# qcrems = The Q-Chem 'rem' variables stored as a list of OrderedDicts
# qctemplate = The Q-Chem template file, not including the coordinates or rem variables
# charge = The net charge of the molecule
# mult = The spin multiplicity of the molecule
MetaVariableNames = set(['fnm', 'ftype', 'qcrems', 'qctemplate', 'psitemplate', 'qcerr', 'charge', 'mult', 'bonds', 'topology', 'molecules'])
# Variable names relevant to quantum calculations explicitly
QuantumVariableNames = set(['qcrems', 'qctemplate', 'psitemplate', 'charge', 'mult', 'qcsuf', 'qm_ghost', 'qm_bondorder'])
# Superset of all variable names.
AllVariableNames = QuantumVariableNames | AtomVariableNames | MetaVariableNames | FrameVariableNames
# OrderedDict requires Python 2.7 or higher
import os, sys, re, copy
import numpy as np
from numpy import sin, cos, arcsin, arccos
import imp
import itertools
from collections import OrderedDict, namedtuple, Counter
from ctypes import *
from warnings import warn
#================================#
# Set up the logger #
#================================#
try:
from output import *
except:
from logging import *
class RawStreamHandler(StreamHandler):
"""Exactly like output.StreamHandler except it does no extra formatting
before sending logging messages to the stream. This is more compatible with
how output has been displayed in ForceBalance. Default stream has also been
changed from stderr to stdout"""
def __init__(self, stream = sys.stdout):
super(RawStreamHandler, self).__init__(stream)
def emit(self, record):
message = record.getMessage()
self.stream.write(message)
self.flush()
logger=getLogger()
logger.handlers = [RawStreamHandler(sys.stdout)]
logger.setLevel(INFO)
module_name = __name__.replace('.molecule','')
# Covalent radii from Cordero et al. 'Covalent radii revisited' Dalton Transactions 2008, 2832-2838.
Radii = [0.31, 0.28, # H and He
1.28, 0.96, 0.84, 0.76, 0.71, 0.66, 0.57, 0.58, # First row elements
0.00, 1.41, 1.21, 1.11, 1.07, 1.05, 1.02, 1.06, # Second row elements
# 1.66, 1.41, 1.21, 1.11, 1.07, 1.05, 1.02, 1.06, # Second row elements
2.03, 1.76, 1.70, 1.60, 1.53, 1.39, 1.61, 1.52, 1.50,
1.24, 1.32, 1.22, 1.22, 1.20, 1.19, 1.20, 1.20, 1.16, # Third row elements, K through Kr
2.20, 1.95, 1.90, 1.75, 1.64, 1.54, 1.47, 1.46, 1.42,
1.39, 1.45, 1.44, 1.42, 1.39, 1.39, 1.38, 1.39, 1.40, # Fourth row elements, Rb through Xe
2.44, 2.15, 2.07, 2.04, 2.03, 2.01, 1.99, 1.98,
1.98, 1.96, 1.94, 1.92, 1.92, 1.89, 1.90, 1.87, # Fifth row elements, s and f blocks
1.87, 1.75, 1.70, 1.62, 1.51, 1.44, 1.41, 1.36,
1.36, 1.32, 1.45, 1.46, 1.48, 1.40, 1.50, 1.50, # Fifth row elements, d and p blocks
2.60, 2.21, 2.15, 2.06, 2.00, 1.96, 1.90, 1.87, 1.80, 1.69] # Sixth row elements
# A list that gives you the element if you give it the atomic number, hence the 'none' at the front.
Elements = ["None",'H','He',
'Li','Be','B','C','N','O','F','Ne',
'Na','Mg','Al','Si','P','S','Cl','Ar',
'K','Ca','Sc','Ti','V','Cr','Mn','Fe','Co','Ni','Cu','Zn','Ga','Ge','As','Se','Br','Kr',
'Rb','Sr','Y','Zr','Nb','Mo','Tc','Ru','Rh','Pd','Ag','Cd','In','Sn','Sb','Te','I','Xe',
'Cs','Ba','La','Ce','Pr','Nd','Pm','Sm','Eu','Gd','Tb','Dy','Ho','Er','Tm','Yb',
'Lu','Hf','Ta','W','Re','Os','Ir','Pt','Au','Hg','Tl','Pb','Bi','Po','At','Rn',
'Fr','Ra','Ac','Th','Pa','U','Np','Pu','Am','Cm','Bk','Cf','Es','Fm','Md','No','Lr','Rf','Db','Sg','Bh','Hs','Mt']
# Dictionary of atomic masses ; also serves as the list of elements (periodic table)
PeriodicTable = OrderedDict([('H' , 1.0079), ('He' , 4.0026),
('Li' , 6.941), ('Be' , 9.0122), ('B' , 10.811), ('C' , 12.0107), ('N' , 14.0067), ('O' , 15.9994), ('F' , 18.9984), ('Ne' , 20.1797),
('Na' , 22.9897), ('Mg' , 24.305), ('Al' , 26.9815), ('Si' , 28.0855), ('P' , 30.9738), ('S' , 32.065), ('Cl' , 35.453), ('Ar' , 39.948),
('K' , 39.0983), ('Ca' , 40.078), ('Sc' , 44.9559), ('Ti' , 47.867), ('V' , 50.9415), ('Cr' , 51.9961), ('Mn' , 54.938), ('Fe' , 55.845), ('Co' , 58.9332),
('Ni' , 58.6934), ('Cu' , 63.546), ('Zn' , 65.39), ('Ga' , 69.723), ('Ge' , 72.64), ('As' , 74.9216), ('Se' , 78.96), ('Br' , 79.904), ('Kr' , 83.8),
('Rb' , 85.4678), ('Sr' , 87.62), ('Y' , 88.9059), ('Zr' , 91.224), ('Nb' , 92.9064), ('Mo' , 95.94), ('Tc' , 98), ('Ru' , 101.07), ('Rh' , 102.9055),
('Pd' , 106.42), ('Ag' , 107.8682), ('Cd' , 112.411), ('In' , 114.818), ('Sn' , 118.71), ('Sb' , 121.76), ('Te' , 127.6), ('I' , 126.9045), ('Xe' , 131.293),
('Cs' , 132.9055), ('Ba' , 137.327), ('La' , 138.9055), ('Ce' , 140.116), ('Pr' , 140.9077), ('Nd' , 144.24), ('Pm' , 145), ('Sm' , 150.36),
('Eu' , 151.964), ('Gd' , 157.25), ('Tb' , 158.9253), ('Dy' , 162.5), ('Ho' , 164.9303), ('Er' , 167.259), ('Tm' , 168.9342), ('Yb' , 173.04),
('Lu' , 174.967), ('Hf' , 178.49), ('Ta' , 180.9479), ('W' , 183.84), ('Re' , 186.207), ('Os' , 190.23), ('Ir' , 192.217), ('Pt' , 195.078),
('Au' , 196.9665), ('Hg' , 200.59), ('Tl' , 204.3833), ('Pb' , 207.2), ('Bi' , 208.9804), ('Po' , 209), ('At' , 210), ('Rn' , 222),
('Fr' , 223), ('Ra' , 226), ('Ac' , 227), ('Th' , 232.0381), ('Pa' , 231.0359), ('U' , 238.0289), ('Np' , 237), ('Pu' , 244),
('Am' , 243), ('Cm' , 247), ('Bk' , 247), ('Cf' , 251), ('Es' , 252), ('Fm' , 257), ('Md' , 258), ('No' , 259),
('Lr' , 262), ('Rf' , 261), ('Db' , 262), ('Sg' , 266), ('Bh' , 264), ('Hs' , 277), ('Mt' , 268)])
def getElement(mass):
return PeriodicTable.keys()[np.argmin([np.abs(m-mass) for m in PeriodicTable.values()])]
def elem_from_atomname(atomname):
""" Given an atom name, attempt to get the element in most cases. """
return re.search('[A-Z][a-z]*',atomname).group(0)
if "forcebalance" in __name__:
#============================#
#| DCD read/write functions |#
#============================#
# Try to load _dcdlib.so either from a directory in the LD_LIBRARY_PATH
# or from the same directory as this module.
try: _dcdlib = CDLL("_dcdlib.so")
except:
try: _dcdlib = CDLL(os.path.join(imp.find_module(__name__.split('.')[0])[1],"_dcdlib.so"))
except:
warn('The dcdlib module cannot be imported (Cannot read/write DCD files)')
#============================#
#| PDB read/write functions |#
#============================#
try: from PDB import *
except:
warn('The pdb module cannot be miported (Cannot read/write PDB files)')
#=============================#
#| Mol2 read/write functions |#
#=============================#
try: import Mol2
except:
warn('The Mol2 module cannot be imported (Cannot read/write Mol2 files)')
#==============================#
#| OpenMM interface functions |#
#==============================#
try:
from simtk.unit import *
from simtk.openmm import *
from simtk.openmm.app import *
except:
warn('The OpenMM modules cannot be imported (Cannot interface with OpenMM)')
#===========================#
#| Convenience subroutines |#
#===========================#
## One bohr equals this many angstroms
bohrang = 0.529177249
def unmangle(M1, M2):
"""
Create a mapping that takes M1's atom indices to M2's atom indices based on position.
If we start with atoms in molecule "PDB", and the new molecule "M"
contains re-numbered atoms, then this code works:
M.elem = list(np.array(PDB.elem)[unmangled])
"""
if M1.na != M2.na:
logger.error("Unmangler only deals with same number of atoms\n")
raise RuntimeError
unmangler = {}
for i in range(M1.na):
for j in range(M2.na):
if np.linalg.norm(M1.xyzs[0][i] - M2.xyzs[0][j]) < 0.1:
unmangler[j] = i
unmangled = [unmangler[i] for i in sorted(unmangler.keys())]
if len(unmangled) != M1.na:
logger.error("Unmangler failed (different structures?)\n")
raise RuntimeError
return unmangled
def nodematch(node1,node2):
# Matching two nodes of a graph. Nodes are equivalent if the elements are the same
return node1['e'] == node2['e']
def isint(word):
"""ONLY matches integers! If you have a decimal point? None shall pass!"""
return re.match('^[-+]?[0-9]+$',word)
def isfloat(word):
"""Matches ANY number; it can be a decimal, scientific notation, integer, or what have you"""
return re.match('^[-+]?[0-9]*\.?[0-9]*([eEdD][-+]?[0-9]+)?$',word)
# Used to get the white spaces in a split line.
splitter = re.compile(r'(\s+|\S+)')
# Container for Bravais lattice vector. Three cell lengths, three angles, three vectors, volume, and TINKER trig functions.
Box = namedtuple('Box',['a','b','c','alpha','beta','gamma','A','B','C','V'])
radian = 180. / np.pi
def CubicLattice(a):
""" This function takes in three lattice lengths and three lattice angles, and tries to return a complete box specification. """
b = a
c = a
alpha = 90
beta = 90
gamma = 90
alph = alpha*np.pi/180
bet = beta*np.pi/180
gamm = gamma*np.pi/180
v = np.sqrt(1 - cos(alph)**2 - cos(bet)**2 - cos(gamm)**2 + 2*cos(alph)*cos(bet)*cos(gamm))
Mat = np.matrix([[a, b*cos(gamm), c*cos(bet)],
[0, b*sin(gamm), c*((cos(alph)-cos(bet)*cos(gamm))/sin(gamm))],
[0, 0, c*v/sin(gamm)]])
L1 = Mat*np.matrix([[1],[0],[0]])
L2 = Mat*np.matrix([[0],[1],[0]])
L3 = Mat*np.matrix([[0],[0],[1]])
return Box(a,b,c,alpha,beta,gamma,np.array(L1).flatten(),np.array(L2).flatten(),np.array(L3).flatten(),v*a*b*c)
def BuildLatticeFromLengthsAngles(a, b, c, alpha, beta, gamma):
""" This function takes in three lattice lengths and three lattice angles, and tries to return a complete box specification. """
alph = alpha*np.pi/180
bet = beta*np.pi/180
gamm = gamma*np.pi/180
v = np.sqrt(1 - cos(alph)**2 - cos(bet)**2 - cos(gamm)**2 + 2*cos(alph)*cos(bet)*cos(gamm))
Mat = np.matrix([[a, b*cos(gamm), c*cos(bet)],
[0, b*sin(gamm), c*((cos(alph)-cos(bet)*cos(gamm))/sin(gamm))],
[0, 0, c*v/sin(gamm)]])
L1 = Mat*np.matrix([[1],[0],[0]])
L2 = Mat*np.matrix([[0],[1],[0]])
L3 = Mat*np.matrix([[0],[0],[1]])
return Box(a,b,c,alpha,beta,gamma,np.array(L1).flatten(),np.array(L2).flatten(),np.array(L3).flatten(),v*a*b*c)
def BuildLatticeFromVectors(v1, v2, v3):
""" This function takes in three lattice vectors and tries to return a complete box specification. """
a = np.linalg.norm(v1)
b = np.linalg.norm(v2)
c = np.linalg.norm(v3)
alpha = arccos(np.dot(v2, v3) / np.linalg.norm(v2) / np.linalg.norm(v3)) * radian
beta = arccos(np.dot(v1, v3) / np.linalg.norm(v1) / np.linalg.norm(v3)) * radian
gamma = arccos(np.dot(v1, v2) / np.linalg.norm(v1) / np.linalg.norm(v2)) * radian
alph = alpha*np.pi/180
bet = beta*np.pi/180
gamm = gamma*np.pi/180
v = np.sqrt(1 - cos(alph)**2 - cos(bet)**2 - cos(gamm)**2 + 2*cos(alph)*cos(bet)*cos(gamm))
Mat = np.matrix([[a, b*cos(gamm), c*cos(bet)],
[0, b*sin(gamm), c*((cos(alph)-cos(bet)*cos(gamm))/sin(gamm))],
[0, 0, c*v/sin(gamm)]])
L1 = Mat*np.matrix([[1],[0],[0]])
L2 = Mat*np.matrix([[0],[1],[0]])
L3 = Mat*np.matrix([[0],[0],[1]])
return Box(a,b,c,alpha,beta,gamma,np.array(L1).flatten(),np.array(L2).flatten(),np.array(L3).flatten(),v*a*b*c)
#===========================#
#| Connectivity graph |#
#| Good for doing simple |#
#| topology tricks |#
#===========================#
have_contact = 0
try:
import networkx as nx
class MyG(nx.Graph):
def __init__(self):
super(MyG,self).__init__()
self.Alive = True
def __eq__(self, other):
# This defines whether two MyG objects are "equal" to one another.
if not self.Alive:
return False
if not other.Alive:
return False
return nx.is_isomorphic(self,other,node_match=nodematch)
def __hash__(self):
''' The hash function is something we can use to discard two things that are obviously not equal. Here we neglect the hash. '''
return 1
def L(self):
''' Return a list of the sorted atom numbers in this graph. '''
return sorted(list(self.nodes()))
def AStr(self):
''' Return a string of atoms, which serves as a rudimentary 'fingerprint' : '99,100,103,151' . '''
return ','.join(['%i' % i for i in self.L()])
def e(self):
''' Return an array of the elements. For instance ['H' 'C' 'C' 'H']. '''
elems = nx.get_node_attributes(self,'e')
return [elems[i] for i in self.L()]
def ef(self):
''' Create an Empirical Formula '''
Formula = list(self.e())
return ''.join([('%s%i' % (k, Formula.count(k)) if Formula.count(k) > 1 else '%s' % k) for k in sorted(set(Formula))])
def x(self):
''' Get a list of the coordinates. '''
coors = nx.get_node_attributes(self,'x')
return np.array([coors[i] for i in self.L()])
try:
import contact
have_contact = 1
except:
warn("'contact' cannot be imported (topology tools will be slow.)")
except:
warn("NetworkX cannot be imported (topology tools won't work). Most functionality should still work though.")
def TopEqual(mol1, mol2):
""" For the nanoreactor project: Determine whether two Molecule objects have the same topologies. """
# Checks to see if the two molecule objects have the same fragments.
GraphEqual = Counter(mol1.molecules) == Counter(mol2.molecules)
# Checks to see whether the molecule objects have the same atoms in each fragment.
AtomEqual = Counter([tuple(m.L()) for m in mol1.molecules]) == Counter([tuple(m.L()) for m in mol2.molecules])
return GraphEqual and AtomEqual
def MolEqual(mol1, mol2):
"""
Determine whether two Molecule objects have the same fragments by
looking at elements and connectivity graphs. This is less strict
than TopEqual (i.e. more often returns True).
"""
if mol1.na != mol2.na : return False
if Counter(mol1.elem) != Counter(mol2.elem) : return False
return Counter(mol1.molecules) == Counter(mol2.molecules)
def format_xyz_coord(element,xyz,tinker=False):
""" Print a line consisting of (element, x, y, z) in accordance with .xyz file format
@param[in] element A chemical element of a single atom
@param[in] xyz A 3-element array containing x, y, z coordinates of that atom
"""
if tinker:
return "%-3s % 13.8f % 13.8f % 13.8f" % (element,xyz[0],xyz[1],xyz[2])
else:
return "%-5s % 15.10f % 15.10f % 15.10f" % (element,xyz[0],xyz[1],xyz[2])
def format_gro_coord(resid, resname, aname, seqno, xyz):
""" Print a line in accordance with .gro file format, with six decimal points of precision
Nine decimal points of precision are necessary to get forces below 1e-3 kJ/mol/nm.
@param[in] resid The number of the residue that the atom belongs to
@param[in] resname The name of the residue that the atom belongs to
@param[in] aname The name of the atom
@param[in] seqno The sequential number of the atom
@param[in] xyz A 3-element array containing x, y, z coordinates of that atom
"""
return "%5i%-5s%5s%5i % 13.9f % 13.9f % 13.9f" % (resid,resname,aname,seqno,xyz[0],xyz[1],xyz[2])
def format_xyzgen_coord(element,xyzgen):
""" Print a line consisting of (element, p, q, r, s, t, ...) where
(p, q, r) are arbitrary atom-wise data (this might happen, for
instance, with atomic charges)
@param[in] element A chemical element of a single atom
@param[in] xyzgen A N-element array containing data for that atom
"""
return "%-5s" + ' '.join(["% 15.10f" % i] for i in xyzgen)
def format_gro_box(box):
""" Print a line corresponding to the box vector in accordance with .gro file format
@param[in] box Box NamedTuple
"""
if box.alpha == 90.0 and box.beta == 90.0 and box.gamma == 90.0:
return ' '.join(["% 13.9f" % (i/10) for i in [box.a, box.b, box.c]])
else:
return ' '.join(["% 13.9f" % (i/10) for i in [box.A[0], box.B[1], box.C[2], box.A[1], box.A[2], box.B[0], box.B[2], box.C[0], box.C[1]]])
def is_gro_coord(line):
""" Determines whether a line contains GROMACS data or not
@param[in] line The line to be tested
"""
sline = line.split()
if len(sline) == 6:
return all([isint(sline[2]),isfloat(sline[3]),isfloat(sline[4]),isfloat(sline[5])])
elif len(sline) == 5:
return all([isint(line[15:20]),isfloat(sline[2]),isfloat(sline[3]),isfloat(sline[4])])
else:
return 0
def is_charmm_coord(line):
""" Determines whether a line contains CHARMM data or not
@param[in] line The line to be tested
"""
sline = line.split()
if len(sline) >= 7:
return all([isint(sline[0]), isint(sline[1]), isfloat(sline[4]), isfloat(sline[5]), isfloat(sline[6])])
else:
return 0
def is_gro_box(line):
""" Determines whether a line contains a GROMACS box vector or not
@param[in] line The line to be tested
"""
sline = line.split()
if len(sline) == 9 and all([isfloat(i) for i in sline]):
return 1
elif len(sline) == 3 and all([isfloat(i) for i in sline]):
return 1
else:
return 0
def add_strip_to_mat(mat,strip):
out = list(mat)
if out == [] and strip != []:
out = list(strip)
elif out != [] and strip != []:
for (i,j) in zip(out,strip):
i += list(j)
return out
def pvec(vec):
return ''.join([' % .10e' % i for i in list(vec.flatten())])
def grouper(n, iterable):
""" Groups a big long iterable into groups of ten or what have you. """
args = [iter(iterable)] * n
return list([e for e in t if e is not None] for t in itertools.izip_longest(*args))
def even_list(totlen, splitsize):
""" Creates a list of number sequences divided as evenly as possible. """
joblens = np.zeros(splitsize,dtype=int)
subsets = []
for i in range(totlen):
joblens[i%splitsize] += 1
jobnow = 0
for i in range(splitsize):
subsets.append(range(jobnow, jobnow + joblens[i]))
jobnow += joblens[i]
return subsets
class MolfileTimestep(Structure):
""" Wrapper for the timestep C structure used in molfile plugins. """
_fields_ = [("coords",POINTER(c_float)), ("velocities",POINTER(c_float)),
("A",c_float), ("B",c_float), ("C",c_float), ("alpha",c_float),
("beta",c_float), ("gamma",c_float), ("physical_time",c_double)]
def both(A, B, key):
return key in A.Data and key in B.Data
def diff(A, B, key):
if not (key in A.Data and key in B.Data) : return False
else:
if type(A.Data[key]) is np.ndarray:
return (A.Data[key] != B.Data[key]).any()
elif key == 'tinkersuf':
return [sorted([int(j) for j in i.split()]) for i in A.Data[key]] != \
[sorted([int(j) for j in i.split()]) for i in B.Data[key]]
else:
return A.Data[key] != B.Data[key]
def either(A, B, key):
return key in A.Data or key in B.Data
#===========================#
#| Alignment subroutines |#
#| Moments added 08/03/12 |#
#===========================#
def EulerMatrix(T1,T2,T3):
""" Constructs an Euler matrix from three Euler angles. """
DMat = np.matrix(np.zeros((3,3)))
DMat[0,0] = np.cos(T1)
DMat[0,1] = np.sin(T1)
DMat[1,0] = -np.sin(T1)
DMat[1,1] = np.cos(T1)
DMat[2,2] = 1
CMat = np.matrix(np.zeros((3,3)))
CMat[0,0] = 1
CMat[1,1] = np.cos(T2)
CMat[1,2] = np.sin(T2)
CMat[2,1] = -np.sin(T2)
CMat[2,2] = np.cos(T2)
BMat = np.matrix(np.zeros((3,3)))
BMat[0,0] = np.cos(T3)
BMat[0,1] = np.sin(T3)
BMat[1,0] = -np.sin(T3)
BMat[1,1] = np.cos(T3)
BMat[2,2] = 1
EMat = BMat*CMat*DMat
return np.matrix(EMat)
def ComputeOverlap(theta,elem,xyz1,xyz2):
"""
Computes an 'overlap' between two molecules based on some
fictitious density. Good for fine-tuning alignment but gets stuck
in local minima.
"""
xyz2R = np.array(EulerMatrix(theta[0],theta[1],theta[2])*np.matrix(xyz2.T)).T
Obj = 0.0
elem = np.array(elem)
for i in set(elem):
for j in np.where(elem==i)[0]:
for k in np.where(elem==i)[0]:
dx = xyz1[j] - xyz2R[k]
dx2 = np.dot(dx,dx)
Obj -= np.exp(-0.5*dx2)
return Obj
def AlignToDensity(elem,xyz1,xyz2,binary=False):
"""
Computes a "overlap density" from two frames.
This function can be called by AlignToMoments to get rid of inversion problems
"""
grid = np.pi*np.array(list(itertools.product([0,1],[0,1],[0,1])))
ovlp = np.array([ComputeOverlap(e, elem, xyz1, xyz2) for e in grid]) # Mao
t1 = grid[np.argmin(ovlp)]
xyz2R = (np.array(EulerMatrix(t1[0],t1[1],t1[2])*np.matrix(xyz2.T)).T).copy()
return xyz2R
def AlignToMoments(elem,xyz1,xyz2=None):
"""Pre-aligns molecules to 'moment of inertia'.
If xyz2 is passed in, it will assume that xyz1 is already
aligned to the moment of inertia, and it simply does 180-degree
rotations to make sure nothing is inverted."""
xyz = xyz1 if xyz2 is None else xyz2
I = np.zeros((3,3))
for i, xi in enumerate(xyz):
I += (np.dot(xi,xi)*np.eye(3) - np.outer(xi,xi))
# This is the original line from MSMBuilder, but we're choosing not to use masses
# I += PeriodicTable[elem[i]]*(np.dot(xi,xi)*np.eye(3) - np.outer(xi,xi))
A, B = np.linalg.eig(I)
# Sort eigenvectors by eigenvalue
BB = B[:, np.argsort(A)]
determ = np.linalg.det(BB)
Thresh = 1e-3
if np.abs(determ - 1.0) > Thresh:
if np.abs(determ + 1.0) > Thresh:
print "in AlignToMoments, determinant is % .3f" % determ
BB[:,2] *= -1
xyzr = np.array(np.matrix(BB).T * np.matrix(xyz).T).T.copy()
if xyz2 is not None:
xyzrr = AlignToDensity(elem,xyz1,xyzr,binary=True)
return xyzrr
else:
return xyzr
def get_rotate_translate(matrix1,matrix2):
# matrix2 contains the xyz coordinates of the REFERENCE
assert np.shape(matrix1) == np.shape(matrix2), 'Matrices not of same dimensions'
# Store number of rows
nrows = np.shape(matrix1)[0]
# Getting centroid position for each selection
avg_pos1 = matrix1.sum(axis=0)/nrows
avg_pos2 = matrix2.sum(axis=0)/nrows
# Translation of matrices
avg_matrix1 = matrix1-avg_pos1
avg_matrix2 = matrix2-avg_pos2
# Covariance matrix
covar = np.dot(avg_matrix1.T,avg_matrix2)
# Do the SVD in order to get rotation matrix
v,s,wt = np.linalg.svd(covar)
v = np.matrix(v)
wt = np.matrix(wt)
# Rotation matrix
# Transposition of v,wt
wvt = wt.T*v.T
# Ensure a right-handed coordinate system
d = np.matrix(np.eye(3))
if np.linalg.det(wvt) < 0:
d[2,2] = -1.0
rot_matrix = np.array((wt.T*d*v.T).T)
# rot_matrix = np.transpose(np.dot(np.transpose(wt),np.transpose(v)))
trans_matrix = avg_pos2-np.dot(avg_pos1,rot_matrix)
return trans_matrix, rot_matrix
def cartesian_product2(arrays):
""" Form a Cartesian product of two NumPy arrays. """
la = len(arrays)
arr = np.empty([len(a) for a in arrays] + [la], dtype=np.int32)
for i, a in enumerate(np.ix_(*arrays)):
arr[...,i] = a
return arr.reshape(-1, la)
def extract_int(arr, avgthre, limthre, label="value", verbose=True):
"""
Get the representative integer value from an array.
The integer value is the rounded mean. Perform sanity
checks to ensure the array isn't overly "non-integer".
Parameters
----------
arr : numpy.ndarray
NumPy array containing a series of floating point values
where we'd like to extract the representative integer value.
avgthre : float
If the average deviates from the closest integer by
more than this amount, do not pass.
limthre : float
If any element in this array deviates from the closest integer by
more than this amount, do not pass.
label : str
Descriptive name of this variable, used only in printout.
verbose : bool
Print information in case array makes excursions larger than the threshold
Returns
-------
int
Representative integer value for the array.
passed : bool
Indicates whether the array mean and/or maximum deviations stayed with the thresholds.
"""
average = np.mean(arr)
maximum = np.max(arr)
minimum = np.min(arr)
rounded = round(average)
passed = True
if abs(average - rounded) > avgthre:
if verbose:
logger.info("Average %s (%f) deviates from integer %s (%i) by more than threshold of %f" % (label, average, label, rounded, avgthre))
passed = False
if abs(maximum - minimum) > limthre:
if verbose:
logger.info("Maximum %s fluctuation (%f) is larger than threshold of %f" % (label, abs(maximum-minimum), limthre))
passed = False
return int(rounded), passed
def extract_pop(M, verbose=True):
"""
Extract our best estimate of charge and spin-z from the comments
section of a Molecule object created with Nanoreactor. Note that
spin-z is 1.0 if there is one unpaired electron (not one/half) because
the unit is in terms of populations.
This function is intended to work on atoms that are *extracted*
from an ab initio MD trajectory, where the Mulliken charge and spin
populations are not exactly integers. It attempts to return the closest
integer but it will sometimes fail.
If the number of electrons and spin-z are inconsistent, then
return -999,-999 (indicates failure).
Parameters
----------
M : Molecule
Molecule object that we're getting charge and spin from, with a known comment syntax
Reaction: formula CHO -> CO+H atoms ['0-2'] -> ['0-1','2'] frame 219482 charge -0.742 sz +0.000 sz^2 0.000
Returns
-------
chg : int
Representative integer net charge of this Molecule, or -999 if inconsistent
spn : int
Representative integer spin-z of this Molecule (one unpaired electron: sz=1), or -999 if inconsistent
"""
# Read in the charge and spin on the whole system.
srch = lambda s : np.array([float(re.search('(?<=%s )[-+]?[0-9]*\.?[0-9]*([eEdD][-+]?[0-9]+)?' % s, c).group(0)) for c in M.comms if all([i in c for i in 'charge', 'sz'])])
Chgs = srch('charge') # An array of the net charge.
SpnZs = srch('sz') # An array of the net Z-spin.
Spn2s = srch('sz\^2') # An array of the sum of sz^2 by atom.
chg, chgpass = extract_int(Chgs, 0.3, 1.0, label="charge")
spn, spnpass = extract_int(abs(SpnZs), 0.3, 1.0, label="spin-z")
# Try to calculate the correct spin.
nproton = sum([Elements.index(i) for i in M.elem])
nelectron = nproton + chg
if not spnpass:
if verbose: logger.info("Going with the minimum spin consistent with charge.")
if nelectron%2 == 0:
spn = 0
else:
spn = 1
# The number of electrons should be odd iff the spin is odd.
if ((nelectron-spn)/2)*2 != (nelectron-spn):
if verbose: logger.info("\x1b[91mThe number of electrons (%i) is inconsistent with the spin-z (%i)\x1b[0m" % (nelectron, spn))
return -999, -999
if verbose: logger.info("%i electrons; charge %i, spin %i" % (nelectron, chg, spn))
return chg, spn
def arc(Mol, begin=None, end=None, RMSD=True, align=True):
"""
Get the arc-length for a trajectory segment.
Uses RMSD or maximum displacement of any atom in the trajectory.
Parameters
----------
Mol : Molecule
Molecule object for calculating the arc length.
begin : int
Starting frame, defaults to first frame
end : int
Ending frame, defaults to final frame
RMSD : bool
Set to True to use frame-to-frame RMSD; otherwise use the maximum displacement of any atom
Returns
-------
Arc : np.ndarray
Arc length between frames in Angstrom, length is n_frames - 1
"""
if align:
Mol.align()
if begin is None:
begin = 0
if end is None:
end = len(Mol)
if RMSD:
Arc = Mol.pathwise_rmsd(align)
else:
Arc = np.array([np.max([np.linalg.norm(Mol.xyzs[i+1][j]-Mol.xyzs[i][j]) for j in range(Mol.na)]) for i in range(begin, end-1)])
return Arc
def EqualSpacing(Mol, frames=0, dx=0, RMSD=True, align=True):
"""
Equalize the spacing of frames in a trajectory with linear interpolation.
This is done in a very simple way, first calculating the arc length
between frames, then creating an equally spaced array, and interpolating
all Cartesian coordinates along this equally spaced array.
This is intended to be used on trajectories with smooth transformations and
ensures that concatenated frames containing both optimization coordinates
and dynamics trajectories don't have sudden changes in their derivatives.
Parameters
----------
Mol : Molecule
Molecule object for equalizing the spacing.
frames : int
Return a Molecule object with this number of frames.
RMSD : bool
Use RMSD in the arc length calculation.
Returns
-------
Mol1 : Molecule
New molecule object, either the same one (if frames > len(Mol))
or with equally spaced frames.
"""
ArcMol = arc(Mol, RMSD=RMSD, align=align)
ArcMolCumul = np.insert(np.cumsum(ArcMol), 0, 0.0)
if frames != 0 and dx != 0:
logger.error("Provide dx or frames or neither")
elif dx != 0:
frames = int(float(max(ArcMolCumul))/dx)
elif frames == 0:
frames = len(ArcMolCumul)
ArcMolEqual = np.linspace(0, max(ArcMolCumul), frames)
xyzold = np.array(Mol.xyzs)
xyznew = np.zeros((frames, Mol.na, 3))
for a in range(Mol.na):
for i in range(3):
xyznew[:,a,i] = np.interp(ArcMolEqual, ArcMolCumul, xyzold[:, a, i])
if len(xyzold) == len(xyznew):
Mol1 = copy.deepcopy(Mol)
else:
# If we changed the number of coordinates, then
# do some integer interpolation of the comments and
# other frame variables.
Mol1 = Mol[np.array([int(round(i)) for i in np.linspace(0, len(xyzold)-1, len(xyznew))])]
Mol1.xyzs = list(xyznew)
return Mol1
class Molecule(object):
""" Lee-Ping's general file format conversion class.
The purpose of this class is to read and write chemical file formats in a
way that is convenient for research. There are highly general file format
converters out there (e.g. catdcd, openbabel) but I find that writing
my own class can be very helpful for specific purposes. Here are some things
this class can do:
- Convert a .gro file to a .xyz file, or a .pdb file to a .dcd file.
Data is stored internally, so any readable file can be converted into
any writable file as long as there is sufficient information to write
that file.
- Accumulate information from different files. For example, we may read
A.gro to get a list of coordinates, add quantum settings from a B.in file,
and write A.in (this gives us a file that we can use to run QM calculations)
- Concatenate two trajectories together as long as they're compatible. This
is done by creating two Molecule objects and then simply adding them. Addition
means two things: (1) Information fields missing from each class, but present
in the other, are added to the sum, and (2) Appendable or per-frame fields
(i.e. coordinates) are concatenated together.
- Slice trajectories using reasonable Python language. That is to
say, MyMolecule[1:10] returns a new Molecule object that contains
frames 1 through 9 (inclusive and numbered starting from zero.)
Special variables: These variables cannot be set manually because
there is a special method associated with getting them.
na = The number of atoms. You'll get this if you use MyMol.na or MyMol['na'].
ns = The number of snapshots. You'll get this if you use MyMol.ns or MyMol['ns'].
Unit system: Angstroms.
"""
def __len__(self):
""" Return the number of frames in the trajectory. """
L = -1
klast = None
Defined = False
for key in self.FrameKeys:
Defined = True
if L != -1 and len(self.Data[key]) != L:
self.repair(key, klast)
L = len(self.Data[key])
klast = key
if not Defined:
return 0
return L
def __getattr__(self, key):
""" Whenever we try to get a class attribute, it first tries to get the attribute from the Data dictionary. """
if key == 'qm_forces':
warn('qm_forces is a deprecated keyword because it actually meant gradients; setting to qm_grads.')
key = 'qm_grads'
if key == 'ns':
return len(self)
elif key == 'na': # The 'na' attribute is the number of atoms.
L = -1
klast = None
Defined = False
for key in self.AtomKeys:
Defined = True
if L != -1 and len(self.Data[key]) != L:
self.repair(key, klast)
L = len(self.Data[key])
klast = key
if Defined:
return L
elif 'xyzs' in self.Data:
return len(self.xyzs[0])
else:
return 0
#raise RuntimeError('na is ill-defined if the molecule has no AtomKeys member variables.')
## These attributes return a list of attribute names defined in this class that belong in the chosen category.
## For example: self.FrameKeys should return set(['xyzs','boxes']) if xyzs and boxes exist in self.Data
elif key == 'FrameKeys':
return set(self.Data) & FrameVariableNames
elif key == 'AtomKeys':
return set(self.Data) & AtomVariableNames
elif key == 'MetaKeys':
return set(self.Data) & MetaVariableNames
elif key == 'QuantumKeys':
return set(self.Data) & QuantumVariableNames
elif key in self.Data:
return self.Data[key]
return getattr(super(Molecule, self), key)
def __setattr__(self, key, value):
""" Whenever we try to get a class attribute, it first tries to get the attribute from the Data dictionary. """
## These attributes return a list of attribute names defined in this class, that belong in the chosen category.
## For example: self.FrameKeys should return set(['xyzs','boxes']) if xyzs and boxes exist in self.Data
if key == 'qm_forces':
warn('qm_forces is a deprecated keyword because it actually meant gradients; setting to qm_grads.')
key = 'qm_grads'
if key in AllVariableNames:
self.Data[key] = value
return super(Molecule,self).__setattr__(key, value)
def __getitem__(self, key):
"""
The Molecule class has list-like behavior, so we can get slices of it.
If we say MyMolecule[0:10], then we'll return a copy of MyMolecule with frames 0 through 9.
"""
if isinstance(key, int) or isinstance(key, slice) or isinstance(key,np.ndarray):
if isinstance(key, int):
key = [key]
New = Molecule()
for k in self.FrameKeys:
if k == 'boxes':
New.Data[k] = [j for i, j in enumerate(self.Data[k]) if i in np.arange(len(self))[key]]
else:
New.Data[k] = list(np.array(copy.deepcopy(self.Data[k]))[key])