Survey Data with Pandas and StatsModels#
Examples and exercises for a tutorial on survey data analysis.
Click here to run this notebook on Colab. For the exercises, use the version with blanks.
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
# Show NumPy scalars as 0.36, not np.float64(0.36)
np.set_printoptions(legacy='1.25')
from os.path import basename, exists
def download(url):
filename = basename(url)
if not exists(filename):
from urllib.request import urlretrieve
local, _ = urlretrieve(url, filename)
print('Downloaded ' + local)
def decorate(**options):
"""Decorate the current axes.
Call decorate with keyword arguments like
decorate(title='Title',
xlabel='x',
ylabel='y')
The keyword arguments can be any of the axis properties
https://matplotlib.org/stable/api/axes_api.html
"""
legend = options.pop("legend", True)
loc = options.pop("loc", "best")
# Pass options to Axis.set
ax = plt.gca()
ax.set(**options)
# Add a legend if there are any labeled elements
handles, labels = ax.get_legend_handles_labels()
if handles and legend:
ax.legend(handles, labels, loc=loc)
# Tight layout is generally a good idea
plt.tight_layout()
Data#
We will use data from the General Social Survey (GSS). The raw dataset is big, so I’ve prepared an extract, which the following cell downloads. The code to generate the extract is in this notebook.
# This dataset is prepared in the GssExtract repository
DATA_PATH = "https://github.com/AllenDowney/GssExtract/raw/main/data/interim/"
filename = "gss_extract_pacs_workshop.hdf"
download(DATA_PATH + filename)
Distributions#
We’ll see two ways to describe distributions: value counts and KDE.
And for joint distributions, we’ll use a cross tabulation.
def value_counts(seq, **options):
"""Make a series of values and the number of times they appear.
Args:
seq: sequence
returns: Pandas Series
"""
return pd.Series(seq).value_counts(**options).sort_index()
Exercise#
Use value_counts to plot the distribution of cohort as a bar plot.
Then use sns.kdeplot to show the same distribution as a continuous quantity.
Cross tabulation#
Cross tabulation is a way to represent the joint distribution of two variables.
Target variable#
I’ve selected four variables you might want to work with. Whichever one you choose, put its cell last.
# https://gssdataexplorer.norc.org/variables/434/vshow
varname = 'happy'
question = """Taken all together, how would you say things are these days--
would you say that you are very happy, pretty happy, or not too happy?
"""
responses = ['Very happy', 'Pretty happy', 'Not too happy']
ylabel = "Percent saying 'very happy'"
# https://gssdataexplorer.norc.org/variables/439/vshow
varname = 'helpful'
question = """Would you say that most of the time people try to be helpful,
or that they are mostly just looking out for themselves?
"""
responses = ['Helpful', 'Look out\nfor themselves', 'Depends']
ylabel = "Percent saying 'helpful'"
# https://gssdataexplorer.norc.org/variables/440/vshow
varname = 'fair'
question = """Do you think most people would try to take advantage of you
if they got a chance, or would they try to be fair?
"""
# notice that the negative response is first here!
responses = ['Take advantage', 'Fair', 'Depends']
ylabel = "Percent saying 'fair'"
# https://gssdataexplorer.norc.org/variables/441/vshow
varname = 'trust'
question = """Generally speaking, would you say that most people can be trusted
or that you can't be too careful in dealing with people?
"""
responses = ['Trusted', "Can't be\ntoo careful", 'Depends']
ylabel = "Percent saying 'trusted'"
Exercise#
Do you think people who say “trust” are also more likely to say “helpful”?
To find out, make a cross tabulation of trust and helpful.
Consider using normalize='index'.
Make a heatmap that shows this joint distribution.
What relationships can you see in the responses?
Time Series#
To see how responses have changed over time, we’ll use cross tabulation again.
Smoothing#
To smooth the time series, we’ll use LOWESS (locally weighted scatterplot smoothing).
def make_lowess(series, frac=0.5):
"""Use LOWESS to compute a smooth line.
series: pd.Series
returns: pd.Series
"""
y = series.to_numpy()
x = series.index.to_numpy()
smooth_array = lowess(y, x, frac)
index, data = np.transpose(smooth_array)
return pd.Series(data, index=index)
Exercise#
Use this function to make smooth_series.
Plot the result along with time_series.
Binary series#
It will be useful to convert the responses to a binary variable represented with 0s and 1s.
Pivot tables#
Now we’ll use a pivot table to put respondents in groups and see how the responses have changed over time.
polviews_map = {
1: 'Liberal',
2: 'Liberal',
3: 'Liberal',
4: 'Moderate',
5: 'Conservative',
6: 'Conservative',
7: 'Conservative',
}
Better colors#
We’ll select colors from a Seaborn color palette.
Exercise#
Loop through the groups in color_map.
For each one, extract a column from table and plot the data points; then plot a smooth line.
Groupby#
Now that we have y as a binary variable, we have another way to compute the time series and the pivot table.
There are pros and cons of each method.
Logistic regression#
We’ll use logistic regression to model changes over time.
def fit_model(data, x_range, x_shift):
formula = 'y ~ x + I(x**2)'
model = smf.logit(formula, data=data).fit(disp=False)
pred_df = pd.DataFrame(dict(x=x_range - x_shift))
pred = model.predict(pred_df) * 100
pred.index = x_range
pred.name = 'model'
return pred
Exercise#
Use subset.groupby('polviews3') to loop through the groups.
For each one, use fit_model to generate and plot a “prediction” for the group.
Group by cohort#
Now let’s group respondents by decade of birth and see how each cohort changes over time.
def round_into_bins(series, bin_width, low=0, high=None):
"""Rounds values down to the bin they belong in.
series: pd.Series
bin_width: number, width of the bins
returns: Series of bin values (with NaN preserved)
"""
if high is None:
high = series.max()
bins = np.arange(low, high + bin_width, bin_width)
indices = np.digitize(series, bins)
result = pd.Series(bins[indices - 1], index=series.index, dtype='float')
result[series.isna()] = np.nan
return result
All the Cohorts#
Now let’s compare cohorts. The following function encapsulates the steps from the previous section.
def plot_cohort(df, cohort, color='C0'):
cohort_df = df.query("cohort10 == @cohort").dropna(subset=['y', 'x'])
cohort_series = cohort_df.groupby('year')['y'].mean() * 100
year_range = cohort_series.index
pred_series = fit_model(cohort_df, year_range, year_shift)
x, y = pred_series.index[-1], pred_series.values[-1]
plt.text(x, y, cohort, ha='left', va='center')
pred_series.plot(label=cohort, color=color)
Exercise#
Plot pred_all.
Then loop through cohorts and colors, and call plot_cohort for each group.
Depending on which target variable you chose, you might get a warning from StatsModels. The youngest cohort, born in the 2000s, is in only two surveys, 2018 and 2022 – not enough data to fit a quadratic model, so the optimizer might not converge. The predictions for those two years are still good.
Simpson’s Paradox#
The composition of the population changes over time.
Sampling weights#
The GSS uses stratified sampling, so some respondents represent more people in the population than others. We can use resampling to correct for stratified sampling.
def resample_by_year(df, column='wtssall'):
"""Resample rows within each year using weighted sampling.
df: DataFrame
column: string name of weight variable
returns: DataFrame
"""
grouped = df.groupby('year')
samples = [group.sample(n=len(group), replace=True, weights=group[column])
for _, group in grouped]
sample = pd.concat(samples, ignore_index=True)
return sample
Copyright 2025–2026 Allen Downey
Code license: MIT License
Text license: Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International