karlnapf / ds3_kernel_testing Goto Github PK

import numpy as np
from tqdm import tqdm_notebook as tqdm
from scipy.spatial.distance import squareform, pdist, cdist
import shogun as sg

data = np.load("blobs.npz")
X_learn = data["X"]
Y_learn = data["Y"]

def two_sample_permutation_test(test_statistic, X, Y, num_permutations, prog_bar=True):
    assert X.ndim == Y.ndim
    
    statistics = np.zeros(num_permutations)
    
    range_ = range(num_permutations)
    if prog_bar:
        range_ = tqdm(range_)
    for i in range_:
        # concatenate samples
        if X.ndim == 1:
            Z = np.hstack((X,Y))
        elif X.ndim == 2:
            Z = np.vstack((X,Y))
            
        # IMPLEMENT: permute samples and compute test statistic
        perm_inds = np.random.permutation(len(Z))
        Z = Z[perm_inds]
        X_ = Z[:len(X)]
        Y_ = Z[len(X):]
        my_test_statistic = test_statistic(X_, Y_)
        statistics[i] = my_test_statistic
    return statistics

def quadratic_time_mmd(X,Y,kernel):
    assert X.ndim == Y.ndim == 2
    K_XX = kernel(X,X)
    K_XY = kernel(X,Y)
    K_YY = kernel(Y,Y)
       
    n = len(K_XX)
    m = len(K_YY)
    
    # IMPLEMENT: unbiased MMD statistic (could also use biased, doesn't matter if we use permutation tests)
    np.fill_diagonal(K_XX, 0)
    np.fill_diagonal(K_YY, 0)
    mmd = np.sum(K_XX) / (n*(n-1))  + np.sum(K_YY) / (m*(m-1))  - 2*np.sum(K_XY)/(n*m)
    return mmd


def gauss_kernel(X, Y=None, sigma=1.0):
    """
    Computes the standard Gaussian kernel k(x,y)=exp(- ||x-y||**2 / (2 * sigma**2))

    X - 2d array, samples on left hand side
    Y - 2d array, samples on right hand side, can be None in which case they are replaced by X
    
    returns: kernel matrix
    """

    # IMPLEMENT: compute squared distances and kernel matrix
    sq_dists = sq_distances(X,Y)
    K = np.exp(-sq_dists / (2 * sigma**2))
    return K

def sq_distances(X,Y=None):
    assert(X.ndim==2)

    # IMPLEMENT: compute pairwise distance matrix. Don't use explicit loops, but the above scipy functions
    # if X=Y, use more efficient pdist call which exploits symmetry
    if Y is None:
        sq_dists = squareform(pdist(X, 'sqeuclidean'))
    else:
        assert(Y.ndim==2)
        assert(X.shape[1]==Y.shape[1])
        sq_dists = cdist(X, Y, 'sqeuclidean')

    return sq_dists


log_sigma=-2
num_permutations=200

# Shogun implementation
feat_p=sg.RealFeatures(X_learn.T.astype(np.float64))
feat_q=sg.RealFeatures(Y_learn.T.astype(np.float64))
kernel=sg.GaussianKernel(2 * (10**log_sigma)**2)

mmd=sg.QuadraticTimeMMD(feat_p,feat_q)
mmd.set_kernel(kernel)

mmd.set_statistic_type(sg.ST_UNBIASED_FULL)
statistic=mmd.compute_statistic()

mmd.set_null_approximation_method(sg.NAM_PERMUTATION)
mmd.set_num_null_samples(num_permutations)


# DS3 summer school implementation
my_kernel = lambda X,Y : gauss_kernel(X,Y,sigma=10**log_sigma)
my_mmd = lambda X,Y : quadratic_time_mmd(X,Y, my_kernel)
my_statistic = my_mmd(X_learn, Y_learn)
statistics = two_sample_permutation_test(my_mmd, X_learn, Y_learn, num_permutations, prog_bar=False)
p_value = np.mean(my_statistic <= np.sort(statistics))

print(statistic)
print(my_statistic)
print(mmd.compute_p_value(statistic))
print(p_value)