Plot witness f for Gaussian and Laplace PDFs

witness.py

#!/usr/bin/env python
import numpy as np
import pylab as pl
from scipy.spatial.distance import cdist

SAMPLES = 1.5e5
SIGMA = 0.2

def main():
    X = np.random.laplace(size=SAMPLES)
    Y = np.random.normal(size=SAMPLES)
    x = np.arange(-6.0, 6.0, 0.01)
    kernel = rbf(SIGMA)

    pl.figure()
    pl.plot(x, witness(X, Y, kernel, x), 'r-', lw=3)
    pl.plot(x, normal(x), 'k--', lw=3)
    pl.plot(x, laplace(x), 'k:', lw=3)
    pl.legend(('f', 'Gauss', 'Laplace'))
    pl.xlabel('X')
    pl.ylabel('PDFs and f')
    pl.show()

def normal(x, mu=0.0, sigma=1.0):
    """Compute the normal PDF"""
    normalization = 1.0 / (sigma * np.sqrt(2 * np.pi))
    return normalization * np.exp( -(x - mu)**2 / (2 * sigma**2))

def laplace(x, loc=0.0, scale=1.0):
    """Compute the laplace PDF"""
    return np.exp(-abs(x-loc) / scale)/(2 * scale)

def witness(X, Y, kernel, x, norm_sample=1e3):
    """
    Compute values of the witness function for two distributions
    at the values in x, given the particular kernel

    Arguments:
    X - 1-D data drawn according to P_X (array-like)
    Y - 1-D data drawn according to P_Y (array-like)
    kernel - kernel for RKHS from which witness is chosen
    x - values for which the witness should be computed (array-like)

    Returns:
    f(x) - array the same size as x
    """
    X = np.reshape(np.asarray(X), (len(X), 1))
    Y = np.reshape(np.asarray(Y), (len(Y), 1))
    x = np.reshape(np.asarray(x), (len(x), 1))
    numerator = (np.average(kernel(X, x), axis=0)
                 - np.average(kernel(Y, x), axis=0))

    # Compute normalization (downsample X and Y since
    # it's not that important for this constant)
    Xs = X[:norm_sample]
    Ys = Y[:norm_sample]
    normalization = np.sqrt(np.average(kernel(Xs, Xs))
                            - 2*np.average(kernel(Xs, Ys))
                            + np.average(kernel(Ys, Ys)))

    if normalization == 0.0:
        return numerator
    else:
        return numerator/normalization
    
def rbf(sigma):
    """Radial Basis Function Kernel"""
    def rbf_kernel(x, y):
        return np.exp(-cdist(x, y, 'sqeuclidean')/(2*sigma**2))
    return rbf_kernel

if __name__ == '__main__':
    main()