Skip to content
Permalink
 
 
Cannot retrieve contributors at this time
'''
Viewing and Coloring Stuff
(c) 2011 Thomas Holder, MPI for Developmental Biology
License: BSD-2-Clause
'''
from pymol import cmd, CmdException
def nice(selection='(all)', simple=1e5):
'''
DESCRIPTION
My favourite representation
Color: by chain (elem C) and by element (others)
Representation: cartoon (polymer), sticks (organic),
spheres (inorganic) and nonbonded (solvent)
USAGE
nice [ selection ]
'''
simple = int(simple)
cmd.util.cbc(selection)
cmd.color('atomic', '(%s) and not elem C' % (selection))
if simple and cmd.count_atoms(selection) >= simple:
cmd.show_as('ribbon', selection)
cmd.show_as('lines', '(%s) and organic' % (selection))
cmd.show_as('nonbonded', '(%s) and inorganic' % (selection))
else:
cmd.show_as('cartoon', selection)
cmd.show_as('sticks', '(%s) and organic' % (selection))
cmd.show_as('spheres', '(%s) and inorganic' % (selection))
cmd.show_as('nonbonded', '(%s) and solvent' % (selection))
def get_color_family(color):
'''
DESCRIPTION
API only. Get the colors list from a color submenu (like all "Reds").
'''
from pymol import menu
colors = []
try:
for (c,_,expr) in getattr(menu, color)(cmd, ''):
if c != 1:
continue
assert expr.startswith('cmd.color(')
col = expr[10:].partition(',')[0]
if col[0] in ['"', "'"]:
col = col[1:-1]
colors.append(col)
except:
raise CmdException(repr(color))
return colors
def cbm(selection='all', first_color=2):
'''
DESCRIPTION
Color by molecule
USAGE
cbm [ selection [, first_color ]]
'''
import itertools
if isinstance(first_color, str) and first_color.endswith('s'):
colors = get_color_family(first_color)
col_it = itertools.cycle(colors)
else:
col_it = itertools.count(int(first_color))
for model in cmd.get_object_list('(' + selection + ')'):
col = next(col_it)
if selection in ['*', 'all']:
# explicit model coloring for cmd.get_object_color_index to work
cmd.color(col, model)
else:
cmd.color(col, '%s and (%s)' % (model, selection))
def cbs(selection='all', first_color=2, quiet=1):
'''
DESCRIPTION
Color by segment
'''
import itertools
if isinstance(first_color, str) and first_color.endswith('s'):
colors = get_color_family(first_color)
col_it = itertools.cycle(colors)
else:
col_it = itertools.count(int(first_color))
segi_colors = dict()
def callback(segi):
if segi not in segi_colors:
segi_colors[segi] = next(col_it)
return segi_colors[segi]
cmd.alter(selection, 'color = callback(segi)', space=locals())
cmd.rebuild()
expression_sc = cmd.Shortcut([
'count',
'resi',
'b',
'q',
'pc',
])
def spectrumany(expression, color_list, selection='(all)', minimum=None, maximum=None, quiet=1):
'''
DESCRIPTION
Define a color spectrum with as many color-stops as you like (at least 2).
Note: Since PyMOL 1.6, the regular `spectrum` command also supports
arbitrary color lists, so this command is obsolete.
USAGE
spectrumany expression, color_list [, selection [, minimum [, maximum ]]]
ARGUMENTS
expression = count, resi, b, q, or pc: respectively, atom count, residue
index, temperature factor, occupancy, or partial charge {default: count}
color_list = string: Space separated list of colors
... all other arguments like with `spectrum` command
EXAMPLE
spectrumany count, forest green yellow white
spectrumany b, red yellow white, (polymer), maximum=100.0
SEE ALSO
spectrum
'''
quiet = int(quiet)
colors = color_list.split()
if len(colors) < 2:
print('failed! please provide at least 2 colors')
return
colvec = [cmd.get_color_tuple(i) for i in colors]
parts = len(colvec) - 1
expression = {'pc': 'partial_charge', 'fc': 'formal_charge',
'count': 'index'}.get(expression, expression)
minmax_expr = {'resi': 'resv'}.get(expression, expression)
discrete_expr = ['index', 'resi']
if cmd.count_atoms(selection) == 0:
print('empty selection')
return
if None in [minimum, maximum]:
e_list = list()
cmd.iterate(selection, 'e_list.append(%s)' % (minmax_expr), space=locals())
if minimum is None:
minimum = min(e_list)
if maximum is None:
maximum = max(e_list)
minimum, maximum = float(minimum), float(maximum)
if not quiet:
print(' Spectrum: range (%.5f to %.5f)' % (minimum, maximum))
if maximum == minimum:
print('no spectrum possible, only equal %s values' % (expression))
return
if expression in discrete_expr:
val_range = int(maximum - minimum + 1)
else:
val_range = maximum - minimum
cmd.color(colors[0], selection)
steps = 60 // parts
steps_total = steps * parts
val_start = minimum
for p in range(parts):
for i in range(steps):
ii = float(i)/steps
col_list = [colvec[p+1][j] * ii + colvec[p][j] * (1.0 - ii) for j in range(3)]
col_name = '0x%02x%02x%02x' % tuple(int(0xFF * v) for v in col_list)
val_end = val_range * (i + 1 + p * steps) // steps_total + minimum
if expression in discrete_expr:
cmd.color(col_name, '(%s) and %s %d-%d' % (selection, expression, val_start, val_end))
else:
cmd.color(col_name, '(%s) and %s > %f' % (selection, expression, val_start))
val_start = val_end
def spectrum_states(selection='all', representations='cartoon ribbon',
color_list='blue cyan green yellow orange red',
first=1, last=0, quiet=1):
'''
DESCRIPTION
Color each state in a multi-state object different.
(c) 2011 Takanori Nakane and Thomas Holder
USAGE
spectrum_states [ selection [, representations [, color_list [, first [, last ]]]]]
ARGUMENTS
selection = string: object names (works with complete objects only)
{default: all}
representations = string: space separated list of representations
{default: cartoon ribbon}
color_list = string: space separated list of colors {default: blue cyan
green yellow orange red}
SEE ALSO
spectrum, spectrumany
'''
from math import floor, ceil
first, last, quiet = int(first), int(last), int(quiet)
colors = color_list.split()
if len(colors) < 2:
raise CmdException('please provide at least 2 colors')
colvec = [cmd.get_color_tuple(i) for i in colors]
# filter for valid <repr>_color settings
settings = []
for r in representations.split():
if r[-1] == 's':
r = r[:-1]
s = r + '_color'
if s in cmd.setting.name_list:
settings.append(s)
elif not quiet:
print(' Warning: no such setting: ' + repr(s))
# object names only
selection = ' '.join(cmd.get_object_list('(' + selection + ')'))
if cmd.count_atoms(selection) == 0:
raise CmdException('empty selection')
if last < 1:
last = cmd.count_states(selection)
val_range = int(last - first + 1)
if val_range < 2:
raise CmdException('no spectrum possible, need more than 1 state')
for i in range(val_range):
p = float(i) / (val_range - 1) * (len(colvec) - 1)
p0, p1 = int(floor(p)), int(ceil(p))
ii = (p - p0)
col_list = [colvec[p1][j] * ii + colvec[p0][j] * (1.0 - ii) for j in range(3)]
col_name = '0x%02x%02x%02x' % tuple(int(0xFF * v) for v in col_list)
for s in settings:
cmd.set(s, col_name, selection, state=i+first)
class scene_preserve(object):
'''
DESCRIPTION
API only. Context manager to restore the current scene on exit.
'''
def __init__(self, **kwargs):
self.kwargs = kwargs
def __enter__(self):
import random
self.name = 'tmp_%d' % (random.randint(0, 1e8))
cmd.scene(self.name, 'store', **self.kwargs)
def __exit__(self, type, value, traceback):
cmd.scene(self.name, 'recall')
cmd.scene(self.name, 'delete')
# commands
cmd.alias('z', 'zoom visible')
cmd.alias('x', 'nice')
cmd.extend('nice', nice)
cmd.extend('cbm', cbm)
cmd.extend('cbs', cbs)
cmd.extend('spectrumany', spectrumany)
cmd.extend('spectrum_states', spectrum_states)
# tab-completion of arguments
cmd.auto_arg[0]['spectrumany'] = [ expression_sc , 'expression' , ', ' ]
cmd.auto_arg[1]['spectrumany'] = [ cmd.auto_arg[0]['color'][0], 'color', ' ' ]
cmd.auto_arg[2]['spectrumany'] = cmd.auto_arg[2]['spectrum']
cmd.auto_arg[0]['spectrum_states'] = cmd.auto_arg[0]['disable']
cmd.auto_arg[1]['spectrum_states'] = [ cmd.auto_arg[0]['show'][0], 'representation', ' ' ]
cmd.auto_arg[2]['spectrum_states'] = cmd.auto_arg[1]['spectrumany']
# vi:expandtab:smarttab