123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105 |
- import matplotlib.pyplot as plt
- import pandas as pd
- import numpy as np
- def percentile_from_pdf(pdf, bin_centers, percentile=0.5):
- cdf = 0
- for pdf_val, bin_low in zip(pdf, bin_centers[:-1]):
- cdf += pdf_val
- if cdf > percentile:
- return bin_low # TODO: Interpolate
- print(pdf, bin_centers)
- raise ValueError(f"couldn't find percentile: {percentile}, cdf: {cdf}" )
- def pdf_stats(pdf, bins):
- from collections import namedtuple
- Stats = namedtuple('Stats', ['hist', 'mean', 'median', 'quart_high', 'quart_low'])
- bin_centers = (bins[:-1] + bins[1:])/2
- mean = np.average(bin_centers, weights=pdf)
- median = percentile_from_pdf(pdf, bin_centers)
- quart_low = percentile_from_pdf(pdf, bin_centers, 0.25)
- quart_high = percentile_from_pdf(pdf, bin_centers, 0.75)
- return Stats((pdf, bins), mean, median, quart_high, quart_low)
- def get_stats_congress(year, age_max=110, parties=None, states=None):
- query = f'''\
- SELECT yob, position, party FROM Member
- WHERE congress={year} AND position IN ("Representative", "Senator")
- '''
- if parties:
- query += ' AND party IN (' + ", ".join(f'"{party}"' for party in parties) + ')'
- if states:
- query += ' AND state IN (' + ", ".join(f'"{state}"' for state in states) + ')'
- data = pd.read_sql_query(query, 'sqlite:///us_congress_members.sqlite3')
- data['age'] = year - data.yob
- pdf, bins = np.histogram(data.age, bins=age_max, range=(0, age_max), density=True)
- return pdf_stats(pdf, bins)
- def get_stats_genpop(year_data, age_max=110):
- pdf, bins = np.histogram(year_data.AGE, bins=age_max, range=(0, age_max), weights=year_data.PERWT, density=True)
- return pdf_stats(pdf, bins)
- def plot_pdf(genpop_stats, congress_stats):
- import matplotlib.pyplot as plt
- genpop_pdf, genpop_bins = genpop_stats.hist
- congress_pdf, congress_bins = congress_stats.hist
- plt.plot(genpop_bins[:-1], genpop_pdf, 'r.', label='U.S. Population')
- plt.plot(congress_bins[:-1], congress_pdf, 'b.', label='Congress')
- plt.legend()
- plt.show()
- def plot_yearly_stats(congress_stats, genpop_stats):
- congress_years = []
- congress_medians = []
- congress_quart_highs = []
- congress_quart_lows = []
- for year, year_stats in congress_stats.items():
- congress_years.append(year)
- congress_medians.append(year_stats.median)
- congress_quart_highs.append(year_stats.quart_high)
- congress_quart_lows.append(year_stats.quart_low)
- genpop_years = []
- genpop_medians = []
- genpop_quart_highs = []
- genpop_quart_lows = []
- for year, year_stats in genpop_stats.items():
- genpop_years.append(year)
- genpop_medians.append(year_stats.median)
- genpop_quart_highs.append(year_stats.quart_high)
- genpop_quart_lows.append(year_stats.quart_low)
- plt.fill_between(genpop_years, genpop_medians, genpop_quart_highs, color='b', alpha=0.3)
- plt.fill_between(genpop_years, genpop_medians, genpop_quart_lows, color='b', alpha=0.3)
- plt.plot(genpop_years, genpop_medians, 'b', label='U.S. Population')
- plt.fill_between(congress_years, congress_medians, congress_quart_highs, color='r', alpha=0.3)
- plt.fill_between(congress_years, congress_medians, congress_quart_lows, color='r', alpha=0.3)
- plt.plot(congress_years, congress_medians, 'r', label='Congress')
- plt.legend()
- plt.grid()
- plt.ylabel('Age')
- plt.xlabel('Year')
- plt.show()
- if __name__ == '__main__':
- stats_genpop = {}
- stats_congress = {}
- for year in range(1850, 2017):
- stats_congress[year] = get_stats_congress(year)
- people = pd.read_csv('usa_00001.csv', usecols=['YEAR', 'AGE', 'PERWT'], index_col='YEAR')
- for year in people.index.unique():
- stats_genpop[year] = get_stats_genpop(people.loc[year])
- plot_yearly_stats(stats_congress, stats_genpop)
|