This is a read-only mirror of pymolwiki.org

Spectrumany

From PyMOL Wiki
Revision as of 14:10, 17 September 2010 by Speleo3 (talk | contribs) (support "resi" as expression)
Jump to navigation Jump to search

This script works similar to the spectrum command, but instead of predefined palettes, any color sequence can be used.

The color sequence is given by a space separated list of colors, so palette "red_white_blue" is the same as color sequence "red white blue".

Example

fetch 2x19

# these two produce the same result
spectrum count, red_white_blue, chain B
spectrumany count, red white blue, chain B

# gradient of different green colors
spectrumany count, smudge palegreen limegreen limon green forest, chain B

The Script

'''
(c) 2010 Thomas Holder
'''

from pymol import cmd, stored

def spectrumany(expression, color_list, selection='(all)', minimum=None, maximum=None):
    '''
DESCRIPTION

    Define a color spectrum with as many color-stops as you like (at least 2).

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
    '''
    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

    count_expr = 'index'
    expression = {'pc': 'partial_charge', 'fc': 'formal_charge',
            'count': count_expr}.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]:
        stored.e = list()
        cmd.iterate(selection, 'stored.e.append(%s)' % (minmax_expr))
        if minimum is None:
            minimum = min(stored.e)
        if maximum is None:
            maximum = max(stored.e)
    minimum, maximum = float(minimum), float(maximum)
    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(i * 255 for i 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

cmd.extend('spectrumany', spectrumany)