Two sided permutation test

permutation_test.py

import numpy as np
def permutation_test(a, b, n=1000):
    """
    Two-sided permutation test
    Input:
    | {a, b}: 1xm arrays of data
    | n: Number of permutations
    Output:
    | p: Two-sided pvalue: mean of |mean(a) - mean(b)| > |mean(perm(a)) - mean(perm((b)))|_i for i permutations
    
    By: Michael Silverstein
    Adapted from: All of Statistics pg 163 by Larry Wasserman
    """
    # Ensure arrays
    a = np.array(a)
    b = np.array(b)
    if (a.size == 0) | (b.size == 0):
        return np.nan
    # Test statistic
    tobs = np.abs(a.mean() - b.mean())
    # Combine into single vector
    combo = np.concatenate((a, b))
    # Permute n times
    perms = np.array([np.random.permutation(combo) for _ in range(n)])
    # Observed statistic for each permutation
    Ts = np.abs(perms[:, :len(a)].mean(1) - perms[:, len(a):].mean(1))
    # Compute pval
    p = (Ts > tobs).mean()
    return p