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:

         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):

    return scs.chi2_contingency(result_list)

test_df = pd.read_csv('test_data_for_chi_square.csv')

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!

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)

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

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')
  • \$\begingroup\$ Great solution! I'll test it next time I touch the code (in the not too distant future ;) )! Thanks a lot. \$\endgroup\$ – Klaster Aug 31 '15 at 21:18
  • 1
    \$\begingroup\$ instead of groupby you can use df.pivot_table(index=['var1', 'var2'], aggfunc=len) \$\endgroup\$ – Ruggero Turra Oct 27 '15 at 0:08

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


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.


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):

To make it a list comprehension is rather easy.

list_ = [
    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([
        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.


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):

         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
0     0  1
1     2  0
2     2  0

So, to summarize:

chi2_contingency(pandas.crosstab(test_df.var2, test_df.var1))
  • 3
    \$\begingroup\$ Welcome to Code Review! Please add some more information to your post. What is this code doing? Why? What is this replacing in the OP's code? What is better about this code? Your post is of very low quality right now. \$\endgroup\$ – SirPython Oct 27 '15 at 0:35
  • \$\begingroup\$ sorry, but I don't understand the downvotes. My answer answers exactly to the question. I have expanded the answer, but still I don't understand why you need such verbosity. To answer to your question: 1. the code is doing what requested. 2. why? 3. everything 4. shorter, faster, stronger... I don't see any difference from my original answer and the one from pbarber, except the fact my code is shorter \$\endgroup\$ – Ruggero Turra Oct 27 '15 at 12:42
  • \$\begingroup\$ If you look at the other answers I think the function you mention is only part of the solution, not the complete thing. Which is fine IMO. \$\endgroup\$ – ferada Oct 27 '15 at 12:43
  • \$\begingroup\$ the last part is to call chi2_contingency, do I really have to specify it? \$\endgroup\$ – Ruggero Turra Oct 27 '15 at 12:45

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.