Files
pymol-open-source/testing/tests/api/selecting.py
2023-04-14 14:30:28 -04:00

390 lines
17 KiB
Python

from pymol import cmd, testing, stored
class TestSelecting(testing.PyMOLTestCase):
@testing.requires_version('2.4')
def test_error_propagation(self):
try:
cmd.center("(all) and no_such_name")
except Exception as e:
self.assertTrue("Invalid selection name" in e.message)
self.assertTrue("no_such_name" in e.message)
return
self.fail("did not raise")
def testDeselect(self):
cmd.load(self.datafile("1oky.pdb.gz"), "m1")
cmd.select("all")
self.assertEquals("sele" in cmd.get_names("all",enabled_only=1), True)
cmd.deselect()
self.assertEquals("sele" in cmd.get_names("all",enabled_only=1), False)
def testIndicate(self):
cmd.load(self.datafile("1oky.pdb.gz"), "m1")
cmd.indicate("polymer")
self.assertEquals(cmd.count_atoms("polymer"),cmd.count_atoms("indicate"))
self.assertEquals(cmd.count_atoms("indicate & solvent & organic & inorganic"), 0)
def testPop(self):
cmd.fragment("ala")
cmd.select("src", "ala")
# iterate across the 10 atoms in ala
cnt = 0
while cmd.pop("tmp","src"):
cnt+=1
self.assertEquals(cnt,10)
def testSelect(self):
cmd.fragment("gly", "m1")
NC = 2
# auto_number_selections=0
cmd.select("elem C")
self.assertEquals(NC, cmd.count_atoms("sele"))
self.assertEquals(0, cmd.get_setting_int("sel_counter"))
# auto_number_selections=1
cmd.set('auto_number_selections', 1)
cmd.set('sel_counter', 3)
cmd.select("elem C")
self.assertEquals(NC, cmd.count_atoms("sel04"))
self.assertEquals(4, cmd.get_setting_int("sel_counter"))
cmd.set('auto_number_selections', 1)
cmd.select(None, "elem C")
self.assertEquals(NC, cmd.count_atoms("sel05"))
self.assertEquals(5, cmd.get_setting_int("sel_counter"))
# name=None always numbers the selection
cmd.set('auto_number_selections', 0)
cmd.select(None, "elem C")
self.assertEquals(NC, cmd.count_atoms("sel06"))
self.assertEquals(6, cmd.get_setting_int("sel_counter"))
# default
cmd.select("foo", "elem C")
self.assertEquals(NC, cmd.count_atoms("foo"))
self.assertEquals(["foo"], cmd.get_names("selections", enabled_only=1))
# merge with non-existing, enable=0
cmd.delete('foo')
cmd.select("foo", "elem C", 0, merge=1)
self.assertEquals(NC, cmd.count_atoms("foo"))
self.assertEquals([], cmd.get_names("selections", enabled_only=1))
# merge, enable=1
cmd.select("foo", "elem N", 1)
self.assertEquals(1, cmd.count_atoms("foo"))
self.assertEquals(["foo"], cmd.get_names("selections", enabled_only=1))
cmd.select("foo", "elem C", -1, merge=1)
self.assertEquals(NC + 1, cmd.count_atoms("foo"))
self.assertEquals(["foo"], cmd.get_names("selections", enabled_only=1))
cmd.select("foo", "elem O", -1, merge=2)
self.assertEquals(NC + 2, cmd.count_atoms("foo"))
self.assertEquals(["foo"], cmd.get_names("selections", enabled_only=1))
# merge, enable=0
cmd.select("foo", "elem N", 0)
self.assertEquals(1, cmd.count_atoms("foo"))
self.assertEquals([], cmd.get_names("selections", enabled_only=1))
cmd.select("foo", "elem C", -1, merge=1)
self.assertEquals(NC + 1, cmd.count_atoms("foo"))
self.assertEquals([], cmd.get_names("selections", enabled_only=1))
cmd.select("foo", "elem O", -1, merge=2)
self.assertEquals(1, cmd.count_atoms("foo"))
self.assertEquals([], cmd.get_names("selections", enabled_only=1))
# state
cmd.create('m1', 'm1 & elem C', 1, 2)
self.assertEquals(7, cmd.select('present'))
self.assertEquals(7, cmd.select('present', state=1))
self.assertEquals(2, cmd.select('present', state=2))
self.assertEquals(0, cmd.select('present', state=3))
# domain
cmd.select("foo", "elem C")
cmd.select("bar", "name CA+N+O", domain="foo")
self.assertEquals(1, cmd.count_atoms("bar"))
def testSelectList(self):
cmd.fragment('ala', 'm1')
cmd.alter('m1', 'rank += 10')
cmd.alter('m1', 'ID += 20')
cmd.select_list('s1', 'm1', [1, 4], mode='index')
self.assertEqual(cmd.count_atoms('s1'), 2)
self.assertEqual(cmd.count_atoms('s1 & elem N'), 1)
self.assertEqual(cmd.count_atoms('s1 & elem O'), 1)
cmd.select_list('s1', 'm1', [15, 16, 17, 18, 19], mode='rank')
self.assertEqual(cmd.count_atoms('s1'), 5)
self.assertEqual(cmd.count_atoms('s1 & elem H'), 5)
cmd.select_list('s1', 'm1', [21, 22, 23], mode='id')
self.assertEqual(cmd.count_atoms('s1'), 3)
self.assertEqual(cmd.count_atoms('s1 & elem C'), 3)
# empty list (use new selection names to verify that named selection
# was crated)
cmd.select_list('s2', 'm1', [], mode='index')
self.assertEqual(cmd.count_atoms('s2'), 0)
cmd.select_list('s3', 'm1', [], mode='rank')
self.assertEqual(cmd.count_atoms('s3'), 0)
cmd.select_list('s4', 'm1', [], mode='id')
self.assertEqual(cmd.count_atoms('s4'), 0)
@testing.requires_version('2.4')
def testSelectList_state(self):
cmd.fragment('ala', 'm1')
cmd.create('m1', 'm1 & elem C', 1, 2)
cmd.select_list('s1', 'm1', [100] * 100 + [2, 3, 4], state=1, mode='rank')
self.assertEqual(cmd.count_atoms('s1'), 3)
cmd.select_list('s1', 'm1', [100] * 100 + [2, 3, 4], state=2, mode='rank')
self.assertEqual(cmd.count_atoms('s1'), 2)
# https://github.com/schrodinger/pymol-open-source/issues/109
@testing.foreach('index', 'rank', 'id')
@testing.requires_version('2.5')
def testSelectList_unmatched(self, mode):
cmd.fragment('ala', 'm1')
cmd.select_list('s1', 'm1', [100], mode=mode)
self.assertEqual(cmd.count_atoms('s1'), 0)
def testMacros(self):
'''
Test selection macros: /model/segi/chain/resn`resi/name`alt
See also tests/jira/PYMOL-2720.py
'''
cmd.load(self.datafile("1oky.pdb.gz"), "m1")
cmd.alter('*', 'segi="X"')
cmd.alter('resi 200-238', 'chain="B"')
cmd.alter('resi 239-360', 'chain="C"')
cmd.alter('resi 150-220', 'segi="Y"')
cmd.alter('resi 220-', 'segi="Z"')
# all atoms
self.assertEqual(cmd.count_atoms('/'), cmd.count_atoms())
self.assertEqual(cmd.count_atoms('/////'), cmd.count_atoms())
# model
self.assertEqual(cmd.count_atoms('/m1'), cmd.count_atoms('all'))
self.assertEqual(cmd.count_atoms('m1////'), cmd.count_atoms('all'))
# segi
self.assertEqual(cmd.count_atoms('//Y'), cmd.count_atoms('segi Y'))
self.assertEqual(cmd.count_atoms('//Y+Z'), cmd.count_atoms('segi Y+Z'))
# chains
self.assertEqual(cmd.count_atoms('///B'), cmd.count_atoms('chain B'))
self.assertEqual(cmd.count_atoms('///A+B'), cmd.count_atoms('chain A+B'))
# resn/resi
self.assertEqual(cmd.count_atoms('////100'), cmd.count_atoms('resi 100'))
self.assertEqual(cmd.count_atoms('////`100'), cmd.count_atoms('resi 100'))
self.assertEqual(cmd.count_atoms('////100-110'), cmd.count_atoms('resi 100-110'))
self.assertEqual(cmd.count_atoms('////ARG`100'), cmd.count_atoms('resi 100'))
self.assertEqual(cmd.count_atoms('////ALA'), cmd.count_atoms('resn ALA'), )
self.assertEqual(cmd.count_atoms('ALA/'), cmd.count_atoms('resn ALA'), )
self.assertEqual(cmd.count_atoms('////ALA`'), cmd.count_atoms('resn ALA'))
self.assertEqual(cmd.count_atoms('////ALA+GLU'), cmd.count_atoms('resn ALA+GLU'), )
# name/alt
self.assertEqual(cmd.count_atoms('/////CG'), cmd.count_atoms('name CG'))
self.assertEqual(cmd.count_atoms('*/CG'), cmd.count_atoms('name CG'))
self.assertEqual(cmd.count_atoms('/////CG`B'), cmd.count_atoms('name CG & alt B'))
self.assertEqual(cmd.count_atoms('/////`B'), cmd.count_atoms('alt B'))
# combies
self.assertEqual(cmd.count_atoms('100/CA'), cmd.count_atoms('resi 100 & name CA'))
self.assertEqual(cmd.count_atoms('A//CA'), cmd.count_atoms('chain A & name CA'))
self.assertEqual(cmd.count_atoms('A//`B'), cmd.count_atoms('chain A & alt B'))
self.assertEqual(cmd.count_atoms('/m1/Y/B/*/CA'), cmd.count_atoms('segi Y & chain B & name CA'))
# or-logic
ref = cmd.count_atoms('(resi 100 & name CA) (resi 110 & name CB)')
self.assertEqual(cmd.count_atoms('100/CA|110/CB'), ref)
self.assertEqual(cmd.count_atoms('100/CA or 110/CB'), ref)
self.assertEqual(cmd.count_atoms('100/CA 110/CB'), ref)
@testing.requires_version('1.8.7')
def test_not_enabled(self):
cmd.fragment('ala')
cmd.fragment('arg')
cmd.fragment('asp')
cmd.fragment('gly')
cmd.fragment('glu')
cmd.delete('not a*') # use "not " prefix
self.assertEqual(cmd.get_names(), ['ala', 'arg', 'asp'])
cmd.disable('arg')
cmd.delete('!enabled') # use "!" prefix
self.assertEqual(cmd.get_names(), ['ala', 'asp'])
cmd.disable('asp')
cmd.delete('enabled')
self.assertEqual(cmd.get_names(), ['asp'])
@testing.requires_version('2.1')
def test_protein_nucleic(self):
npro = 340
nnuc = 112
cmd.load(self.datafile("1oky-frag.pdb"), "1p")
cmd.load(self.datafile('1ehz-5.pdb'), "1n")
self.assertEqual(cmd.count_atoms('polymer'), npro + nnuc)
self.assertEqual(cmd.count_atoms('polymer.protein'), npro)
self.assertEqual(cmd.count_atoms('polymer.nucleic'), nnuc)
@testing.requires_version('2.2')
def test_wildcards(self):
cmd.pseudoatom('foo')
cmd.pseudoatom('foobar')
cmd.pseudoatom('foo_bar')
cmd.pseudoatom('foo_bar_bla')
cmd.pseudoatom('bar')
cmd.pseudoatom('_bar')
self.assertEqual(1, cmd.count_atoms('foo'))
self.assertEqual(1, cmd.count_atoms('*foo'))
self.assertEqual(4, cmd.count_atoms('*foo*'))
self.assertEqual(1, cmd.count_atoms('bar'))
self.assertEqual(4, cmd.count_atoms('*bar'))
self.assertEqual(5, cmd.count_atoms('*bar*'))
self.assertEqual(2, cmd.count_atoms('foo*bar'))
self.assertEqual(3, cmd.count_atoms('*foo*bar*'))
self.assertEqual(1, cmd.count_atoms('_*'))
self.assertEqual(6, cmd.count_atoms('*'))
self.assertEqual(6, cmd.count_atoms('*'))
cmd.alter('*', 'name = "ABCD"')
cmd.alter('foo', 'name = "ABED"')
cmd.alter('bar', 'name = "ABEC"')
self.assertEqual(6, cmd.count_atoms('name A*'))
self.assertEqual(2, cmd.count_atoms('name ABE*'))
self.assertEqual(1, cmd.count_atoms('name ABED*'))
self.assertEqual(1, cmd.count_atoms('name AB*C'))
self.assertEqual(5, cmd.count_atoms('name AB*C*'))
self.assertEqual(5, cmd.count_atoms('name AB*D'))
self.assertEqual(5, cmd.count_atoms('name AB*D*'))
@testing.requires_version('2.2')
def test_wildcard_sets_ranges(self):
for i in range(10): cmd.pseudoatom('m%d' % i, chain=chr(65 + i))
cmd.alter('m5', 'chain = "AB"')
cmd.alter('m6', 'chain = "BA"')
cmd.alter('m7', 'chain = "CC"')
cmd.alter('m8', 'chain = "ZA"')
cmd.alter('m9', 'chain = "ABA"')
# A patterns
self.assertEqual(1, cmd.count_atoms('chain A'))
self.assertEqual(3, cmd.count_atoms('chain A*'))
self.assertEqual(4, cmd.count_atoms('chain *A'))
self.assertEqual(5, cmd.count_atoms('chain *A*'))
self.assertEqual(1, cmd.count_atoms('chain A*A'))
# B patterns
self.assertEqual(2, cmd.count_atoms('chain B*'))
self.assertEqual(2, cmd.count_atoms('chain *B'))
self.assertEqual(4, cmd.count_atoms('chain *B*'))
# X patterns (no matches)
self.assertEqual(0, cmd.count_atoms('chain X*'))
self.assertEqual(0, cmd.count_atoms('chain *X'))
self.assertEqual(0, cmd.count_atoms('chain *X*'))
# list with wildcards
self.assertEqual(5, cmd.count_atoms('chain B*+A*'))
self.assertEqual(3, cmd.count_atoms('chain B*+A*A'))
self.assertEqual(3, cmd.count_atoms('chain B*+A*A+*X'))
# lexicographical alpha ranges, A:C, will match AB (A <= AB <= C) but not CC (C < CC)
# no wildcards in alpha ranges possible
self.assertEqual(6, cmd.count_atoms('chain A:C'))
self.assertEqual(6, cmd.count_atoms('chain A:CA'))
self.assertEqual(7, cmd.count_atoms('chain A:CX')) # include CC
self.assertEqual(6, cmd.count_atoms('chain A:C+Z'))
self.assertEqual(7, cmd.count_atoms('chain A:C+Z*'))
def test_sets_ranges(self):
cmd.fab('ACDEFGHIKL')
cmd.alter('resi 9', 'resi="9A"') # insertion code
self.assertEqual(3, cmd.count_atoms('guide & resi 2-4'))
self.assertEqual(3, cmd.count_atoms('guide & resi 2:4'))
self.assertEqual(2, cmd.count_atoms('guide & resi 2+4'))
self.assertEqual(4, cmd.count_atoms('guide & resi 2-4+6'))
self.assertEqual(6, cmd.count_atoms('guide & resi 2-4+6-8'))
self.assertEqual(0, cmd.count_atoms('guide & resi 9'))
self.assertEqual(1, cmd.count_atoms('guide & resi 9A'))
self.assertEqual(1, cmd.count_atoms('guide & resi 10'))
self.assertEqual(0, cmd.count_atoms('guide & resi 10A'))
self.assertEqual(2, cmd.count_atoms('guide & resi 9-10'))
self.assertEqual(2, cmd.count_atoms('guide & resi 9A-10A'))
self.assertEqual(10 + 9, cmd.count_atoms('name CA+CB'))
self.assertEqual(10 + 9, cmd.count_atoms('name CA+CB+XYZ'))
self.assertEqual(10, cmd.count_atoms('name C'))
self.assertEqual(50, cmd.count_atoms('name C*'))
self.assertEqual(10, cmd.count_atoms('name H'))
self.assertEqual(24, cmd.count_atoms('name H*'))
self.assertEqual(10, cmd.count_atoms('name *H'))
self.assertEqual(74, cmd.count_atoms('name *H*'))
self.assertEqual(20, cmd.count_atoms('name C+N'))
self.assertEqual(23, cmd.count_atoms('name C+N*'))
self.assertEqual(60, cmd.count_atoms('name C*+N'))
self.assertEqual(63, cmd.count_atoms('name C*+N*'))
def test_guide(self):
cmd.fragment("gly")
self.assertEqual(0, cmd.count_atoms("organic"))
self.assertEqual(1, cmd.count_atoms("guide"))
self.assertEqual(1, cmd.count_atoms("organic | guide"))
@testing.foreach('origin', 'center')
def test_dummy_selectors(self, op):
self.assertEqual(cmd.count_atoms('all'), 0)
self.assertEqual(cmd.count_atoms(op), 0)
self.assertEqual(cmd.count_atoms(op + ' around 2'), 0)
cmd.fragment('arg')
self.assertEqual(cmd.count_atoms('all'), 24)
self.assertEqual(cmd.count_atoms(op), 0)
self.assertEqual(cmd.count_atoms(op + ' around 2'), 8)
self.assertEqual(cmd.count_atoms(op + ' expand 2'), 8)
self.assertEqual(cmd.count_atoms(op + ' within 2 of arg'), 0)
self.assertEqual(cmd.count_atoms('arg within 2 of ' + op), 8)
self.assertEqual(cmd.count_atoms('arg near_to 2 of ' + op), 8)
self.assertEqual(cmd.count_atoms('arg beyond 3 of ' + op), 13)
self.assertEqual(cmd.count_atoms(op + ' gap 3'), 5)
def test_offcenter_origin_selector(self):
cmd.fragment('arg')
cmd.origin('elem O')
self.assertEqual(cmd.count_atoms('origin around 2'), 2)
self.assertEqual(cmd.count_atoms('origin expand 2'), 2)
self.assertEqual(cmd.count_atoms('arg within 2 of origin'), 2)
self.assertEqual(cmd.count_atoms('arg near_to 2 of origin'), 2)
self.assertEqual(cmd.count_atoms('arg beyond 3 of origin'), 19)
self.assertEqual(cmd.count_atoms('origin gap 3'), 14)
# TODO Also test 'center' selector here. I'd expect that it gives the
# same result as the test_dummy_selectors test, but numbers are
# different!
# don't select center/origin with "all" keyword
def test_no_all_dummy_selection(self):
cmd.pseudoatom('p1', pos=(-1, 0, 0))
cmd.pseudoatom('p2', pos=(1, 0, 0))
self.assertEqual(cmd.count_atoms('center around 1.5'), 2)
self.assertEqual(cmd.count_atoms('(p1|center) around 1.5'), 1)
self.assertEqual(cmd.count_atoms('all within 1.5 of center'), 2)
self.assertEqual(cmd.count_atoms('all within 1.5 of (all within 1.5 of p1)'), 1)
self.assertEqual(cmd.count_atoms('all within 1.5 of (center within 1.5 of p1)'), 2)
self.assertEqual(cmd.count_atoms('all near_to 1.5 of center'), 2)
self.assertEqual(cmd.count_atoms('all near_to 1.5 of (all near_to 1.5 of p1)'), 0)
self.assertEqual(cmd.count_atoms('all near_to 1.5 of (center near_to 1.5 of p1)'), 2)
self._test_no_implicit_dummy_selection()
@testing.requires_version('2.5')
def _test_no_implicit_dummy_selection(self):
self.assertEqual(cmd.count_atoms('(p1 around 1.5) around 1.5'), 0)