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.
