<jupyter_start><jupyter_text>Logistic regression model  Setup<jupyter_code>!git clone --depth 1 https://github.com/probml/pyprobml /pyprobml &> /dev/null
%cd -q /pyprobml/scripts
!pip install -q optax
!pip install -q blackjax
!pip install -q sgmcmcjax
%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np

import itertools
import warnings
from functools import partial

import jax
import jax.numpy as jnp 
from jax.random import uniform, normal, bernoulli, split
from jax import jit, grad, value_and_grad, vmap
from jax.experimental import optimizers
from jax.scipy.special import logsumexp

from blackjax import nuts, stan_warmup
import optax

import sgmcmc_subspace_lib as sub
from sgmcmc_utils import build_optax_optimizer, build_nuts_sampler
from sgmcmcjax.samplers import *<jupyter_output><empty_output><jupyter_text>Generate Data This part is based on https://github.com/jeremiecoullon/SGMCMCJax/blob/master/docs/nbs/models/logistic_regression.py<jupyter_code>#ignore by GPU/TPU message (generated by jax module)
warnings.filterwarnings("ignore", message='No GPU/TPU found, falling back to CPU.')


# Sample initial beta values from random normal
def init_params(rng_key, d):
    return normal(rng_key, (d,))

def gen_cov_mat(key, d, rho):
    Sigma0 = np.diag(np.ones(d))
    for i in range(1,d):
        for j in range(0, i):
            Sigma0[i,j] = (uniform(key)*2*rho - rho)**(i-j)
            Sigma0[j,i] = Sigma0[i,j]

    return jnp.array(Sigma0)

def logistic(theta, x):
    return 1/(1+jnp.exp(-jnp.dot(theta, x)))

def gen_data(key, dim, N):
    """
    Generate data with dimension `dim` and `N` data points
    Parameters
    ----------
    key: uint32
        random key
    dim: int
        dimension of data
    N: int
        Size of dataset
    Returns
    -------
    theta_true: ndarray
        Theta array used to generate data
    X: ndarray
        Input data, shape=(N,dim)
    y: ndarray
        Output data: 0 or 1s. shape=(N,)
    """
    key, theta_key, cov_key, x_key = split(key, 4)
    rho = 0.4
    print(f"Generating data, with N={N} and dim={dim}")
    theta_true = normal(theta_key, shape=(dim, ))*jnp.sqrt(10)
    covX = gen_cov_mat(cov_key, dim, rho)
    X = jnp.dot(normal(x_key, shape=(N,dim)), jnp.linalg.cholesky(covX))

    p_array = batch_logistic(theta_true, X)
    keys = split(key, N)
    y = batch_benoulli(keys, p_array).astype(jnp.int32)
    return theta_true, X, y
@jit
def predict(params, inputs):
  return batch_logistic(params, inputs) > 0.5

@jit
def accuracy(params, batch):
  inputs, targets = batch
  predicted_class = predict(params, inputs)
  return jnp.mean(predicted_class == targets)

@jit
def loglikelihood(theta, x_val, y_val):
    return -logsumexp(jnp.array([0., (1.-2.*y_val)*jnp.dot(theta, x_val)]))

@jit
def logprior(theta):
    return -(0.5/10)*jnp.dot(theta,theta)
batch_logistic = jit(vmap(logistic, in_axes=(None, 0)))
batch_benoulli = vmap(bernoulli, in_axes=(0, 0))
batch_loglikelihood = vmap(loglikelihood, in_axes=(None, 0, 0))
dim = 10 # Choose a dimension for the parameters (10, 50,100 in the paper)
subspace_dim = 2 # Choose a dimension for the subspace parameters
ndata = 10000 # Number of data points
nwarmup = 1000 # Number of iterations during warmup phase
nsamples = 10000 # Number of SGMCMC iterations
nsamplesCV = nsamples // 2 
key = jax.random.PRNGKey(42)
theta_true, X, y = gen_data(key, dim, ndata)
batch_size = int(0.01*X.shape[0])
data = (X, y)

init_key, key = split(key)
theta_init = init_params(init_key, dim)<jupyter_output>Generating data, with N=10000 and dim=10<jupyter_text>SGD<jupyter_code>niters = 5000

learning_rate = 6e-5
opt = optax.sgd(learning_rate=learning_rate)
optimizer = build_optax_optimizer(opt, loglikelihood, logprior,
                                  data, batch_size, pbar=False)

opt_key, key = split(key)
sgd_params, logpost_array = optimizer(opt_key, niters, theta_init)
train_accuracy = accuracy(sgd_params, data)
print("Training set accuracy {}".format(train_accuracy))
plt.plot(-logpost_array, color='tab:orange')
plt.show()<jupyter_output><empty_output><jupyter_text>Subspace Model<jupyter_code>sub_opt_key, key = split(key)
sgd_sub_params, _, opt_log_post_trace, _ = sub.subspace_optimizer(
    sub_opt_key, loglikelihood, logprior, theta_init, data, batch_size,
    subspace_dim, nwarmup, nsamples, opt, pbar=False)
train_accuracy = accuracy(sgd_sub_params, data)
print("Training set accuracy {}".format(train_accuracy))
plt.plot(-opt_log_post_trace, color='tab:pink')
plt.show()<jupyter_output><empty_output><jupyter_text>NUTS<jupyter_code>nuts_sampler = build_nuts_sampler(nwarmup, loglikelihood, logprior,
                                  data, batch_size=ndata, pbar=False)
nuts_key, key = split(key)
nuts_params = nuts_sampler(nuts_key, nsamples//10, theta_init)
train_accuracy = accuracy(jnp.mean(nuts_params, axis=0), data)
print("Training set accuracy {}".format(train_accuracy))<jupyter_output>Training set accuracy 0.9384999871253967<jupyter_text>Subspace Model<jupyter_code>build_nuts_sampler_partial = partial(build_nuts_sampler, nwarmup=nwarmup)
nuts_key, key = split(key)
nuts_sub_params = sub.subspace_sampler(nuts_key, loglikelihood, logprior,
                                       theta_init, build_nuts_sampler_partial,
                                       data, batch_size, subspace_dim, nsamples,
                                       nsteps_full=0, nsteps_sub=0, use_cv=False, opt=None, pbar=False)<jupyter_output><empty_output><jupyter_text>SGLD<jupyter_code>dt = 1e-5
# Run sampler
sgld_sampler = build_sgld_sampler(dt, loglikelihood, logprior, data, batch_size)
sgld_key, key = split(key)
sgld_output = sgld_sampler(sgld_key, nsamples, theta_init)<jupyter_output><empty_output><jupyter_text>Subspace Sampler<jupyter_code>build_sgld_sampler_partial = partial(build_sgld_sampler, dt=dt)
sgld_key, key = split(key)
sgld_sub_output = sub.subspace_sampler(sgld_key, loglikelihood, logprior,
                                       theta_init, build_sgld_sampler_partial,
                                       data, batch_size, subspace_dim, nsamples, 
                                       nsteps_full=0, nsteps_sub=0, use_cv=False, 
                                       opt=None, pbar=False)<jupyter_output><empty_output><jupyter_text>SGLDCV<jupyter_code>dt = 1e-5
sgldCV_sampler = build_sgldCV_sampler(dt, loglikelihood, logprior,
                                      data, batch_size, sgd_params)
sgldCV_key, key = split(key)
sgldCV_output = sgldCV_sampler(sgldCV_key, nsamplesCV, sgd_params)<jupyter_output><empty_output><jupyter_text>Subspace Sampler<jupyter_code>build_sgldCV_sampler_partial = partial(build_sgldCV_sampler, dt=dt)
sgldCV_key, key = split(key)
sgldCV_sub_output = sub.subspace_sampler(sgldCV_key, loglikelihood, logprior,
                                         theta_init, build_sgldCV_sampler_partial,
                                         data, batch_size, subspace_dim, nsamples, 
                                         nsteps_full=niters//2, nsteps_sub=niters//2, 
                                         use_cv=True, opt=opt, pbar=False)<jupyter_output><empty_output><jupyter_text>SGHMC<jupyter_code>L = 5
dt = 1e-6
sghmc_sampler = build_sghmc_sampler(dt, L, loglikelihood,
                                    logprior, data, batch_size)
sghmc_key, key = split(key)
sghmc_output = sghmc_sampler(sghmc_key, nsamples, theta_init)<jupyter_output><empty_output><jupyter_text>Subspace Sampler<jupyter_code>build_sghmc_sampler_partial = partial(build_sghmc_sampler, dt=dt, L=L)
sghmc_key, key = split(key)
sghmc_sub_output = sub.subspace_sampler(sghmc_key, loglikelihood, logprior,
                                        theta_init, build_sghmc_sampler_partial,
                                        data, batch_size, subspace_dim,
                                        nsamples, nsteps_full=0, nsteps_sub=0,
                                        use_cv=False, opt=None, pbar=False)<jupyter_output><empty_output><jupyter_text>SGHMCCV<jupyter_code>dt = 1e-7    # step size parameter
L = 5
sghmcCV_sampler = build_sghmcCV_sampler(dt, L, loglikelihood, logprior, data, batch_size, sgd_params)
sghmcCV_key, key = split(key)
sghmcCV_output = sghmcCV_sampler(sghmcCV_key, nsamplesCV, sgd_params)<jupyter_output><empty_output><jupyter_text>Subspace Sampler<jupyter_code>build_sghmcCV_sampler_partial = partial(build_sghmcCV_sampler, dt=dt, L=L)
sghmcCV_key, key = split(key)
sghmcCV_sub_output = sub.subspace_sampler(sghmcCV_key, loglikelihood, logprior,
                                          theta_init, build_sghmcCV_sampler_partial,
                                          data, batch_size, subspace_dim, nsamples,
                                          nsteps_full=niters//2, nsteps_sub=niters//2, 
                                          use_cv=True, opt=opt, pbar=False)<jupyter_output><empty_output><jupyter_text>SGNHT<jupyter_code>dt = 1e-6    # step size parameter
a = 0.02
sgnht_sampler = build_sgnht_sampler(dt, loglikelihood, logprior,
                                    data, batch_size, a=a)
sgnht_key, key = split(key)
sgnht_output = sgnht_sampler(sgnht_key, nsamples, theta_init)<jupyter_output><empty_output><jupyter_text>Subspace Sampler<jupyter_code>build_sgnht_sampler_partial = partial(build_sgnht_sampler, dt=dt, a=a)
sgnht_key, key = split(key)
sgnht_sub_output = sub.subspace_sampler(sgnht_key, loglikelihood, logprior,
                                        theta_init, build_sgnht_sampler_partial,
                                        data, batch_size, subspace_dim, nsamples,
                                        nsteps_full=0, nsteps_sub=0, use_cv=False,
                                        opt=opt, pbar=False)<jupyter_output><empty_output><jupyter_text>SGHNTCV<jupyter_code>dt = 1e-6     # step size parameter
a = 0.02
sgnhtCV_sampler = build_sgnhtCV_sampler(dt, loglikelihood, logprior,
                                        data, batch_size, sgd_params, a=a)
sgnhtCV_key, key = split(key)
sgnhtCV_output = sgnhtCV_sampler(sgnhtCV_key, nsamplesCV, sgd_params)<jupyter_output><empty_output><jupyter_text>Subspace Sampler<jupyter_code>build_sgnhtCV_sampler_partial = partial(build_sgnhtCV_sampler, dt=dt, a=a)
sgnhtCV_key, key = split(key)
sgnhtCV_sub_output = sub.subspace_sampler(sgnhtCV_key, loglikelihood, logprior,
                                          theta_init, build_sgnhtCV_sampler_partial,
                                          data, batch_size, subspace_dim, nsamples,
                                          nsteps_full=niters//2, nsteps_sub=niters//2,
                                          use_cv=True, opt=opt, pbar=False)<jupyter_output><empty_output><jupyter_text>ULA - SGLD with the full dataset<jupyter_code>dt = 4e-5  # step size parameter
ula_sampler = build_sgld_sampler(dt, loglikelihood, logprior,
                                 data, batch_size=ndata)
ula_key, key = split(key)
ula_output = ula_sampler(ula_key, nsamples, theta_init)<jupyter_output><empty_output><jupyter_text>Subspace Sampler<jupyter_code>build_sgld_sampler_partial = partial(build_sgld_sampler, dt=dt)
ula_key, key = split(key)
ula_sub_output = sub.subspace_sampler(ula_key, loglikelihood, logprior,
                                      theta_init, build_sgld_sampler_partial,
                                      data, ndata, subspace_dim, nsamples,
                                      nsteps_full=0, nsteps_sub=0, use_cv=False,
                                      opt=None, pbar=False)<jupyter_output><empty_output><jupyter_text>Trace plots<jupyter_code>def trace_plot(outs):
  nrows, ncols = 2, 4
  fig, axes = plt.subplots(nrows=nrows, ncols=ncols, figsize=(16, 12))

  for ax, (title, out) in zip(axes.flatten(), outs.items()):
    ax.plot(out)
    ax.set_title(title)
    ax.set_xlabel("Iteration")
    ax.grid(color='white', linestyle='-', linewidth=2)
    ax.set_axisbelow(True)
    ax.set_facecolor("#EAEBF0")

  plt.tight_layout()
  plt.savefig("traceplot.pdf", dpi=300)
  plt.show()
outs = {"STAN": nuts_params, "SGLD": sgld_output, "SGLDCV": sgldCV_output, "ULA": ula_output, "SGHMC": sghmc_output, "SGHMCCV":sghmcCV_output, "SGNHT": sgnht_output, "SGNHTCV": sgnhtCV_output}
trace_plot(outs)
subspace_outs = {"STAN": nuts_sub_params, "SGLD": sgld_sub_output, "SGLDCV": sgldCV_sub_output, "ULA": ula_sub_output, "SGHMC": sghmc_sub_output, "SGHMCCV":sghmcCV_sub_output, "SGNHT": sgnht_sub_output, "SGNHTCV": sgnhtCV_sub_output}
trace_plot(subspace_outs)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.10194020356234097, 'fraction_numerical': 0.01184796437659033, 'mean_word_length': 3.377654020187957, 'pattern_counts': {'":': 16, '<': 83, '<?xml version=': 0, '>': 85, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 50, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '21', 'dir': 'jupyter-structured-clean-dedup', 'id': '14288', 'license': 'MIT', 'n_tokens_mistral': 4635, 'n_tokens_neox': 4403, 'n_words': 810, 'path': 'notebooks/logistic_regression_sgmcmc.ipynb', 'repo_name': 'susnato/probml-notebooks', 'size': '12576'}