This is a read-only mirror of pymolwiki.org

Spectrumany

From PyMOL Wiki
Revision as of 07:09, 26 July 2010 by Speleo3 (talk | contribs) (minimum/maximum arguments swapped)
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

    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)

    if None in [minimum, maximum]:
        stored.e = list()
        cmd.iterate(selection, 'stored.e.append(%s)' % (expression))
        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 == count_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_name = 'ss_%d_%d' % (p, i)
            col_list = [colvec[p+1][j] * ii + colvec[p][j] * (1.0 - ii) for j in range(3)]
            cmd.set_color(col_name, col_list)
            val_end = val_range * (i + 1 + p * steps) / steps_total + minimum
            if expression == count_expr:
                cmd.color(col_name, '(%s) and index %d-%d' % (selection, 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)