Subsetting data
Almost every statistical modeling project I have worked on has required removing some data from the analysis. Often, this is because of missing values or outliers. Sometimes, there are theoretical reasons for limiting our analysis to a subset of the data. For example, we have weather data going back to 1600, but our analysis goals only involve changes in weather since 1900. Fortunately, the subsetting tools in pandas are quite powerful and flexible. We will work with data from the United States National Longitudinal Survey (NLS) of Youth in this section.
Note
The NLS of Youth is conducted by the United States Bureau of Labor Statistics. This survey started with a cohort of individuals in 1997 who were born between 1980 and 1985, with annual follow-ups each year through 2017. For this recipe, I pulled 89 variables on grades, employment, income, and attitudes toward government from the hundreds of data items on the survey. Separate files for SPSS, Stata, and SAS can be downloaded from the repository. The NLS data is available for public use at https://www.nlsinfo.org/investigator/pages/search.
Let's start subsetting the data using pandas:
- We will start by loading the NLS data. We also set an index:
import pandas as pd import numpy as np nls97 = pd.read_csv("data/nls97.csv") nls97.set_index("personid", inplace=True) - Let's select a few columns from the NLS data. The following code creates a new DataFrame that contains some demographic and employment data. A useful feature of pandas is that the new DataFrame retains the index of the old DataFrame, as shown here:
democols = ['gender','birthyear','maritalstatus', 'weeksworked16','wageincome','highestdegree'] nls97demo = nls97[democols] nls97demo.index.name 'personid'
- We can use slicing to select rows by position.
nls97demo[1000:1004]selects every row, starting from the row indicated by the integer to the left of the colon (1000, in this case) up to, but not including, the row indicated by the integer to the right of the colon (1004). The row at1000is the 1,001st row because of zero-based indexing. Each row appears as a column in the output since we have transposed the resulting DataFrame:nls97demo[1000:1004].T personid 195884 195891 195970\ gender Male Male Female birthyear 1981 1980 1982 maritalstatus NaN Never-married Never-married weeksworked16 NaN 53 53 wageincome NaN 14,000 52,000 highestdegree 4.Bachelors 2.High School 4.Bachelors personid 195996 gender Female birthyear 1980 maritalstatus NaN weeksworked16 NaN wageincome NaN highestdegree 3.Associates
- We can also skip rows over the interval by setting a value for the step after the second colon. The default value for the step is 1. The value for the following step is 2, which means that every other row between
1000and1004will be selected:nls97demo[1000:1004:2].T personid 195884 195970 gender Male Female birthyear 1981 1982 maritalstatus NaN Never-married weeksworked16 NaN 53 wageincome NaN 52,000 highestdegree 4.Bachelors 4. Bachelors
- If we do not include a value to the left of the colon, row selection will start with the first row. Notice that this returns the same DataFrame as the
headmethod does:nls97demo[:3].T personid 100061 100139 100284 gender Female Male Male birthyear 1980 1983 1984 maritalstatus Married Married Never-married weeksworked16 48 53 47 wageincome 12,500 120,000 58,000 highestdegree 2.High School 2. High School 0.None nls97demo.head(3).T personid 100061 100139 100284 gender Female Male Male birthyear 1980 1983 1984 maritalstatus Married Married Never-married weeksworked16 48 53 47 wageincome 12,500 120,000 58,000 highestdegree 2.High School 2.High School 0. None
- If we use a negative number, -n, to the left of the colon, the last n rows of the DataFrame will be returned. This returns the same DataFrame as the
tailmethod does:nls97demo[-3:].T personid 999543 999698 999963 gender Female Female Female birthyear 1984 1983 1982 maritalstatus Divorced Never-married Married weeksworked16 0 0 53 wageincome NaN NaN 50,000 highestdegree 2.High School 2.High School 4. Bachelors nls97demo.tail(3).T personid 999543 999698 999963 gender Female Female Female birthyear 1984 1983 1982 maritalstatus Divorced Never-married Married weeksworked16 0 0 53 wageincome NaN NaN 50,000 highestdegree 2.High School 2.High School 4. Bachelors
- We can select rows by index value using the
locaccessor. Recall that for thenls97demoDataFrame, the index ispersonid. We can pass a list of the index labels to thelocaccessor, such asloc[[195884,195891,195970]], to get the rows associated with those labels. We can also pass a lower and upper bound of index labels, such asloc[195884:195970], to retrieve the indicated rows:nls97demo.loc[[195884,195891,195970]].T personid 195884 195891 195970 gender Male Male Female birthyear 1981 1980 1982 maritalstatus NaN Never-married Never-married weeksworked16 NaN 53 53 wageincome NaN 14,000 52,000 highestdegree 4.Bachelors 2.High School 4.Bachelors nls97demo.loc[195884:195970].T personid 195884 195891 195970 gender Male Male Female birthyear 1981 1980 1982 maritalstatus NaN Never-married Never-married weeksworked16 NaN 53 53 wageincome NaN 14,000 52,000 highestdegree 4.Bachelors 2.High School 4.Bachelors
- To select rows by position, rather than by index label, we can use the
ilocaccessor. We can pass a list of position numbers, such asiloc[[0,1,2]], to the accessor to get the rows at those positions. We can pass a range, such asiloc[0:3], to get rows between the lower and upper bound, not including the row at the upper bound. We can also use theilocaccessor to select the last n rows.iloc[-3:]selects the last three rows:nls97demo.iloc[[0,1,2]].T personid 100061 100139 100284 gender Female Male Male birthyear 1980 1983 1984 maritalstatus Married Married Never-married weeksworked16 48 53 47 wageincome 12,500 120,000 58,000 highestdegree 2.High School 2.High School 0. None nls97demo.iloc[0:3].T personid 100061 100139 100284 gender Female Male Male birthyear 1980 1983 1984 maritalstatus Married Married Never-married weeksworked16 48 53 47 wageincome 12,500 120,000 58,000 highestdegree 2.High School 2.High School 0. None nls97demo.iloc[-3:].T personid 999543 999698 999963 gender Female Female Female birthyear 1984 1983 1982 maritalstatus Divorced Never-married Married weeksworked16 0 0 53 wageincome NaN NaN 50,000 highestdegree 2.High School 2.High School 4. Bachelors
Often, we need to select rows based on a column value or the values of several columns. We can do this in pandas by using Boolean indexing. Here, we pass a vector of Boolean values (which can be a Series) to the loc accessor or the bracket operator. The Boolean vector needs to have the same index as the DataFrame.
- Let's try this using the
nightlyhrssleepcolumn on the NLS DataFrame. We want a Boolean Series that isTruefor people who sleep 6 or fewer hours a night (the 33rd percentile) andFalseifnightlyhrssleepis greater than 6 or is missing.sleepcheckbool = nls97.nightlyhrssleep<=lowsleepthresholdcreates the boolean Series. If we display the first few values ofsleepcheckbool, we will see that we are getting the expected values. We can also confirm that thesleepcheckboolindex is equal to thenls97index:nls97.nightlyhrssleep.head() personid 100061 6 100139 8 100284 7 100292 nan 100583 6 Name: nightlyhrssleep, dtype: float64 lowsleepthreshold = nls97.nightlyhrssleep.quantile(0.33) lowsleepthreshold 6.0 sleepcheckbool = nls97.nightlyhrssleep<=lowsleepthreshold sleepcheckbool.head() personid 100061 True 100139 False 100284 False 100292 False 100583 True Name: nightlyhrssleep, dtype: bool sleepcheckbool.index.equals(nls97.index) True
Since the sleepcheckbool Series has the same index as nls97, we can just pass it to the loc accessor to create a DataFrame containing people who sleep 6 hours or less a night. This is a little pandas magic here. It handles the index alignment for us:
lowsleep = nls97.loc[sleepcheckbool] lowsleep.shape (3067, 88)
- We could have created the
lowsleepsubset of our data in one step, which is what we would typically do unless we need the Boolean Series for some other purpose:lowsleep = nls97.loc[nls97.nightlyhrssleep<=lowsleepthreshold] lowsleep.shape (3067, 88)
- We can pass more complex conditions to the
locaccessor and evaluate the values of multiple columns. For example, we can select rows wherenightlyhrssleepis less than or equal to the threshold andchildathome(number of children living at home) is greater than or equal to3:lowsleep3pluschildren = \ nls97.loc[(nls97.nightlyhrssleep<=lowsleepthreshold) & (nls97.childathome>=3)] lowsleep3pluschildren.shape (623, 88)
Each condition in nls97.loc[(nls97.nightlyhrssleep<=lowsleepthreshold) & (nls97.childathome>3)] is placed in parentheses. An error will be generated if the parentheses are excluded. The & operator is the equivalent of and in standard Python, meaning that both conditions have to be True for the row to be selected. We could have used | for or if we wanted to select the row if either condition was True.
- Finally, we can select rows and columns at the same time. The expression to the left of the comma selects rows, while the list to the right of the comma selects columns:
lowsleep3pluschildren = \ nls97.loc[(nls97.nightlyhrssleep<=lowsleepthreshold) & (nls97.childathome>=3), ['nightlyhrssleep','childathome']] lowsleep3pluschildren.shape (623, 2)
We used three different tools to select columns and rows from a pandas DataFrame in the last two sections: the [] bracket operator and two pandas-specific accessors, loc and iloc. This will be a little confusing if you are new to pandas, but it becomes clear which tool to use in which situation after just a few months. If you came to pandas with a fair bit of Python and NumPy experience, you will likely find the [] operator most familiar. However, the pandas documentation recommends against using the [] operator for production code. The loc accessor is used for selecting rows by Boolean indexing or by index label, while the iloc accessor is used for selecting rows by row number.
This section was a brief primer on selecting columns and rows with pandas. Although we did not go into too much detail on this, most of what you need to know to subset data was covered, as well as everything you need to know to understand the pandas-specific material in the rest of this book. We will start putting some of that to work in the next two sections by creating frequencies and summary statistics for our features.
