matplotlib python utility functions

util.py

from matplotlib.pyplot import *
from numpy import *

__author__ = 'Flaviu Tamas'


def add_trendline(x_data, y_data, axis_, degree=1, color="black", decimals=2):
    trendfunc, r2 = trendline(x_data, y_data, degree)
    axis_.plot(x_data, trendfunc(x_data), color=color, label=trendfunc_to_latex(trendfunc, decimals))


def add_scatter(x_data, y_data, axis_, title=None, x_label=None, y_label=None, data_label=None, color="black"):
    if title is not None:
        axis_.set_title(title)
    if x_label is not None:
        axis_.set_xlabel(x_label)
    if y_label is not None:
        axis_.set_ylabel(y_label)
    if data_label is not None:
        axis_.scatter(x_data, y_data, color=color, label=data_label)
    else:
        axis_.scatter(x_data, y_data, color=color)


def trendline(x_data, y_data, degree=1):
    trendfunc = poly1d(polyfit(x_data, y_data, degree))

    # http://stackoverflow.com/a/895063
    yhat = trendfunc(x_data)
    ybar = sum(y_data) / len(y_data)
    ssreg = sum((yhat - ybar) ** 2)
    sstot = sum((y_data - ybar) ** 2)
    r2 = ssreg / sstot

    return trendfunc, r2


def trendfunc_to_latex(func, decimals=2):
    coeffs = list(reversed(func.coeffs.tolist()))
    power = 0
    accm = ""
    for coeff in coeffs:
        temp = ""
        if coeff > 0 and power < len(coeffs) - 1:
            temp += temp + "+"
        temp += ("%." + str(decimals) + "f") % coeff
        if power == 1:
            temp += "x"
        elif power != 0:
            temp += "x^" + str(power)
        accm = temp + accm
        power += 1
    return '$' + accm + '$'


def add_legend(axes_, pos="lower center"):
    axes_.legend(loc=pos)


def new_fig_axis(num=111):
    fig = figure()
    return fig, fig.add_subplot(num)