lineid_plot_utils

lineid_utils.py

import matplotlib as mpl
import lineid_plot


def color_text_boxes(ax, labels, colors, color_arrow=True):
    assert len(labels) == len(colors), \
        "Equal no. of colors and lables must be given"
    boxes = ax.findobj(mpl.text.Annotation)
    box_lables = lineid_plot.unique_labels(labels)
    for box in boxes:
        l = box.get_label()
        try:
            loc = box_labels.index(l)
        except ValueError:
            continue  # No changes for this box
        box.set_color(colors[loc])
        if color_arrow:
            box.arrow_patch.set_color(colors[loc])

    ax.figure.canvas.draw()


def color_lines(ax, labels, colors):
    assert len(labels) == len(colors), \
        "Equal no. of colors and lables must be given"
    lines = ax.findobj(mpl.lines.Line2D)
    line_labels = [i + "_line" for i in lineid_plot.unique_labels(labels)]
    for line in lines:
        l = line.get_label()
        try:
            loc = line_labels.index(l)
        except ValueError:
            continue  # No changes for this line
        line.set_color(colors[loc])

    ax.figure.canvas.draw()