Source code for utils.plotXML

import matplotlib
import numpy as np
import textwrap


class AxisError(Exception):
    def __init__(self):
        Exception.__init__(self, "Veritas plotting error")



def plot_fig(ax, root, x_axis, y_axis, mlc_x, mlc_y):
    '''
    Plots Gantry Rotation (GantryRtn) as a function as a function of Monitor Units (MUs)
    Also adds vertical lines where the kV or MV imaging points are scheduled.
    '''
    # Generate the plot
    if root is None:
        raise TypeError
    ax.clear()
    x_points = list()
    y_points = list()

    cp_search_string = ".//ControlPoints/Cp/"
    if mlc_x:
        x_axis = "Mlc/" + x_axis
    x_tags_lst = root.findall(cp_search_string + x_axis)

    if mlc_y:
        y_axis = "Mlc/" + y_axis

    y_tags_lst = root.findall(cp_search_string + y_axis)
    if (not y_tags_lst) or (not x_tags_lst):
        raise AxisError
    for cp_tags in root.findall(".//ControlPoints/Cp"):
        x_tag = cp_tags.find(x_axis)
        y_tag = cp_tags.find(y_axis)

        if (y_tag is not None) and (x_tag is not None):
            x_points = add_x_point(x_points, x_tag, mlc_x)
            y_points = add_y_point(y_points, y_tag, mlc_y, same=False)
        elif y_tag is None:
            # When y_points is not moving record the previous value
            x_points = add_x_point(x_points, x_tag, mlc_x)
            y_points = add_y_point(y_points, y_tag, mlc_y, same=True)

    # Extract imaging x_points values and imaging types
    ip = list()
    imagingType = list()

    for d in root.findall(".//ImagingPoint"):
        ip.append(float(d.find('Cp').text))

        if d.find('.//KV') is not None:
            imagingType.append("kV")
        elif d.find('.//MV') is not None:
            imagingType.append("MV")

    # Plot imaging points
    plotColor = {"kV": "blue", "MV": "red"}
    label_index = 0; # For stopping legend repetition
    if ip is not None:
        for value, flag in zip(ip, imagingType):
            idx, frac = int(value), value-int(value)
            cur_pt = x_points[int(idx)] + frac*(abs(x_points[int(idx)] - x_points[idx+1]))
            ax.plot(cur_pt*np.ones(len(y_points)), x_points, color=plotColor[flag],
                    label=flag if label_index < 2 else "")
            # tenary operator is trick to stop repetition of legend label
            label_index += 1
            ax.hold(True)


    # Plot control points
    axis_mapping = {"Mu": "Monitor Units (MUs)", "GantryRtn": "Gantry Rotation",
                    "CollRtn": "Column Rotation", "CouchVrt": "Couch Vertical",
                    "CouchLng": "Couch Longitude",
                    "CouchLat": "Couch Latitude", "CouchRtn": "Couch Rotation",
                    "X1": "X1", "X2":"X2", "Y1": "Y1", "Y2":"Y2",
                    "Mlc/A": "MLC A", "Mlc/B": "MLC B"}
    x_label = axis_mapping[x_axis]
    y_label = axis_mapping[y_axis]
    ax.plot(x_points, y_points, color="green", label=textwrap.fill(y_label, 10))
    # Now the labeling part
    ax.set_xlabel(x_label)
    ax.set_ylabel(y_label)

    title = {"Mu": "MU", "GantryRtn": "Gantry",
             "CollRtn": "CollRtn", "CouchVrt": "CouchVrt",
             "CouchLat": "CouchLat", "CouchRtn": "CouchRtn",
             "CouchLng": "CouchLng",
             "X1": "X1", "X2":"X2", "Y1": "Y1", "Y2":"Y2",
             "Mlc/A": "MLC A", "Mlc/B": "MLC B"}
    ax.set_title('{0} Vs. {1} Plots'.format(title[y_axis], title[x_axis]))
    handles, labels = ax.get_legend_handles_labels()
    ax.legend(handles, labels)
    ax.hold(False)
    return ax

def add_x_point(points, tag, mlc_no):
    if mlc_no:
        mlc_points = tag.text.split()
        result = mlc_points[mlc_no-1]
    else:
        result = tag.text
    points.append(float(result))
    return points



def add_y_point(points, tag, mlc_no, same=False):
    if same:
        points.append(points[-1])
        return points

    if mlc_no:
        mlc_points = tag.text.split()
        result = mlc_points[mlc_no-1]
    else:
        result = tag.text
    points.append(float(result))
    return points