Add log_lik experiments.

import time

import numpy as np

from scipy import linalg as sp_linalg

from scipy import special as sp_special

from scipy import stats as sp_stats

def sqrt_2x2(A):

	s = np.sqrt(A[:,0,0]*A[:,1,1]-A[:,0,1]**2)

	t = np.sqrt(A[:,0,0]+A[:,1,1] + 2*s).reshape((-1,1,1))

	A[:,0,0] += s

	A[:,1,1] += s

	return A/t

def batch_matrix_sqrt(A):

	if A.shape[1] == 1:

		return np.sqrt(A)

	if A.shape[1] == 2:

		return sqrt_2x2(A)

	return np.stack([sp_linalg.sqrtm(A[i,:,:]) for i in range(n_A)], axis=0)

class bottleneck_nngp_prior(object):

	def __init__(self, X, v_b=1.0, v_w=1.0, v_n=1e-4, output_dims=1, depths=[0], widths=[], bottleneck_activation=False, bottleneck_noise=1e-9):

		self.n_inputs = X.shape[0]

		self.X = X

		self.v_b = v_b

		self.v_w = v_w

		self.v_n = v_n

		self.output_dims = output_dims

		self.depths = depths

		self.widths = widths

		self.bottleneck_activation = bottleneck_activation

		self.bottleneck_noise = bottleneck_noise

	def K(self, Gram, depth, noise=0.0):

		Gram = np.atleast_2d(Gram)

		if Gram.ndim == 2:

			Gram = np.expand_dims(Gram, 0)

		K_XX = self.v_b + self.v_w*Gram

		for _ in range(depth):

			A = np.sqrt(np.einsum("ijj,ikk->ijk", K_XX, K_XX))

			R = np.maximum(np.minimum(K_XX/A, 1.0), -1.0)

			K_XX = self.v_b + self.v_w*A/np.pi*(np.sqrt(1-R**2) + (np.pi-np.arccos(R))*R)

		return K_XX + noise*np.expand_dims(np.eye(self.n_inputs), 0)

	def sample(self, n_samples=1):

		grams = self.X.dot(self.X.T)

		for (depth, width) in zip(self.depths[:-1], self.widths):

			if not self.bottleneck_activation:

				sqrt_K_grams = batch_matrix_sqrt(self.K(grams, depth))

				if width >= self.n_inputs:

					wis = sp_stats.wishart(df=width, scale=np.eye(self.n_inputs))

					samples = wis.rvs((n_samples))/width

				else:

					samples = np.random.normal(0, 1, size=(n_samples, self.n_inputs, width))

					samples = np.einsum("sik,sjk->sij", samples, samples)/width

				if samples.ndim == 1:

					samples = samples.reshape((-1, 1, 1))

				grams = np.einsum("sik,skj->sij", sqrt_K_grams, samples)

				grams = np.einsum("sik,skj->sij", grams, sqrt_K_grams)

			else:

				sqrt_K_grams = batch_matrix_sqrt(self.K(grams, depth))

				samples = np.random.normal(0, 1, size=(n_samples, self.n_inputs, width))

				samples = np.maximum(0, np.einsum("sik,skj->sij", sqrt_K_grams, samples))

				grams = 2*np.einsum("sik,sjk->sij", samples, samples)/width

		sqrt_K_grams = batch_matrix_sqrt(self.K(grams, self.depths[-1]))

		samples = np.random.normal(0, 1, size=(n_samples, self.n_inputs, self.output_dims))

		Y = np.einsum("sik,skj->sij", sqrt_K_grams, samples)

		return Y

	def log_likelihood(self, Y, n_samples=1):

		grams = np.expand_dims(self.X.dot(self.X.T), 0)

		for (depth, width) in zip(self.depths[:-1], self.widths):

			Ks = self.K(grams, depth, noise=self.bottleneck_noise)

			Ls = np.stack([np.linalg.cholesky(K) for K in Ks], 0)

			samples = np.random.normal(0, 1, size=(n_samples, self.n_inputs, width))

			samples = np.einsum("sik,skj->sij", Ls, samples)

			if self.bottleneck_activation:

				samples = np.sqrt(2)*np.maximum(0, samples)

			grams = np.einsum("sik,sjk->sij", samples, samples)/width

		Ks = self.K(grams, self.depths[-1], noise=self.v_n)

		Ls = np.stack([np.linalg.cholesky(K) for K in Ks], 0)

		LYs = np.stack([sp_linalg.solve_triangular(L, Y, lower=True) for L in Ls], 0)

		logdets = 2*np.sum(np.log(np.diagonal(Ls, axis1=1, axis2=2)), 1)

		logliks = -1/2*np.sum(np.sum(LYs**2, 1), 1) - Y.shape[1]/2*logdets - Y.shape[1]*self.n_inputs/2*np.log(2*np.pi)

		loglik = sp_special.logsumexp(logliks) - np.log(n_samples)

		return loglik

import os

import argparse

import time

import pickle

import random

import itertools

import numpy as np

import pandas as pd

from sklearn import datasets

import models

random.seed(123)

np.random.seed(456)

parser = argparse.ArgumentParser(description="Convergence of bottleneck NNGP to NNGP in log-likelihood.")

parser.add_argument("--depths", "-d", dest="depths", type=lambda s: [int(i) for i in s.split(",")], default="1,1", help="comma-separated depths of bottleneck components.")

args = parser.parse_args()

# generate rings data

thetas = np.pi/6*np.arange(12)

zs = 0.25*np.arange(5)-0.5

X_1 = np.array([[np.cos(theta), np.sin(theta), z] for (theta, z) in itertools.product(thetas, zs)])

X_2 = X_1.dot(np.array([[0, 0, -1], [0, 1, 0], [1, 0, 0]]))

X_2[:,1] += 1

X = np.concatenate((X_1, X_2), axis=0)

Y = np.zeros((X.shape[0]))

Y[X.shape[0]//2:] = 1

Y = Y.reshape((-1, 1))

#Y = np.tile(Y, [1, 2])

# standardize

X = (X - np.mean(X, axis=0, keepdims=True))/np.std(X, axis=0, keepdims=True)

Y = (Y - np.mean(Y, axis=0, keepdims=True))/np.std(Y, axis=0, keepdims=True)

# add noise

Y = Y + np.random.normal(0, 1e-2, size=Y.shape)

n_samples = 10**3

widths = [1, 2, 3, 4, 5, 10, 50, 100, 500, np.inf]

depths = args.depths

# variance hyperparameters

v_b = 1.0

v_w = 1.0

v_n = 1e-4

liks = []

for width in widths:

	print("width:", width)

	ds = [sum(depths)+len(depths)-1] if width == np.inf else depths

	ws = [] if width == np.inf else [width]*(len(depths)-1)

	prior = models.bottleneck_nngp_prior(X, output_dims=Y.shape[1], v_b=v_b, v_w=v_w, v_n=v_n, depths=ds, widths=ws, bottleneck_activation=True)

	l = prior.log_likelihood(Y, n_samples=n_samples)/X.shape[0]

	liks.append(l)

liks = list(map(lambda x: np.round(x, 3), liks))

data = pd.DataFrame.from_dict({"width": widths, "lik": liks})[["width", "lik"]]

# write to CSV

os.makedirs("results/rings", exist_ok=True)

filename = "results/rings/depth"+"-".join(map(str, depths))+".csv"

data.to_csv(filename, index=False)

print("Done!")

 No newline at end of file