plotting_functions.py

'''
This file contains functions to generate line graphs, histograms, scatterplots (with line of best fit), boxplots, and bar graphs using matplotlib.pyplot.
'''

import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
from textwrap import wrap


def plot_function(x, y, title, xlabel, ylabel, axis=111, c=None):
    fig = plt.figure()
    ax = fig.add_subplot(axis)
    ax.set_xlim(0, max(x))
    ax.set_ylim(0, max(y))
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_title(title)
    ax.plot(x, y, c=c, lw=2)
    plt.savefig('../images/{}.png'.format(title), dpi=300, bbox_inches = 'tight')


def plot_histogram(x, title, xlabel, ylabel='Frequencies', axis=111, c=None):
    fig = plt.figure()
    ax = fig.add_subplot(axis)
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_title(title)
    ax.hist(x, c=c, rwidth=0.95)
    plt.savefig('../images/{}.png'.format(title), dpi=300, bbox_inches = 'tight')


def plot_scatterplot(x, y, title, xlabel, ylabel, axis=111, c=None):
    fig = plt.figure()
    ax = fig.add_subplot(axis)
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_title(title)
    ax.scatter(x, y, c=c)
    slope, intercept = np.polyfit(x, y, 1)
    ax.plot(x, x * slope + intercept, c=c, lw=2)
    plt.savefig('../images/{}.png'.format(title), dpi=300, bbox_inches = 'tight')


def plot_boxplot(x, title, xlabel, ylabel='Observed values', axis=111, c=None):
    fig = plt.figure()
    ax = fig.add_subplot(axis)
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_title(title)
    bplot = ax.boxplot(x, medianprops={'c': 'k'}, patch_artist=True)
    # Fills boxes with gray
    [patch.set_facec('gray') for patch in bplot['boxes']]
    for line in bplot['medians']:
        (x, y) = line.get_xydata()[1]
        ax.annotate(y, (x, y))
    plt.savefig('../images/{}.png'.format(title), dpi=300, bbox_inches = 'tight')


def plot_horizontal_bargraph(x_column, y_column, df, title='', axis=111):
    # Sort df by numerical column
    df.sort_values(by=x_column, ascending=False, inplace=True)

    # Set up matplotlib axis
    fig = plt.figure()
    ax = fig.add_subplot(axis)

    # Change labels to title case and underscores to spaces
    x_column_formatted = x_column.replace('_', ' ').title()
    y_column_formatted = y_column.replace('_', ' ').title()

    # Formatting text labels
    df[y_column] = df[y_column].apply(lambda x : '\n'.join(wrap(x, 15)))

    sns.barplot(x_column, y_column, data=df, ax=ax, orient='h')

    # Set labels on axes & title
    ax.set_xlabel(x_column_formatted)
    ax.set_ylabel(y_column_formatted)
    if title == '':
        title = '{} by {}'.format(x_column_formatted, y_column_formatted)
    ax.set_title(title)

    widths = [p.get_width() for p in ax.patches]
    heights = [p.get_y() + p.get_height()/2.0 for p in ax.patches]
    coord_list = zip(widths, heights)
    # Adding value labels to bars
    [ax.text(coords[0]+0.04, coords[1] + 0.05, '{:1.2f}'.format(coords[0]), ha='center') for coords in coord_list]

    for i in range(df.shape[0]):
        plt.hlines(heights[i], df['min'].iloc[i], df['max'].iloc[i])

    plt.savefig('../images/{}.png'.format(title), dpi=300, bbox_inches = 'tight')