Scripts: colordist.py

#
# Change pseudobond colors based on the bond length. This is for showing distance constraints when
# interactively moving one molecule relative to another.
#
# To color all distance pseudobonds, or selected pseudobonds, or stop coloring all or selected:
#
#    runscript colordist.py 8.5 green red
#    runscript colordist.py 8.5 green red sel
#    runscript colordist.py
#    runscript colordist.py sel
#
class color_by_distance:

    def __init__(self):
        self.pseudobond_colors = {}
        self.tid = None

    def color_pseudobonds(self, pseudobonds, length, color_short, color_long):

        pbc = self.pseudobond_colors
        for pb in pseudobonds:
            pbc[pb] = (length, color_short, color_long)

        self.recolor()

        if self.tid is None:
            from chimera import triggers
            self.tid = triggers.addHandler('OpenState', self.moved_cb, None)

    def moved_cb(self, trigger_name, data, changes):
        if 'transformation change' in changes.reasons:
            self.recolor()

    def recolor(self):
        for pb, (length, color_short, color_long) in self.pseudobond_colors.items():
            pb.color = (color_long if pb.length() >= length else color_short)

    def stop_coloring(self, pseudobonds = None):
        pbc = self.pseudobond_colors
        if pseudobonds is None:
            pbc.clear()
        else:
            for pb in pseudobonds:
                if pb in pbc:
                    pbc.pop(pb)
        if len(pbc) == 0 and not self.tid is None:
            from chimera import triggers
            triggers.deleteHandler('OpenState', self.tid)
            self.tid = None

def parse_pseudobonds(pb_spec):

    from chimera import specifier
    sel = specifier.evalSpec(pb_spec)
    pblist = sel.pseudobonds()
    return pblist

def all_pseudobonds():

    pblist = []
    from chimera import openModels, PseudoBondGroup
    for m in openModels.list(hidden = True, modelTypes = [PseudoBondGroup]):
        if m.category == 'distance monitor':
            pblist.extend(m.pseudoBonds)
    return pblist

cbd = None

def parse_command(args):

    if args is None or len(args) not in (0,1,3,4):
        from Commands import CommandError
        raise CommandError('Wrong number of arguments: runscript colorconstraints.py 8.5 green red sel')

    # Get color by distance updater.
    if hasattr(chimera, 'color_by_distance'):
        cbd = chimera.color_by_distance
    else:
        chimera.color_by_distance = cbd = color_by_distance()

    # Stop color updating if no arguments given.
    if len(args) == 0:
        cbd.stop_coloring()
        return

    # Uncolor pseudobonds if just one argument was given.
    if len(args) == 1:
        cbd.stop_coloring(parse_pseudobonds(args[0]))
        return
    
    # Get distance
    from Commands import float_arg
    length = float_arg(args[0])

    # Get colors
    color1, color2 = args[1:3]
    from Commands import color_arg
    srgba = color_arg(color1)
    lrgba = color_arg(color2)
    from chimera import specifier, MaterialColor
    short_color = MaterialColor(*srgba)
    long_color = MaterialColor(*lrgba)

    # Get pseudobonds
    pblist = parse_pseudobonds(args[3]) if len(args) == 4 else all_pseudobonds()
    if len(pblist) == 0:
        from Commands import CommandError
        raise CommandError('No pseudobonds to color')

    # Setup automatic coloring
    cbd.color_pseudobonds(pblist, length, short_color, long_color)

parse_command(globals().get('arguments'))