-
Notifications
You must be signed in to change notification settings - Fork 674
/
Copy pathtest_atomselections.py
1193 lines (951 loc) · 40.3 KB
/
test_atomselections.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
# -*- Mode: python; tab-width: 4; indent-tabs-mode:nil; coding:utf-8 -*-
# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4 fileencoding=utf-8
#
# MDAnalysis --- https://www.mdanalysis.org
# Copyright (c) 2006-2017 The MDAnalysis Development Team and contributors
# (see the file AUTHORS for the full list of names)
#
# Released under the GNU Public Licence, v2 or any higher version
#
# Please cite your use of MDAnalysis in published work:
#
# R. J. Gowers, M. Linke, J. Barnoud, T. J. E. Reddy, M. N. Melo, S. L. Seyler,
# D. L. Dotson, J. Domanski, S. Buchoux, I. M. Kenney, and O. Beckstein.
# MDAnalysis: A Python package for the rapid analysis of molecular dynamics
# simulations. In S. Benthall and S. Rostrup editors, Proceedings of the 15th
# Python in Science Conference, pages 102-109, Austin, TX, 2016. SciPy.
# doi: 10.25080/majora-629e541a-00e
#
# N. Michaud-Agrawal, E. J. Denning, T. B. Woolf, and O. Beckstein.
# MDAnalysis: A Toolkit for the Analysis of Molecular Dynamics Simulations.
# J. Comput. Chem. 32 (2011), 2319--2327, doi:10.1002/jcc.21787
#
import os
import itertools
import numpy as np
from numpy.testing import(
assert_equal,
)
import MDAnalysis
import MDAnalysis as mda
import MDAnalysis.core.selection
from MDAnalysis.lib.distances import distance_array
from MDAnalysis.core.selection import Parser
from MDAnalysis import SelectionError
from MDAnalysis.tests.datafiles import (
PSF, DCD,
PRMpbc, TRJpbc_bz2,
PSF_NAMD, PDB_NAMD,
GRO, RNA_PSF, NUCLsel, TPR, XTC,
TRZ_psf, TRZ,
PDB_icodes,
PDB_HOLE,
)
from MDAnalysisTests import make_Universe
import pytest
class TestSelectionsCHARMM(object):
@pytest.fixture(scope='class')
def universe(self):
"""Set up the standard AdK system in implicit solvent.
Geometry here is orthogonal
"""
return MDAnalysis.Universe(PSF, DCD)
@pytest.fixture()
def universe_copy(self, universe):
return MDAnalysis.Universe(PSF, DCD)
def test_segid(self, universe):
sel = universe.select_atoms('segid 4AKE')
assert_equal(sel.n_atoms, 3341, "failed to select segment 4AKE")
assert_equal(sorted(sel.indices),
sorted(universe.select_atoms('segid 4AKE').indices),
"selected segment 4AKE is not the same as auto-generated segment s4AKE")
def test_protein(self, universe):
sel = universe.select_atoms('protein')
assert_equal(sel.n_atoms, 3341, "failed to select protein")
assert_equal(sorted(sel.indices),
sorted(universe.select_atoms('segid 4AKE').indices),
"selected protein is not the same as auto-generated protein segment s4AKE")
@pytest.mark.parametrize('resname', sorted(MDAnalysis.core.selection.ProteinSelection.prot_res))
def test_protein_resnames(self, resname):
u = make_Universe(('resnames',))
# set half the residues' names to the resname we're testing
myprot = u.residues[::2]
# Windows note: the parametrized test input string objects
# are actually of type 'numpy.str_' and coercion to str
# proper is needed for unit test on Windows
myprot.resnames = str(resname)
# select protein
sel = u.select_atoms('protein')
# check that contents (atom indices) are identical afterwards
assert_equal(myprot.atoms.ix, sel.ix)
def test_backbone(self, universe):
sel = universe.select_atoms('backbone')
assert_equal(sel.n_atoms, 855)
def test_resid_single(self, universe):
sel = universe.select_atoms('resid 100')
assert_equal(sel.n_atoms, 7)
assert_equal(sel.residues.resnames, ['GLY'])
def test_resid_range(self, universe):
sel = universe.select_atoms('resid 100:105')
assert_equal(sel.n_atoms, 89)
assert_equal(sel.residues.resnames,
['GLY', 'ILE', 'ASN', 'VAL', 'ASP', 'TYR'])
def test_selgroup(self, universe):
sel = universe.select_atoms('not resid 100')
sel2 = universe.select_atoms('not group notr100', notr100=sel)
assert_equal(sel2.n_atoms, 7)
assert_equal(sel2.residues.resnames, ['GLY'])
def test_fullselgroup(self, universe):
sel1 = universe.select_atoms('resid 101')
sel2 = universe.select_atoms('resid 100')
sel3 = sel1.select_atoms('global group r100', r100=sel2)
assert_equal(sel2.n_atoms, 7)
assert_equal(sel2.residues.resnames, ['GLY'])
# resnum selections are boring here because we haven't really a mechanism yet
# to assign the canonical PDB resnums
def test_resnum_single(self, universe):
sel = universe.select_atoms('resnum 100')
assert_equal(sel.n_atoms, 7)
assert_equal(sel.residues.resids, [100])
assert_equal(sel.residues.resnames, ['GLY'])
def test_resnum_range(self, universe):
sel = universe.select_atoms('resnum 100:105')
assert_equal(sel.n_atoms, 89)
assert_equal(sel.residues.resids, range(100, 106))
assert_equal(sel.residues.resnames,
['GLY', 'ILE', 'ASN', 'VAL', 'ASP', 'TYR'])
def test_resname(self, universe):
sel = universe.select_atoms('resname LEU')
assert_equal(sel.n_atoms, 304,
"Failed to find all 'resname LEU' atoms.")
assert_equal(sel.n_residues, 16,
"Failed to find all 'resname LEU' residues.")
assert_equal(sorted(sel.indices),
sorted(universe.select_atoms('segid 4AKE and resname LEU').indices),
"selected 'resname LEU' atoms are not the same as auto-generated s4AKE.LEU")
def test_name(self, universe):
sel = universe.select_atoms('name CA')
assert_equal(sel.n_atoms, 214)
def test_atom(self, universe):
sel = universe.select_atoms('atom 4AKE 100 CA')
assert_equal(len(sel), 1)
assert_equal(sel.resnames, ['GLY'])
assert_equal(
sel.positions,
np.array([[20.38685226, -3.44224262, -5.92158318]],
dtype=np.float32))
def test_atom_empty(self, universe):
sel = universe.select_atoms('atom 4AKE 100 XX') # Does not exist
assert_equal(len(sel), 0)
def test_type(self, universe):
sel = universe.select_atoms("type 1")
assert_equal(len(sel), 253)
def test_and(self, universe):
sel = universe.select_atoms('resname GLY and resid 100')
assert_equal(len(sel), 7)
def test_or(self, universe):
sel = universe.select_atoms('resname LYS or resname ARG')
assert_equal(sel.n_residues, 31)
def test_not(self, universe):
sel = universe.select_atoms('not backbone')
assert_equal(len(sel), 2486)
@pytest.mark.parametrize('selstr', [
'around 4.0 bynum 1943',
'around 4.0 index 1942'
])
def test_around(self, universe, selstr):
sel = universe.select_atoms(selstr)
assert_equal(len(sel), 32)
@pytest.mark.parametrize('selstr', [
'sphlayer 4.0 6.0 bynum 1281',
'sphlayer 4.0 6.0 index 1280'
])
def test_sphlayer(self, universe, selstr):
sel = universe.select_atoms(selstr)
assert_equal(len(sel), 66)
@pytest.mark.parametrize('selstr', [
'sphzone 6.0 bynum 1281',
'sphzone 6.0 index 1280'
])
def test_sphzone(self, universe, selstr):
sel = universe.select_atoms(selstr)
assert_equal(len(sel), 86)
@pytest.mark.parametrize('selstr', [
'cylayer 4.0 6.0 10 -10 bynum 1281',
'cylayer 4.0 6.0 10 -10 index 1280'
])
def test_cylayer(self, universe, selstr):
sel = universe.select_atoms(selstr)
assert_equal(len(sel), 88)
@pytest.mark.parametrize('selstr', [
'cyzone 6.0 10 -10 bynum 1281',
'cyzone 6.0 10 -10 index 1280'
])
def test_cyzone(self, universe, selstr):
sel = universe.select_atoms(selstr)
assert_equal(len(sel), 166)
def test_point(self, universe):
ag = universe.select_atoms('point 5.0 5.0 5.0 3.5')
d = distance_array(np.array([[5.0, 5.0, 5.0]], dtype=np.float32),
universe.atoms.positions,
box=universe.dimensions)
idx = np.where(d < 3.5)[1]
assert_equal(set(ag.indices), set(idx))
def test_prop(self, universe):
sel = universe.select_atoms('prop y <= 16')
sel2 = universe.select_atoms('prop abs z < 8')
assert_equal(len(sel), 3194)
assert_equal(len(sel2), 2001)
def test_bynum(self, universe):
"Tests the bynum selection, also from AtomGroup instances (Issue 275)"
sel = universe.select_atoms('bynum 5')
assert_equal(sel[0].index, 4)
sel = universe.select_atoms('bynum 1:10')
assert_equal(len(sel), 10)
assert_equal(sel[0].index, 0)
assert_equal(sel[-1].index, 9)
subsel = sel.select_atoms('bynum 5')
assert_equal(subsel[0].index, 4)
subsel = sel.select_atoms('bynum 2:5')
assert_equal(len(subsel), 4)
assert_equal(subsel[0].index, 1)
assert_equal(subsel[-1].index, 4)
def test_index(self, universe):
"Tests the index selection, also from AtomGroup instances (Issue 275)"
sel = universe.select_atoms('index 4')
assert_equal(sel[0].index, 4)
sel = universe.select_atoms('index 0:9')
assert_equal(len(sel), 10)
assert_equal(sel[0].index, 0)
assert_equal(sel[-1].index, 9)
subsel = sel.select_atoms('index 4')
assert_equal(subsel[0].index, 4)
subsel = sel.select_atoms('index 1:4')
assert_equal(len(subsel), 4)
assert_equal(subsel[0].index, 1)
assert_equal(subsel[-1].index, 4)
@pytest.mark.parametrize('selstr,expected', (
['byres bynum 0:5', 19],
['byres index 0:19', 43]
))
def test_byres(self, universe, selstr, expected):
sel = universe.select_atoms(selstr)
assert_equal(len(sel), expected)
def test_same_resname(self, universe):
"""Test the 'same ... as' construct (Issue 217)"""
sel = universe.select_atoms("same resname as resid 10 or resid 11")
assert_equal(len(sel), 331,
("Found a wrong number of atoms with same resname as "
"resids 10 or 11"))
target_resids = np.array([7, 8, 10, 11, 12, 14, 17, 25, 32, 37, 38,
42, 46, 49, 55, 56, 66, 73, 80, 85, 93, 95,
99, 100, 122, 127, 130, 144, 150, 176, 180,
186, 188, 189, 194, 198, 203, 207, 214])
assert_equal(sel.residues.resids, target_resids,
("Found wrong residues with same resname as "
"resids 10 or 11"))
def test_same_segment(self, universe_copy):
"""Test the 'same ... as' construct (Issue 217)"""
SNew_A = universe_copy.add_Segment(segid='A')
SNew_B = universe_copy.add_Segment(segid='B')
SNew_C = universe_copy.add_Segment(segid='C')
universe_copy.residues[:100].segments = SNew_A
universe_copy.residues[100:150].segments = SNew_B
universe_copy.residues[150:].segments = SNew_C
target_resids = np.arange(100) + 1
sel = universe_copy.select_atoms("same segment as resid 10")
assert_equal(len(sel), 1520,
"Found a wrong number of atoms in the same segment of resid 10")
assert_equal(sel.residues.resids,
target_resids,
"Found wrong residues in the same segment of resid 10")
target_resids = np.arange(100, 150) + 1
sel = universe_copy.select_atoms("same segment as resid 110")
assert_equal(len(sel), 797,
"Found a wrong number of atoms in the same segment of resid 110")
assert_equal(sel.residues.resids, target_resids,
"Found wrong residues in the same segment of resid 110")
target_resids = np.arange(150, universe_copy.atoms.n_residues) + 1
sel = universe_copy.select_atoms("same segment as resid 160")
assert_equal(len(sel), 1024,
"Found a wrong number of atoms in the same segment of resid 160")
assert_equal(sel.residues.resids, target_resids,
"Found wrong residues in the same segment of resid 160")
def test_empty_same(self, universe):
ag = universe.select_atoms('resname MET')
# No GLY, so 'as resname GLY' is empty
ag2 = ag.select_atoms('same mass as resname GLY')
assert len(ag2) == 0
def test_empty_selection(self, universe):
"""Test that empty selection can be processed (see Issue 12)"""
# no Trp in AdK
assert_equal(len(universe.select_atoms('resname TRP')), 0)
def test_parenthesized_expression(self, universe):
sel = universe.select_atoms(
'( name CA or name CB ) and resname LEU')
assert_equal(len(sel), 32)
def test_no_space_around_parentheses(self, universe):
"""Test that no space is needed around parentheses (Issue 43)."""
# note: will currently be ERROR because it throws a ParseError
sel = universe.select_atoms('(name CA or name CB) and resname LEU')
assert_equal(len(sel), 32)
# TODO:
# test for checking ordering and multiple comma-separated selections
def test_concatenated_selection(self, universe):
E151 = universe.select_atoms('segid 4AKE').select_atoms('resid 151')
# note that this is not quite phi... HN should be C of prec. residue
phi151 = E151.atoms.select_atoms('name HN', 'name N', 'name CA',
'name CB')
assert_equal(len(phi151), 4)
assert_equal(phi151[0].name, 'HN',
"wrong ordering in selection, should be HN-N-CA-CB")
def test_global(self, universe):
"""Test the `global` modifier keyword (Issue 268)"""
ag = universe.select_atoms("resname LYS and name NZ")
# Lys amines within 4 angstrom of the backbone.
ag1 = universe.select_atoms(
"resname LYS and name NZ and around 4 backbone")
ag2 = ag.select_atoms("around 4 global backbone")
assert_equal(ag2.indices, ag1.indices)
@pytest.mark.parametrize('selstring, wildstring', [
('resname TYR THR', 'resname T*R'),
('resname ASN GLN', 'resname *N'),
('resname ASN ASP', 'resname AS*'),
('resname TYR THR', 'resname T?R'),
('resname ASN ASP HSD', 'resname *S?'),
('resname LEU LYS', 'resname L**'),
('resname MET', 'resname *M*'),
('resname GLN GLY', 'resname GL[NY]'),
('resname GLU', 'resname GL[!NY]'),
])
def test_wildcard_selection(self, universe, selstring, wildstring):
ag = universe.select_atoms(selstring)
ag_wild = universe.select_atoms(wildstring)
assert ag == ag_wild
class TestSelectionsAMBER(object):
@pytest.fixture()
def universe(self):
return MDAnalysis.Universe(PRMpbc, TRJpbc_bz2)
def test_protein(self, universe):
sel = universe.select_atoms('protein')
assert_equal(sel.n_atoms, 22, "failed to select protein")
def test_backbone(self, universe):
sel = universe.select_atoms('backbone')
assert_equal(sel.n_atoms, 7)
def test_type(self, universe):
sel = universe.select_atoms('type HC')
assert_equal(len(sel), 6)
assert_equal(sel.names, ['HH31', 'HH32', 'HH33', 'HB1', 'HB2', 'HB3'])
class TestSelectionsNAMD(object):
@pytest.fixture()
def universe(self):
return MDAnalysis.Universe(PSF_NAMD, PDB_NAMD)
def test_protein(self, universe):
# must include non-standard residues
sel = universe.select_atoms(
'protein or resname HAO or resname ORT')
assert_equal(sel.n_atoms, universe.atoms.n_atoms,
"failed to select peptide")
assert_equal(sel.n_residues, 6,
"failed to select all peptide residues")
def test_resid_single(self, universe):
sel = universe.select_atoms('resid 12')
assert_equal(sel.n_atoms, 26)
assert_equal(sel.residues.resnames, ['HAO'])
def test_type(self, universe):
sel = universe.select_atoms('type H')
assert_equal(len(sel), 5)
# note 4th HH
assert_equal(sel.names, ['HN', 'HN', 'HN', 'HH', 'HN'])
class TestSelectionsGRO(object):
@pytest.fixture()
def universe(self):
return MDAnalysis.Universe(GRO)
def test_protein(self, universe):
sel = universe.select_atoms('protein')
assert_equal(sel.n_atoms, 3341, "failed to select protein")
def test_backbone(self, universe):
sel = universe.select_atoms('backbone')
assert_equal(sel.n_atoms, 855)
def test_resid_single(self, universe):
sel = universe.select_atoms('resid 100')
assert_equal(sel.n_atoms, 7)
assert_equal(sel.residues.resnames, ['GLY'])
@pytest.mark.parametrize('selstr', [
'same x as bynum 1 or bynum 10',
'same x as bynum 1 10',
'same x as index 0 or index 9',
'same x as index 0 9'
])
def test_same_coordinate(self, universe, selstr):
"""Test the 'same ... as' construct (Issue 217)"""
sel = universe.select_atoms(selstr)
errmsg = ("Found a wrong number of atoms with same "
"x as ids 1 or 10")
assert_equal(len(sel), 12, errmsg)
target_ids = np.array([0, 8, 9, 224, 643, 3515, 11210, 14121,
18430, 25418, 35811, 43618])
assert_equal(sel.indices, target_ids,
"Found wrong atoms with same x as ids 1 or 10")
@pytest.mark.parametrize('selstr', [
'cylayer 10 20 20 -20 bynum 3554',
'cylayer 10 20 20 -20 index 3553'
])
def test_cylayer(self, universe, selstr):
"""Cylinder layer selections with tricilinic periodicity (Issue 274)"""
atgp = universe.select_atoms('name OW')
sel = atgp.select_atoms(selstr)
assert_equal(len(sel), 1155)
@pytest.mark.parametrize('selstr', [
'cyzone 20 20 -20 bynum 3554',
'cyzone 20 20 -20 index 3553'
])
def test_cyzone(self, universe, selstr):
"""Cylinder zone selections with tricilinic periodicity (Issue 274)"""
atgp = universe.select_atoms('name OW')
sel = atgp.select_atoms(selstr)
assert_equal(len(sel), 1556)
class TestSelectionsTPR(object):
@staticmethod
@pytest.fixture(scope='class')
def universe():
return MDAnalysis.Universe(TPR,XTC)
@pytest.mark.parametrize('selstr', [
'same fragment as bynum 1',
'same fragment as index 0'
])
def test_same_fragment(self, universe, selstr):
"""Test the 'same ... as' construct (Issue 217)"""
# This test comes here because it's a system with solvent,
# and thus multiple fragments.
sel = universe.select_atoms(selstr)
errmsg = ("Found a wrong number of atoms "
"on the same fragment as id 1")
assert_equal(len(sel), 3341, errmsg)
errmsg = ("Found a differ set of atoms when using the 'same "
"fragment as' construct vs. the .fragment property")
assert_equal(sel.indices, universe.atoms[0].fragment.indices, errmsg)
def test_moltype(self, universe):
sel = universe.select_atoms("moltype NA+")
ref = np.array([47677, 47678, 47679, 47680], dtype=np.int32)
assert_equal(sel.ids, ref)
@pytest.mark.parametrize(
'selection_string,reference',
(('molnum 1', [3341, 3342, 3343, 3344]),
('molnum 2:4', [3345, 3346, 3347, 3348, 3349, 3350,
3351, 3352, 3353, 3354, 3355, 3356]),
)
)
def test_molnum(self, universe, selection_string, reference):
sel = universe.select_atoms(selection_string)
assert_equal(sel.ids, np.array(reference, dtype=np.int32))
class TestSelectionRDKit(object):
def setup_class(self):
pytest.importorskip("rdkit.Chem")
@pytest.fixture
def u(self):
smi = "Cc1cNcc1"
u = MDAnalysis.Universe.from_smiles(smi, addHs=False,
generate_coordinates=False)
return u
@pytest.mark.parametrize("sel_str, n_atoms", [
("aromatic", 5),
("not aromatic", 1),
("type N and aromatic", 1),
("type C and aromatic", 4),
])
def test_selection(self, u, sel_str, n_atoms):
sel = u.select_atoms(sel_str)
assert sel.n_atoms == n_atoms
class TestSelectionsNucleicAcids(object):
@pytest.fixture(scope='class')
def universe(self):
return MDAnalysis.Universe(RNA_PSF)
@pytest.fixture(scope='class')
def universe2(self):
return MDAnalysis.Universe(NUCLsel)
def test_nucleic(self, universe):
rna = universe.select_atoms("nucleic")
assert_equal(rna.n_atoms, 739)
assert_equal(rna.n_residues, 23)
def test_nucleic_all(self, universe2):
sel = universe2.select_atoms('nucleic')
assert len(sel) == 34
def test_nucleicbackbone(self, universe):
rna = universe.select_atoms("nucleicbackbone")
assert_equal(rna.n_residues, 23)
assert_equal(rna.n_atoms, rna.n_residues * 5 - 1)
# -1 because this is a single strand of RNA and on P is missing at the 5' end
# todo: need checks for other selection resnames such as DT DA DG DC DU
def test_nucleicbase(self, universe):
rna = universe.select_atoms("nucleicbase")
assert_equal(rna.n_residues, 23)
assert_equal(rna.n_atoms, 214)
def test_nucleicsugar(self, universe):
rna = universe.select_atoms("nucleicsugar")
assert_equal(rna.n_residues, 23)
assert_equal(rna.n_atoms, rna.n_residues * 5)
class BaseDistanceSelection(object):
"""Both KDTree and distmat selections on orthogonal system
Selections to check:
- Around
- SphericalLayer
- SphericalZone
- Point
Cylindrical methods don't use KDTree
"""
@pytest.mark.parametrize('periodic', (True, False))
def test_around(self, u, periodic):
sel = Parser.parse('around 5.0 resid 1', u.atoms)
if periodic:
sel.periodic = True
else:
sel.periodic = False
result = sel.apply(u.atoms)
r1 = u.select_atoms('resid 1')
cog = r1.center_of_geometry().reshape(1, 3)
box = u.dimensions if periodic else None
d = distance_array(u.atoms.positions, r1.positions,
box=box)
ref = set(np.where(d < 5.0)[0])
# Around doesn't include atoms from the reference group
ref.difference_update(set(r1.indices))
assert ref == set(result.indices)
@pytest.mark.parametrize('periodic', (True, False))
def test_spherical_layer(self, u, periodic):
sel = Parser.parse('sphlayer 2.4 6.0 resid 1', u.atoms)
if periodic:
sel.periodic = True
else:
sel.periodic = False
result = sel.apply(u.atoms)
r1 = u.select_atoms('resid 1')
box = u.dimensions if periodic else None
cog = r1.center_of_geometry().reshape(1, 3)
d = distance_array(u.atoms.positions, cog, box=box)
ref = set(np.where((d > 2.4) & (d < 6.0))[0])
assert ref == set(result.indices)
@pytest.mark.parametrize('periodic', (True, False))
def test_spherical_zone(self, u, periodic):
sel = Parser.parse('sphzone 5.0 resid 1', u.atoms)
if periodic:
sel.periodic = True
else:
sel.periodic = False
result = sel.apply(u.atoms)
r1 = u.select_atoms('resid 1')
box = u.dimensions if periodic else None
cog = r1.center_of_geometry().reshape(1, 3)
d = distance_array(u.atoms.positions, cog, box=box)
ref = set(np.where(d < 5.0)[0])
assert ref == set(result.indices)
@pytest.mark.parametrize('periodic', (True, False))
def test_point(self, u, periodic):
sel = Parser.parse('point 5.0 5.0 5.0 3.0', u.atoms)
if periodic:
sel.periodic = True
else:
sel.periodic = False
result = sel.apply(u.atoms)
box = u.dimensions if periodic else None
d = distance_array(np.array([[5.0, 5.0, 5.0]], dtype=np.float32),
u.atoms.positions,
box=box)
ref = set(np.where(d < 3.0)[1])
assert ref == set(result.indices)
class TestOrthogonalDistanceSelections(BaseDistanceSelection):
@pytest.fixture()
def u(self):
return mda.Universe(TRZ_psf, TRZ)
@pytest.mark.parametrize('meth, periodic', [
('distmat', True),
('distmat', False)
])
def test_cyzone(self, u, meth, periodic):
sel = Parser.parse('cyzone 5 4 -4 resid 2', u.atoms)
sel.periodic = periodic
result = sel.apply(u.atoms)
other = u.select_atoms('resid 2')
pos = other.center_of_geometry()
vecs = u.atoms.positions - pos
if periodic:
box = u.dimensions[:3]
vecs -= box * np.rint(vecs / box)
mask = (vecs[:, 2] > -4) & (vecs[:, 2] < 4)
radii = vecs[:, 0] ** 2 + vecs[:, 1] ** 2
mask &= radii < 5 ** 2
ref = set(u.atoms[mask].indices)
assert ref == set(result.indices)
@pytest.mark.parametrize('periodic,expected', ([True, 33], [False, 25]))
def test_sphzone(self, u, periodic, expected):
sel = u.select_atoms('sphzone 5.0 resid 1', periodic=periodic)
assert len(sel) == expected
class TestTriclinicDistanceSelections(BaseDistanceSelection):
@pytest.fixture()
def u(self):
return mda.Universe(GRO)
class TestTriclinicSelections(object):
"""Non-KDTree based selections
This system has triclinic geometry so won't use KDTree based selections
"""
@pytest.fixture()
def u(self):
return mda.Universe(GRO)
def test_around(self, u):
r1 = u.select_atoms('resid 1')
ag = u.select_atoms('around 5.0 resid 1')
d = distance_array(u.atoms.positions, r1.positions, box=u.dimensions)
idx = set(np.where(d < 5.0)[0])
# Around doesn't include atoms from the reference group
idx.difference_update(set(r1.indices))
assert idx == set(ag.indices)
def test_sphlayer(self, u):
r1 = u.select_atoms('resid 1')
cog = r1.center_of_geometry().reshape(1, 3)
ag = u.select_atoms('sphlayer 2.4 6.0 resid 1')
d = distance_array(u.atoms.positions, cog, box=u.dimensions)
idx = set(np.where((d > 2.4) & (d < 6.0))[0])
assert idx == set(ag.indices)
def test_sphzone(self, u):
r1 = u.select_atoms('resid 1')
cog = r1.center_of_geometry().reshape(1, 3)
ag = u.select_atoms('sphzone 5.0 resid 1')
d = distance_array(u.atoms.positions, cog, box=u.dimensions)
idx = set(np.where(d < 5.0)[0])
assert idx == set(ag.indices)
def test_point_1(self, u):
# The example selection
ag = u.select_atoms('point 5.0 5.0 5.0 3.5')
d = distance_array(np.array([[5.0, 5.0, 5.0]], dtype=np.float32),
u.atoms.positions,
box=u.dimensions)
idx = np.where(d < 3.5)[1]
assert_equal(set(ag.indices), set(idx))
def test_point_2(self, u):
ag1 = u.atoms[:10000]
ag2 = ag1.select_atoms('point 5.0 5.0 5.0 3.5')
d = distance_array(np.array([[5.0, 5.0, 5.0]], dtype=np.float32),
ag1.positions,
box=u.dimensions)
idx = np.where(d < 3.5)[1]
assert_equal(set(ag2.indices), set(idx))
def gen_sel_strings(prop, oper):
"""Generate all possible combinations of spaces in selection strings
ie:
'prop x < 1.5'
'prop x< 1.5'
'prop x <1.5'
'prop x<1.5'
"""
for x, y in itertools.product([' ', ''], [' ', '']):
yield 'prop {prop}{spc1}{oper}{spc2}1.5'.format(
prop=prop, spc1=x, oper=oper, spc2=y)
class TestPropSelection(object):
plurals = {'mass': 'masses',
'charge': 'charges'}
op_funcs = {
'<': np.less,
'<=': np.less_equal,
'>': np.greater,
'>=': np.greater_equal,
'==': np.equal,
'!=': np.not_equal
}
opposites = {
'==': '==', '!=': '!=',
'>': '<=', '<=': '>',
'<': '>=', '>=': '<',
}
@pytest.fixture(params=[slice(None, None), slice(None, 100)])
def ag(self, request):
u = make_Universe(('masses', 'charges'))
u.atoms[::2].masses = 1.5
u.atoms[::2].charges = 1.5
return u.atoms[request.param]
@pytest.mark.parametrize('prop, selstr', [
(prop, sel)
for prop in ['mass', 'charge']
for sel in gen_sel_strings(prop, '<')
])
def test_lt(self, prop, selstr, ag):
sel = ag.select_atoms(selstr)
assert_equal(set(sel.indices),
set(ag[getattr(ag, self.plurals[prop]) < 1.5].indices))
@pytest.mark.parametrize('prop, selstr', [
(prop, sel)
for prop in ['mass', 'charge']
for sel in gen_sel_strings(prop, '<=')
])
def test_le(self, prop, selstr, ag):
sel = ag.select_atoms(selstr)
assert_equal(set(sel.indices),
set(ag[getattr(ag,
self.plurals[prop]) <= 1.5].indices))
@pytest.mark.parametrize('prop, selstr', [
(prop, sel)
for prop in ['mass', 'charge']
for sel in gen_sel_strings(prop, '>')
])
def test_gt(self, prop, selstr, ag):
sel = ag.select_atoms(selstr)
assert_equal(set(sel.indices),
set(ag[getattr(ag, self.plurals[prop]) > 1.5].indices))
@pytest.mark.parametrize('prop, selstr', [
(prop, sel)
for prop in ['mass', 'charge']
for sel in gen_sel_strings(prop, '>=')
])
def test_ge(self, prop, selstr, ag):
sel = ag.select_atoms(selstr)
assert_equal(set(sel.indices),
set(ag[getattr(ag,
self.plurals[prop]) >= 1.5].indices))
@pytest.mark.parametrize('prop, selstr', [
(prop, sel)
for prop in ['mass', 'charge']
for sel in gen_sel_strings(prop, '==')
])
def test_eq(self, prop, selstr, ag):
sel = ag.select_atoms(selstr)
assert_equal(set(sel.indices),
set(ag[getattr(ag,
self.plurals[prop]) == 1.5].indices))
@pytest.mark.parametrize('prop, selstr', [
(prop, sel)
for prop in ['mass', 'charge']
for sel in gen_sel_strings(prop, '!=')
])
def test_ne(self, prop, selstr, ag):
sel = ag.select_atoms(selstr)
assert_equal(set(sel.indices),
set(ag[getattr(ag,
self.plurals[prop]) != 1.5].indices))
@pytest.mark.parametrize('prop, op', [
(prop, op)
for prop in ['mass', 'charge']
for op in ('<', '>', '<=', '>=', '==', '!=')
])
def test_flip(self, prop, ag, op):
func = self.op_funcs[op]
# reference group, doing things forwards
ref = ag[func(getattr(ag, self.plurals[prop]), 1.5)]
selstr = 'prop 1.5 {op} {prop}'.format(
op=self.opposites[op], prop=prop)
sel = ag.select_atoms(selstr)
assert_equal(set(ref.indices), set(sel.indices))
class TestBondedSelection(object):
@pytest.fixture()
def u(self):
return mda.Universe(PSF, DCD)
def test_bonded_1(self, u):
ag = u.select_atoms('type 2 and bonded name N')
assert len(ag) == 3
def test_nobonds_warns(self, u):
u = make_Universe(('names',))
# empty bond topology attr
batt = mda.core.topologyattrs.Bonds([])
u.add_TopologyAttr(batt)
with pytest.warns(UserWarning):
u.select_atoms('bonded name AAA')
class TestSelectionErrors(object):
@staticmethod
@pytest.fixture()
def universe():
return make_Universe(
('names', 'masses', 'resids', 'resnames', 'resnums'))
@pytest.mark.parametrize('selstr', [
'name and H', # string selection
'name )',
'resid abcd', # resid arg parsing selection
'resnum 7a7', # rangeselection arg parsing
'resid 1-',
'prop chicken == tasty',
'prop chicken <= 7.4',
'prop mass ^^ 12.0',
'same this as resid 1', # same selection
'same resid resname mass 5.0', # same / expect
'name H and', # check all tokens used
'naem H', # unkonwn (misplet) opertaor
'resid and name C', # rangesel not finding vals
'resnum ',
'bynum or protein',
'index or protein',
'prop mass < 4.0 hello', # unused token
'prop mass > 10. and group this', # missing group
])
def test_selection_fail(self, selstr, universe):
with pytest.raises(SelectionError):
universe.select_atoms(selstr)
def test_segid_and_resid():
u = make_Universe(('segids', 'resids'))
ag = u.select_atoms('segid SegB and resid 1-100')
ref = ag.select_atoms('segid SegB').select_atoms('resid 1-100')
assert_equal(ag.indices, ref.indices)
class TestImplicitOr(object):
@staticmethod
@pytest.fixture()
def universe():
return make_Universe(
('names', 'types', 'resids', 'resnums', 'resnames', 'segids'))
def _check_sels(self, ref, sel, universe):
ref = universe.select_atoms(ref)
sel = universe.select_atoms(sel)
assert_equal(ref.indices, sel.indices)
@pytest.mark.parametrize('ref, sel', [
('name NameABA or name NameACA or name NameADA',
'name NameABA NameACA NameADA'),
('type TypeE or type TypeD or type TypeB', 'type TypeE TypeD TypeB'),
('resname RsC or resname RsY', 'resname RsC RsY'),
('name NameAB* or name NameACC', 'name NameAB* NameACC'),
('segid SegA or segid SegC', 'segid SegA SegC'),
('(name NameABC or name NameABB) and (resname RsD or resname RsF)',
'name NameABC NameABB and resname RsD RsF'),
])
def test_string_selections(self, ref, sel, universe):
self._check_sels(ref, sel, universe)
@pytest.mark.parametrize("seltype", ['resid', 'resnum', 'bynum', 'index'])
@pytest.mark.parametrize('ref, sel', [
('{typ} 1 or {typ} 2', '{typ} 1 2'),
('{typ} 1:10 or {typ} 22', '{typ} 1:10 22'),
('{typ} 1:10 or {typ} 20:30', '{typ} 1:10 20:30'),
('{typ} 1-5 or {typ} 7', '{typ} 1-5 7'),
('{typ} 1-5 or {typ} 7:10 or {typ} 12', '{typ} 1-5 7:10 12'),
('{typ} 1 or {typ} 3 or {typ} 5:10', '{typ} 1 3 5:10'),
])
def test_range_selections(self, seltype, ref, sel, universe):
self._check_sels(ref.format(typ=seltype), sel.format(typ=seltype),