'''
(c) 2010-2012 Thomas Holder, MPI for Developmental Biology
License: BSD-2-Clause
'''
from pymol import cmd, CmdException
_auto_arg0_align = cmd.auto_arg[0]['align']
_auto_arg1_align = cmd.auto_arg[1]['align']
_auto_arg1_select = cmd.auto_arg[1]['select']
def _assert_package_import():
if not __name__.endswith('.selecting'):
raise CmdException("Must do 'import psico.selecting' instead of 'run ...'")
[docs]
def select_pepseq(pattern, selection='all', name='sele', state=1, quiet=1,
cutoff=4.0, one_letter=None, *, _self=cmd):
'''
DESCRIPTION
Find a amino acid sequence pattern (regular expression) in given atom
selection. Does not span gaps (unless matched by a wildcard).
USAGE
select_pepseq pattern [, selection [, name [, state ]]]
ARGUMENTS
pattern = string: amino acid sequence in one letter code, can be a
regular expression pattern.
selection = string: atom selection of protein (non protein atoms in
selection are silently ignored) {default: all})
name = a unique name for the selection {default: sele}
EXAMPLE
fetch 1a00, async=0
select_pepseq AL[EG]R, chain A+B, sele1
select_pepseq ([LIVAM]{3,}), all, sele2
SEE ALSO
There is the pepseq/ps. selection operator. Example:
select actsite, protein and ps. ADFG
Similar scripts:
http://pldserver1.biochem.queensu.ca/~rlc/work/pymol/seq_select.py
http://pymolwiki.org/index.php/FindSeq
'''
import re
from chempy import cpv
if not one_letter:
_assert_package_import()
from . import one_letter
state, quiet = int(state), int(quiet)
cutoff = float(cutoff)
seq_list = []
idx_list = []
prev = [1e300, 1e300, 1e300]
def callback(model, index, resn, coord):
if cpv.distance(coord, prev) > cutoff:
seq_list.append('#')
idx_list.append(None)
seq_list.append(one_letter.get(resn, '#'))
idx_list.append((model, index))
prev[:] = coord
_self.iterate_state(state, '(%s) and guide' % (selection),
'callback(model, index, resn, (x, y, z))', space=locals())
matches = list(re.finditer(pattern.upper(), ''.join(seq_list)))
if not quiet:
if len(matches) == 0:
print(' select_pepseq: Pattern not found in selection')
else:
print(' select_pepseq: Pattern found %d time(s)' % (len(matches)))
sel_list = []
for m in matches:
start, stop = m.span()
sel_list.extend('%s`%d' % idx for idx in idx_list[start:stop] if idx is not None)
return _self.select(name, '(' + selection + ') and byres (none ' + ' '.join(sel_list) + ')')
[docs]
def select_nucseq(pattern, selection='all', name='sele', state=1, quiet=1, *, _self=cmd):
'''
DESCRIPTION
Find a nucleic acid sequence pattern in given atom selection.
'''
one_letter = dict(["AA", "CC", "TT", "GG", "UU"])
return select_pepseq(pattern, selection, name, state, quiet, 6.5, one_letter, _self=_self)
[docs]
def select_sspick(selection, name=None, caonly=0, quiet=0, *, _self=cmd):
'''
DESCRIPTION
Extend selection by connected secondary structure elements.
Also available as wizard (type: "wizard sspick").
USAGE
select_sspick selection [, name [, caonly [, quiet ]]]
ARGUMENTS
selection = string: selection-expression
name = string: create a named atom selection if not None {default: None}
'''
caonly, quiet = int(caonly), int(quiet)
qkeys = set()
_self.iterate('bycalpha (%s)' % (selection),
'qkeys.add(((model,segi,chain,ss), resv))', space={'qkeys': qkeys})
def in_intervals(i, intervals):
for interval in intervals:
if interval[0] <= i <= interval[1]:
return True
return False
elements = dict()
for key, resv in qkeys:
element = elements.setdefault(key, [])
if in_intervals(resv, element):
continue
resv_set = set()
_self.iterate('/%s/%s/%s//CA and ss "%s"' % key, 'resv_set.add(resv)',
space={'resv_set': resv_set})
resv_min = resv
resv_max = resv
while (resv_min - 1) in resv_set:
resv_min -= 1
while (resv_max + 1) in resv_set:
resv_max += 1
element.append((resv_min, resv_max))
sele_list = []
ss_names = {'S': 'Strand', 'H': 'Helix', '': 'Loop', 'L': 'Loop'}
for key in elements:
model, segi, chain, ss = key
for resv_min, resv_max in elements[key]:
sele = '/%s/%s/%s/%d-%d' % (model, segi, chain, resv_min, resv_max)
if caonly:
sele += '/CA'
sele_list.append(sele)
if not quiet:
print("%-6s %s" % (ss_names.get(ss, ss), sele))
sele = ' '.join(sele_list)
if name is not None:
_self.select(name, sele)
elif not quiet:
print(' Selection: ' + sele)
return sele
[docs]
def diff(sele1, sele2, byres=1, name=None, operator='in', quiet=0, *, _self=cmd):
'''
DESCRIPTION
Difference between two molecules
ARGUMENTS
sele1 = string: atom selection
sele2 = string: atom selection
byres = 0/1: report residues, not atoms (does not affect selection)
{default: 1}
operator = in/like/align: operator to match atoms {default: in}
SEE ALSO
symdiff
'''
byres, quiet = int(byres), int(quiet)
if name is None:
name = _self.get_unused_name('diff')
if operator == 'align':
alnobj = _self.get_unused_name('__aln')
_self.align(sele1, sele2, cycles=0, transform=0, object=alnobj)
sele = '(%s) and not %s' % (sele1, alnobj)
_self.select(name, sele)
_self.delete(alnobj)
else:
sele = '(%s) and not ((%s) %s (%s))' % (sele1, sele1, operator, sele2)
_self.select(name, sele)
if not quiet:
if byres:
seleiter = 'byca ' + name
expr = 'print("/%s/%s/%s/%s`%s" % (model,segi,chain,resn,resi))'
else:
seleiter = name
expr = 'print("/%s/%s/%s/%s`%s/%s" % (model,segi,chain,resn,resi,name))'
_self.iterate(seleiter, expr)
return name
[docs]
def symdiff(sele1, sele2, byres=1, name=None, operator='in', quiet=0, *, _self=cmd):
'''
DESCRIPTION
Symmetric difference between two molecules
SEE ALSO
diff
'''
byres, quiet = int(byres), int(quiet)
if name is None:
name = _self.get_unused_name('symdiff')
tmpname = _self.get_unused_name('__tmp')
diff(sele1, sele2, byres, name, operator, quiet)
diff(sele2, sele1, byres, tmpname, operator, quiet)
_self.select(name, tmpname, merge=1)
_self.delete(tmpname)
return name
[docs]
def collapse_resi(selection='(sele)', quiet=1, *, _self=cmd):
'''
DESCRIPTION
Returns a compact selection macro for the given selection on
residue number (resi) level.
Rewrite of http://pymolwiki.org/index.php/CollapseSel
ARGUMENTS
selection = string: atom selection {default: (sele)}
'''
from collections import defaultdict
s_dict = defaultdict(set)
_self.iterate(selection, 's_dict[model,segi,chain].add(resv)', space=locals())
r_all = []
for key, s in s_dict.items():
s = sorted(s)
r = [[s[0], s[0]]]
for i in s[1:]:
if i <= r[-1][1] + 1:
r[-1][1] = i
else:
r.append([i, i])
resi = '+'.join(('%d-%d' % (f, t) if f != t else '%d' % (f)) for (f, t) in r)
r_all.append('/%s/%s/%s/' % key + resi)
if not int(quiet):
for r in r_all:
print(' collapse_resi: ' + str(r))
return ' '.join(r_all)
[docs]
def wait_for(name, state=0, quiet=1, *, _self=cmd):
'''
DESCRIPTION
Wait for "name" to be available as selectable object.
'''
if _self.count_atoms('?' + name, 1, state) == 0:
s = _self.get_setting_boolean('suspend_updates')
if s:
_self.set('suspend_updates', 0)
_self.refresh()
if s:
_self.set('suspend_updates')
[docs]
def select_distances(names='', name='sele', state=1, selection='all', cutoff=-1, quiet=1, *, _self=cmd):
'''
DESCRIPTION
Turns a distance object into a named atom selection.
ARGUMENTS
names = string: names of distance objects (no wildcards!) {default: all
measurement objects}
name = a unique name for the selection {default: sele}
state = int: object state (-1: current, 0: all states) {default: 1}
SEE ALSO
get_raw_distances
'''
from collections import defaultdict
_assert_package_import()
from .querying import get_raw_distances
state, cutoff, quiet = int(state), float(cutoff), int(quiet)
states = [state] if state else list(range(1, _self.count_states(selection) + 1))
sele_dict = defaultdict(set)
for state in states:
distances = get_raw_distances(names, state, selection, _self=_self)
for idx1, idx2, dist in distances:
if cutoff <= 0.0 or dist <= cutoff:
sele_dict[idx1[0]].add(idx1[1])
sele_dict[idx2[0]].add(idx2[1])
_self.select(name, 'none')
tmp_name = _self.get_unused_name('_')
r = 0
for model in sele_dict:
_self.select_list(tmp_name, model, list(sele_dict[model]), mode='index')
r = _self.select(name, tmp_name, merge=1)
_self.delete(tmp_name)
if not quiet:
print(' Selector: selection "%s" defined with %d atoms.' % (name, r))
return r
[docs]
def select_range(name="", selection="", merge=1, *, _self=cmd):
"""
DESCRIPTION
Select all atoms between the first and last atom in the
current (or given) selection.
ARGUMENTS
name = str: Named selection to create {default: current active selection}
selection = str: Selection expression {default: name}
"""
if not name:
name = _self.get_names("selections", enabled_only=1)[0]
if not selection:
selection = name
tmpsele = _self.get_unused_name("_tmpsele")
_self.select(tmpsele, selection)
try:
for model in _self.get_object_list(tmpsele):
idx_first = _self.index(f"first (%{model} & {tmpsele})")[0]
idx_last = _self.index(f"last (%{model} & {tmpsele})")[0]
assert idx_first[0] == idx_last[0]
_self.select(name,
f"%{model} & index {idx_first[1]}-{idx_last[1]}",
merge=merge)
merge = 1
finally:
_self.delete(tmpsele)
[docs]
class select_temporary:
'''
DESCRIPTION
Context manager for creating a temporary named selection.
>>> with select_temporary(sele_expr) as named_sele:
... assert named_sele in cmd.get_names()
'''
def __init__(self, sele, prefix="_sele", *, _self=cmd):
self._self = _self
self.sele = sele
self.prefix = prefix
def __enter__(self):
self.name = self._self.get_unused_name(self.prefix)
self._self.select(self.name, self.sele)
return self.name
def __exit__(self, type, value, traceback):
self._self.delete(self.name)
# commands
cmd.extend('select_pepseq', select_pepseq)
cmd.extend('select_nucseq', select_nucseq)
cmd.extend('select_sspick', select_sspick)
cmd.extend('symdiff', symdiff)
cmd.extend('diff', diff)
cmd.extend('collapse_resi', collapse_resi)
cmd.extend('select_distances', select_distances)
cmd.extend('select_range', select_range)
# autocompletion
cmd.auto_arg[0].update([
('select_sspick', _auto_arg0_align),
('symdiff', _auto_arg0_align),
('diff', _auto_arg0_align),
('collapse_resi', cmd.auto_arg[0]['zoom']),
('select_distances', [
lambda: cmd.Shortcut(cmd.get_names_of_type('object:measurement')),
'distance object', '']),
])
cmd.auto_arg[1].update([
('select_pepseq', _auto_arg1_select),
('select_nucseq', _auto_arg1_select),
('symdiff', _auto_arg1_align),
('diff', _auto_arg1_align),
('select_range', _auto_arg1_select),
])
# vi:expandtab:smarttab