A/B testing using chi-square to calculate the significance in an elegant way

Question

The definition of ABtest

Goal:

For input data, each pair row is a group such as 0 and 1, 2 and 3. I want to test whether each group is significant using chi-square.

• If result is significant and any value in group is more than or equals 5 , it returns p-value, else it returns 'not significant'.
• If any value in group is less than 5, it returns 'sample size too small'.

My question is:

• My code seems so long and not so elegant and want to be more simple.

Notice: please don't use column grp as raw input to rewrite like groupby method because the raw data is not named a1 ca1, b1 cb1 . But I make sure each pair row (e.g. 0 and 1, 2 and 3) is a group.

Code:

import numpy as np
from scipy.stats import chi2_contingency
from itertools import chain, repeat

def abtest(df,tot,convert):
    df = df.copy()
    df['nonconvert'] = df[tot] - df[convert]
    grp = np.split(df.index.values,df.index.size//2)
    rst = []
    for g in grp:
        obs = df.loc[g][[convert,'nonconvert']].values
        if (obs>=5).all():
             _, p, _, _=chi2_contingency(obs)
             if p<0.05:
                rst.append(p)
             else:
                rst.append('not significant')
        else:
            rst.append('sample size too small')
    rate = tot + '_' + convert + '_'+'test'
    df[rate] = list(chain.from_iterable(zip(*repeat(rst, 2))))
    del df['nonconvert']
    return df 

df = abtest(df=df,tot='tot',convert='read')
df = abtest(df=df,tot='tot',convert='lgn')
df = abtest(df=df,tot='lgn',convert='read')

Input:

   grp    tot  lgn  read
0   a1   6300  360    30
1  ca1   2300   60     7
2   b1  26300  148     6
3  cb1  10501   15     3
4   c1  74600   36     2
5  cc1  26000    6     1

Output:

   grp    tot  lgn  read          tot_read_test     tot_lgn_test  \
0   a1   6300  360    30        not significant      4.68208e-09   
1  ca1   2300   60     7        not significant      4.68208e-09   
2   b1  26300  148     6  sample size too small      7.01275e-08   
3  cb1  10501   15     3  sample size too small      7.01275e-08   
4   c1  74600   36     2  sample size too small  not significant   
5  cc1  26000    6     1  sample size too small  not significant   

           lgn_read_test  
0        not significant  
1        not significant  
2  sample size too small  
3  sample size too small  
4  sample size too small  
5  sample size too small 

Answer 1

When you have a finite number of members in a group A and B. Instead of split into groups, hstack the DataFrame like this:

pd.concat(
    [
        df[df.index % 2 == 0].add_prefix('a_').reset_index(),
        df[df.index % 2 == 1].add_prefix('b_').reset_index()
    ], axis=1
)

    a_grp   a_tot   a_lgn   a_read  index   b_grp   b_tot   b_lgn   b_read
0   0   a1  6300    360 30  1   ca1 2300    60  7
1   2   b1  26300   148 6   3   cb1 10501   15  3
2   4   c1  74600   36  2   5   cc1 26000   6   1

Now you can replace the for-loop with 'a_' and 'b_' and df.apply() like

df.apply(apply_chi2)

def apply_chi2(df_ab):
  if df_ab['a_'+convert] > df_ab['a_'+'nonconvert']:
     return ...
  obs = df_ab.a
  return _, p, _, _=chi2_contingency(obs)

Answer 2

One minor thing that you can do is make your function reduce-friendly:

import numpy as np
from scipy.stats import chi2_contingency
from functools import reduce
from itertools import chain, repeat

def abtest(df, args):
    tot, convert = args
    df = df.copy()
    df['nonconvert'] = df[tot] - df[convert]
    grp = np.split(df.index.values,df.index.size//2)
    rst = []

    for g in grp:
        obs = df.loc[g][[convert,'nonconvert']].values
        if (obs>=5).all():
             _, p, _, _=chi2_contingency(obs)
             if p<0.05:
                rst.append(p)
             else:
                rst.append('not significant')
        else:
            rst.append('sample size too small')

    rate = tot + '_' + convert + '_'+'test'
    df[rate] = list(chain.from_iterable(zip(*repeat(rst, 2))))
    del df['nonconvert']
    return df 

df = reduce(abtest,
  (
    ('tot', 'read'),
    ('tot', 'lgn'),
    ('lgn', 'read')
  ),
  df
)