Peter's example of varbvs surface visualization
About
This is a Python reproduction of Peter's example of variational objective surface. See the original example here: https://pcarbo.github.io/pcarbo/visualize_varbvs_surface.html
import numpy as np
from mrashpen.inference.penalized_regression import PenalizedRegression as PLR
from mrashpen.inference.mrash_wrapR import MrASHR
from mrashpen.models.plr_ash import PenalizedMrASH
from mrashpen.models.normal_means_ash_scaled import NormalMeansASHScaled
from mrashpen.inference.ebfit import ebfit
from mrashpen.utils import R_utils
import sys
sys.path.append('/home/saikat/Documents/work/sparse-regression/simulation/eb-linreg-dsc/dsc/functions')
import simulate
import matplotlib.pyplot as plt
from matplotlib import cm
from mpl_toolkits.axes_grid1 import make_axes_locatable
from pymir import mpl_stylesheet
from pymir import mpl_utils
mpl_stylesheet.banskt_presentation()# (splinecolor = 'black')
import rpy2.robjects as robj
import rpy2.robjects.vectors as rvec
from rpy2.robjects.packages import importr
from rpy2.robjects import numpy2ri
numpy2ri.activate()
n = 40
xcor = 0.99
btrue = np.array([3, 0])
alpha0 = np.array([[0, 1],
[1, 0],
[1, 1]])
mu0 = np.array([[0, 3],
[2.5, 0],
[1, 1.5]])
S = np.array([[1, xcor], [xcor, 1]])
mu_seq = np.arange(-1, 3.51, 0.01)
'''
get exact same X and y with the same seed as used by Peter.
https://pcarbo.github.io/pcarbo/visualize_varbvs_surface.html
'''
seed = 1
robj.r('set.seed({})'.format(seed))
r_mass = importr('MASS')
X = r_mass.mvrnorm(n, rvec.FloatVector(np.zeros(2)), robj.r.matrix(S, nrow = 2, ncol = 2))
X = np.array(X)
X = X - np.mean(X, axis = 0)
yerr = robj.r.rnorm(n)
y = np.dot(X, btrue) + np.array(yerr)
y = y - np.mean(y)
def fit_varbvs(X, y, alpha, mu):
n = X.shape[0]
r_X = robj.r.matrix(X, nrow = n, ncol = 2)
r_y = rvec.FloatVector(y)
r_null = robj.NULL
r_alpha = rvec.FloatVector(alpha)
r_mu = rvec.FloatVector(mu)
r_varbvs = importr('varbvs')
fit1 = r_varbvs.varbvs(r_X, r_null, r_y, sigma = 1, sa = 1, logodds = 0, tol = 0,
alpha = r_alpha, mu = r_mu, verbose = False)
return R_utils.robj2dict(fit1)
def sigmoid(x):
return 1 / (1 + np.exp(-x))
def compute_alpha(mu, s):
return sigmoid(np.log(s)/2 + np.square(mu)/(2 * s))
def betavar(p, mu, s):
return p*(s + (1 - p) * np.square(mu))
def sqnorm2(x):
return np.sum(np.square(x), axis = 0)
def mdot(x, y):
return np.sum(x * y, axis = 0)
def compute_kl(X, y, a, mu, s):
'''
Calculate in grid, dimension 3.
'''
n = y.shape[0]
eps = 1e-15
d = np.sum(np.square(X), axis = 0).reshape(-1, 1, 1)
'''
Convert 1d entries to 3d.
There shouldn't be 2d entries.
'''
if a.ndim == 1: a = a[:, np.newaxis, np.newaxis]
if mu.ndim == 1: mu = mu[:, np.newaxis, np.newaxis]
if s.ndim == 1: s = s[:, np.newaxis, np.newaxis]
r = a * mu
v = betavar(a, mu, s)
ylessXr = y.reshape(-1, 1, 1) - np.einsum('ij, jkl->ikl', X, r)
kl = n * np.log(2 * np.pi) / 2 \
+ np.log(n) / 2 \
+ sqnorm2(ylessXr) / 2 \
+ mdot(d, v) / 2 \
+ np.sum(a * np.log(2 * a + eps) + (1 - a) * np.log(2 * (1 - a) + eps), axis = 0) \
- mdot(a, 1 + np.log(s) - (s + np.square(mu))) / 2
return kl
def compute_kl_grid(X, y, mugrid):
dX = np.sum(np.square(X), axis = 0).reshape(-1, 1, 1)
sgrid = 1 / (dX + 1)
alphagrid = compute_alpha(mugrid, sgrid)
klgrid = compute_kl(X, y, alphagrid, mugrid, sgrid)
betagrid = mugrid * alphagrid
return alphagrid, betagrid, klgrid
def compute_kl_grid_r(X, y, mu):
'''
mu is a grid of values, dimension (2, M, M)
M is the number of mu values on each axis.
'''
path="utils/visualize_varbvs_surface_functions.R"
r_func = robj.r
r_func.source(path)
#r_X = robj.r.matrix(X, nrow = X.shape[0], ncol = X.shape[1])
#r_mu_grid = robj.r.matrix(X, nrow = mu.shape[0], ncol = mu.shape[1])
r_mu = mu.reshape(2, -1).T
r_y = rvec.FloatVector(y)
p = r_func.compute_kl_grid(X, r_y, r_mu)
parr = np.array(p)
#return parr
'''
change array dimensions for use in Python
'''
ndim = int(np.sqrt(parr.shape[0]))
mu1 = parr[:, 0]
mu2 = parr[:, 1]
alpha1 = parr[:, 2]
alpha2 = parr[:, 3]
alphagrid = np.array([alpha1, alpha2])
betagrid = np.array([(mu1 * alpha1).reshape(ndim, ndim), (mu2 * alpha2).reshape(ndim, ndim)])
klgrid = parr[:, 4].reshape(ndim, ndim)
return alphagrid, betagrid, klgrid
fit1 = fit_varbvs(X, y, alpha0[0, :], mu0[0, :])
fit2 = fit_varbvs(X, y, alpha0[1, :], mu0[1, :])
fit3 = fit_varbvs(X, y, alpha0[2, :], mu0[2, :])
pts = list([fit1['beta'], fit2['beta'], fit3['beta']])
pts
'''
Python version
'''
mu1, mu2 = np.meshgrid(mu_seq, mu_seq)
mu_grid = np.stack((mu1, mu2))
alpha_grid, beta_grid, kl_grid \
= compute_kl_grid(X, y, mu_grid)
beta_grid = mu_grid * alpha_grid
'''
R version
'''
alpha_grid_r, beta_grid_r, kl_grid_r \
= compute_kl_grid_r(X, y, mu_grid)
# parr = compute_kl_grid_r(X, y, mu_grid)
def plot_contours(ax, X, Y, Z, betas, norm, cstep = 10,
xlabel = "", ylabel = "", zlabel = "",
show_beta = True, show_Zmax = False, show_Zmin = False,
show_cbar = False, show_cmap = True
):
zmin = np.min(Z) - 1 * np.std(Z)
zmax = np.max(Z) + 1 * np.std(Z)
if show_Zmin:
min_ind = np.unravel_index(np.argmin(Z, axis=None), Z.shape)
if show_Zmax:
max_ind = np.unravel_index(np.argmax(Z, axis=None), Z.shape)
levels = np.linspace(zmin, zmax, 200)
clevels = np.linspace(zmin, zmax, 50)
cmap = cm.YlOrRd_r
if show_cmap:
if norm:
cset1 = ax.contourf(X, Y, Z, levels, norm = norm,
cmap=cm.get_cmap(cmap, len(levels) - 1))
else:
cset1 = ax.contourf(X, Y, Z, levels,
cmap=cm.get_cmap(cmap, len(levels) - 1))
cset2 = ax.contour(X, Y, Z, clevels, colors='k')
else:
cset2 = ax.contour(X, Y, Z, clevels, colors='dodgerblue')
for c in cset2.collections:
c.set_linestyle('solid')
ax.set_aspect("equal")
if show_beta:
for beta in betas:
ax.scatter(beta[0], beta[1], color = 'k', marker = '+', s = 100, zorder = 100)
if show_Zmax:
ax.scatter(X[max_ind], Y[max_ind], color = 'k', s = 100, zorder = 100)
if show_Zmin:
ax.scatter(X[min_ind], Y[min_ind], color = 'k', s = 100, zorder = 100)
if xlabel:
ax.set_xlabel(xlabel)
if ylabel:
ax.set_ylabel(ylabel)
if show_cbar:
divider = make_axes_locatable(ax)
cax = divider.append_axes("right", size="5%", pad=0.2)
cbar = plt.colorbar(cset1, cax=cax)
ytickpos = np.arange(int(zmin / cstep) * cstep, zmax, cstep)
cbar.set_ticks(ytickpos)
if zlabel:
cax.set_ylabel(zlabel)
xx, yy = beta_grid
zz = np.log10(kl_grid- np.min(kl_grid) + 1e-8)
fig = plt.figure()
ax1 = fig.add_subplot(111)
norm = None
# norm = cm.colors.Normalize(vmin = np.min(zz), vmax=np.max(zz))
# norm = cm.colors.TwoSlopeNorm(vmin=np.min(ELBO), vcenter=np.max(ELBO)-100, vmax=np.max(ELBO))
plot_contours(ax1, xx, yy, zz, pts, norm, show_cbar = False, show_cmap = False)
plt.show()
