StackOverflow Compensation Survey Plotting

compensation_comparison.py

# -*- coding: utf-8 -*-
"""
To use this, drop the file

'Full Results - Stack Overflow Developer Survey - 2015.csv'

from

https://drive.google.com/file/d/0Bzd_CzYvUxE5U1NSWnA2SFVKX00/view

into the same directory, then run the script.
"""

import scipy as sp
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

def get_data():
    data_path = 'Full Results - Stack Overflow Developer Survey - 2015.csv'
    data = pd.read_csv(data_path, skiprows=1)
    data = data[data['Country'] == 'United Kingdom']

    return data[['Compensation', 'Occupation']].dropna()

def get_mean_compensations(data):
    salary_map = {'Rather not say':    sp.nan,
                  'Unemployed':        0,
                  'Less than $20,000': 10000,
                  '$20,000 - $40,000': 30000,
                  '$40,000 - $60,000': 50000,
                  '$60,000 - $80,000': 70000,
                  '$80,000 - $100,000': 90000,
                  '$100,000 - $120,000': 110000,
                  '$120,000 - $140,000': 130000,
                  '$140,000 - $160,000': 150000,
                  'More than $160,000': 170000}

    data = data.copy()
    data['Compensation'] = data['Compensation'].apply(salary_map.get).dropna()
    groups = data.groupby('Occupation')
    return groups.aggregate(sp.mean), len(groups)

def plot_mean_compensations(mean_compensation, count):
    ax = mean_compensation.sort('Compensation').plot(kind='barh')
    ax.set_title('UK compensation by occupation \nn = {:,}'.format(count))
    ax.set_xlim([0, 150000])
    ax.set_xticklabels(['${:,}k'.format(int(x)/1000) for x in plt.gca().xaxis.get_majorticklocs()])
    ax.set_xlabel('Mean compensation (USD)')
    ax.figure.set_size_inches(10, 10)
    ax.legend_.remove()

    return ax

ax = plot_mean_compensations(*get_mean_compensations(get_data()))
ax.figure.savefig('uk-compensation.png', bbox_inches='tight')