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
[array([0.02278028, 2.54848756]),
 array([2.61971259, 0.02711257]),
 array([1.25071779, 1.40954537])]
'''
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()