content: import os
  
  import numpy as np
  import pandas as pd
  from pandas.tools.plotting import scatter_matrix
  import pylab
  import statsmodels.formula.api as smf
  import statsmodels.stats.api as sms
  
  from industry_to_days import get_census_mapper
  
  """
  USAGE:
  python full_2011.py
  
  CREATES:
  results/2011/clean.csv
  results/2011/corr.txt
  results/2011/het_breushpagan.txt
  results/2011/ols1.txt
  results/2011/ols2.txt
  results/2011/scatter_matrix.png
  results/2011/summary.txt
  """
  
  COL_ORDER = ['vacation', 'paid_vacation', 'age', 'fam_size', 'is_female',
               'income10', 'salary', 'is_employed']
  
  PSID_CSV = os.path.abspath(
          os.path.join(
              os.path.dirname(__file__),
              'psid', '2011.csv'))
  
  
  def get_f_path(fname):
      return os.path.abspath(
              os.path.join(
                  os.path.dirname(__file__),
                  'results', '2011', fname))
  
  
  CLEAN_CSV = get_f_path('clean.csv')
  CORR_TXT = get_f_path('corr.txt')
  HET_BP_TXT = get_f_path('het_breushpagan.txt')
  OLS1_TXT = get_f_path('ols1.txt')
  OLS2_TXT = get_f_path('ols2.txt')
  SCAT_MATRIX_PNG = get_f_path('scatter_matrix.png')
  SUMMARY_TXT = get_f_path('summary.txt')
  
  f_exists = (lambda file_: os.path.isfile(file_))
  
  
  def _calc_vacation(row):
      took, days, weeks, months = (row['took_vac'], row['days_vac'],
                                   row['weeks_vac'], row['months_vac'])
      if took in [8, 9] or (days in [998, 999]) or (months in [98, 99]) or (
          weeks in [98, 99]):
          return np.nan
      elif took == 5:
          return 0
      else:
          return days + (5 * weeks) + (22 * months)
  
  
  def _calc_salary(row):
      amt, unit = row['salary_amt'], row['salary_unit']
      if amt in [0.0, 9999998.0] or unit in [0, 7, 8, 9]:
          return np.nan
      if unit == 3:  # one week
          scalar = 52.0
      elif unit == 4:  # two weeks
          scalar = 26.0
      elif unit == 5:  # one month
          scalar = 12.0
      elif unit == 6:  # one year
          scalar = 1.0
      return scalar * amt
  
  
  def clean(df):
      # make sex into dummy for is_female
      df['is_female'] = df['sex'] - 1
      # remove unknown age values
      df.age = df.age.replace(999, np.nan)
      # figure out total vacation taken
      df['vacation'] = df.apply(_calc_vacation, axis=1)
      # fix salary to be annual amount
      df['salary'] = df.apply(_calc_salary, axis=1)
      # remove outliers
      df.ix[df.salary < 1e3] = np.nan
      df.ix[df.salary >= 400e3] = np.nan
      df.ix[df.income10 < 1e3] = np.nan
      df.ix[df.income10 >= 400e3] = np.nan
      # make employment into dummy for is_employed
      df['is_employed'] = df.employment
      # remove all those not working
      for i in range(2,10) + [99]:
          df.is_employed.replace(i, 0, inplace=True)
      # merge industry data
      df['paid_vacation'] = df.industry.map(get_census_mapper())
      # drop old values
      for col in ['took_vac', 'days_vac', 'weeks_vac', 'months_vac', 'industry',
                  'salary_amt', 'salary_unit', 'sex', 'employment']:
          df.drop(col, axis=1, inplace=True)
      df = df.reindex_axis(sorted(df.columns, key=COL_ORDER.index), axis=1)
      return df
  
  
  def do_stats(df):
      # Only view those that received vacation and are employed
      df.is_employed.replace(0.0, np.nan, inplace=True)
      df.paid_vacation.replace(0.0, np.nan, inplace=True)
      df.dropna(inplace=True)
      # No longer need this dummy
      df.drop('is_employed', axis=1, inplace=True)
  
      # Summary stats
      if not f_exists(SUMMARY_TXT):
          summary = df.describe().T
          summary = np.round(summary, decimals=3)
          with open(SUMMARY_TXT, 'w') as f:
              f.write(summary.to_string())
  
      # Test for autocorrelation: scatter matrix, correlation, run OLS
      if not f_exists(SCAT_MATRIX_PNG):
          scatter_matrix(df, alpha=0.2, figsize=(64, 64), diagonal='hist')
          pylab.savefig(SCAT_MATRIX_PNG, bbox_inches='tight')
      if not f_exists(CORR_TXT):
          corr = df.corr()
          corr = corr.reindex_axis(
              sorted(corr.columns, key=COL_ORDER.index), axis=0)
          corr = corr.reindex_axis(
              sorted(corr.columns, key=COL_ORDER.index), axis=1)
          for i, k in enumerate(corr):
              row = corr[k]
              for j in range(len(row)):
                  if j > i:
                      row[j] = np.nan
          with open(CORR_TXT, 'w') as f:
              f.write(np.round(corr, decimals=3).to_string(na_rep=''))
      if not f_exists(OLS1_TXT):
          ols_results = smf.ols(
              formula='vacation ~ paid_vacation + np.square(paid_vacation) + '
                      'age + fam_size + is_female + income10 + salary + '
                      'np.square(salary)',
              data=df).fit()
          with open(OLS1_TXT, 'w') as f:
              f.write(str(ols_results.summary()))
              f.write('\n\nCondition Number: {}'.format(
                  np.linalg.cond(ols_results.model.exog)))
  
      # Need to drop salary, too much autocorrelation
      df.drop('salary', axis=1, inplace=True)
  
      # Test for autocorrelation: scatter matrix, correlation, run OLS
      if not f_exists(HET_BP_TXT):
          ols_results = smf.ols(
              formula='vacation ~ paid_vacation + np.square(paid_vacation) + '
                      'age + fam_size + is_female + income10',
              data=df).fit()
          names = ['LM', 'LM P val.', 'F Stat.', 'F Stat. P val.']
          test = sms.het_breushpagan(ols_results.resid, ols_results.model.exog)
          f_p = test[3]
          with open(HET_BP_TXT, 'w') as f:
              str_ =  '\n'.join('{}: {}'.format(n, v)
                                for n, v in zip(names, test))
              f.write(str_ + '\n\n')
              if f_p < .01:
                  f.write('No Heteroskedasticity found.\n')
              else:
                  f.write('Warning: Heteroskedasticity found!\n')
  
      # no Heteroskedasticity found
      # final OLS results with robust standard errors
      if not f_exists(OLS2_TXT):
          ols_results = smf.ols(
              formula='vacation ~ paid_vacation + np.square(paid_vacation) + '
                      'age + fam_size + is_female + income10',
              data=df).fit().get_robustcov_results(cov_type='HAC', maxlags=1)
          with open(OLS2_TXT, 'w') as f:
              f.write(str(ols_results.summary()))
              f.write('\n\nCondition Number: {}'.format(
                  np.linalg.cond(ols_results.model.exog)))
      return df
  
  
  def main():
      df = None
      if f_exists(CLEAN_CSV):
          df = pd.io.parsers.read_csv(CLEAN_CSV)
          df.drop('Unnamed: 0', axis=1, inplace=True)
      else:
          with open(PSID_CSV) as csv:
              df = pd.io.parsers.read_csv(csv)
          df = clean(df)
          # write output to a file
          with open(CLEAN_CSV, 'w+') as csv:
              df.to_csv(path_or_buf=csv)
      return do_stats(df)
  
  
  if __name__ == '__main__':
      main()
      print '2011 succeeds! :)'
  