Chi Square Independence Test for Two Pandas DF columns

Question

I want to calculate the scipy.stats.chi2_contingency() for two columns of a pandas DataFrame. The data is categorical, like this:

var1    var2
0       1
1       0
0       2
0       1
0       2

Here is the example data: TU Berlin Server

The task is to build the crosstable sums (contingency table) of each category-relationship. Example:

         var1
         0    1
---------------------
     0 | 0    1
var2 1 | 2    0
     2 | 2    0

I'm not really a coder, but this is what I got (working):

def create_list_sum_of_categories(df, var, cat, var2):
    list1 = []
    for cat2 in range(int(df[var2].min()), int(df[var2].max())+1):
            list1.append( len(df[ (df[var] == cat) & (df[var2] == cat2) ]))   
    return list1

def chi_square_of_df_cols(df,col1,col2):
    ''' for each category of col1 create list with sums of each category of col2'''
    result_list = []
    for cat in range(int(df[col1].min()), int(df[col1].max())+1):
        result_list.append(create_list_sum_of_categories(df,col1,cat,col2)) 

    return scs.chi2_contingency(result_list)


test_df = pd.read_csv('test_data_for_chi_square.csv')
print(chi_square_of_df_cols(test_df,'var1','var2'))

My question gears towards two things:

1. Can you confirm that this actually does what I want?
2. If you have suggestions to make this code more beautiful (e.g. include everything in one function), please go ahead!

Answer 1

AFAIK it does what you want (but on this site you should generally be sure that your code does what you want beforehand).

Beauty, eye of the Beholder, ...; that said, this code can be rewritten in a very concise manner:

import pandas as pd
import scipy.stats as scs


def categories(series):
    return range(int(series.min()), int(series.max()) + 1)


def chi_square_of_df_cols(df, col1, col2):
    df_col1, df_col2 = df[col1], df[col2]

    result = [[sum((df_col1 == cat1) & (df_col2 == cat2))
               for cat2 in categories(df_col2)]
              for cat1 in categories(df_col1)]

    return scs.chi2_contingency(result)


test_df = pd.read_csv('test_data_for_chi_square.csv')
print(chi_square_of_df_cols(test_df, 'var1', 'var2'))

Basically introducing one abstraction (categories) to make your intentions a bit clearer and then precomputing some results (df_col1, df_col2), using sum to count the number of matches instead of indexing into the data frame again.

I doubt numpy functions would give a speedup unless your arrays are a bit bigger, but you might profile that anyway if you need more speed.

However I'd also rather use the following instead in order to save some more CPU cycles by not recomputing categories and df_col1 == cat1 all the time:

def chi_square_of_df_cols(df, col1, col2):
    df_col1, df_col2 = df[col1], df[col2]
    cats1, cats2 = categories(df_col1), categories(df_col2)

    def aux(is_cat1):
        return [sum(is_cat1 & (df_col2 == cat2))
                for cat2 in cats2]

    result = [aux(df_col1 == cat1)
              for cat1 in cats1]

    return scs.chi2_contingency(result)

Answer 2

I would try to use existing pandas features where possible to keep this code minimal - this aids readability and reduces the possibility of bugs being introduced in complicated loop structures.

import pandas
from scipy.stats import chi2_contingency

def chisq_of_df_cols(df, c1, c2):
    groupsizes = df.groupby([c1, c2]).size()
    ctsum = groupsizes.unstack(c1)
    # fillna(0) is necessary to remove any NAs which will cause exceptions
    return(chi2_contingency(ctsum.fillna(0)))

test_df = pandas.DataFrame([[0, 1], [1, 0], [0, 2], [0, 1], [0, 2]], columns=['var1', 'var2'])
chisq_of_df_cols(test_df, 'var1', 'var2')

Answer 3

Using pandas.crosstab, this can be done in a single step:

pandas.crosstab(index=test_df['var1'],columns=test_df['var2'])

It will give this desired result:

var1    | 0 1
--------------------
    var2    |   
--------------------     
    0   |0  1
    1   |2  0
    2   |2  0

You can name the index and colnames and also get the row totals and column totals:

new_test_df = pandas.crosstab(index=test_df['var2'],columns=test_df['var1'],margins=True)
new_test_df.index = ['var2_0','var2_1','var2_2','coltotal']
new_test_df.columns= ['var1_0','var1_1','rowtotal']

Margins gives the totals. Columns and indexes can be used to name the columns.

Answer 4

Input (copy and paste from the original question):

test_df = pandas.DataFrame([[0, 1], [1, 0], [0, 2], [0, 1], [0, 2]],
                            columns=['var1', 'var2'])

Desired output (copy and paste from the original question):

        var1
         0    1
---------------------
     0 | 0    1
var2 1 | 2    0
     2 | 2    0

One line solution using crosstab:

pandas.crosstab(test_df.var2, test_df.var1)

Output (copy and paste from the python console):

var1  0  1
var2      
0     0  1
1     2  0
2     2  0

So, to summarize:

chi2_contingency(pandas.crosstab(test_df.var2, test_df.var1))

Answer 5

I personally can't answer your first concern.

---

For your second one, two list comprehensions would help.
List comprehensions are good when you have a list and a for loop to populate the list.

list_ = []
for i in range(10):
    list_.append(i)

To make it a list comprehension is rather easy.

list_ = [
    i
    for i in range(10)
]

While this is as very rudimentary example as you could just do list(range(10)), it is to show how it works simply.

---

If we change your code to use it then it is simpler.

def create_list_sum_of_categories(df, var, cat, var2):
    return [
        len(df[(df[var] == cat) & (df[var2] == cat2)])
        for cat2 in range(int(df[var2].min()), int(df[var2].max()) + 1)
    ]

def chi_square_of_df_cols(df,col1,col2):
    return scs.chi2_contingency([
        create_list_sum_of_categories(df,col1,cat,col2)
        for cat in range(int(df[col1].min()), int(df[col1].max())+1)
    ])

---

As you asked for ways to make it 'nicer looking'. Merging them into one function does makes it simpler, and it looks nicer.

def chi_square_of_df_cols(df,col1,col2):
    return scs.chi2_contingency([
        [
            len(df[(df[col1] == cat) & (df[col2] == cat2)])
            for cat2 in range(int(df[col1].min()), int(df[col1].max()) + 1)
        ]
        for cat in range(int(df[col2].min()), int(df[col2].max()) + 1)
    ])

Some style recommendations

• It is also recommended that code does not exceed 79 characters.
  The exception to this are comments and docstrings, at 72.

  And that is why I removed the docstring in the examples.

• Use descriptive variable names. data_file or category, rather than df and cat.

• Be careful with white-space append( len( is generally not accepted.
  If you do do this, then you should do the same amount of white space on both sides. As you did to index df. But it is best to avoid this.

Apart from the above your code is quite nice.