# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Logistic model quasi-Laplace with Mr.Ash
#
# - toc: true 
# - badges: true
# - comments: true
# - categories: [jupyter]

# ### About
#
# This is a proof of concept implementation of [quasi-Laplace approximation]({% post_url 2020-06-22-intuition-for-quasi-Laplace %}) with Mr.Ash. The quasi-Laplace approximation introduces a regularizer on the likelihood and approximates the posterior with a Gaussian distribution. This approach can be used for generalized linear model with any link function. 
#
# Bayesian logistic regression with a Gaussian prior (ridge regression) uses the Laplace approximation to approximate the posterior with a Gaussian distribution. The variational treatment of logistic model by Jordan and Jaakkola (2000) also leads to a Gaussian approximation of the posterior distribution. The greater flexibility of the variational approximation leads to improved accuracy compared to Laplace method.
#
# Here, I compare the quasi-Laplace approach with cross-validated Lasso ($L_1$ penalty). I used the $R^2$ statistic by Tjur{% fn 1 %} to quantify the quality of classification, and the Gini index{% fn 2 %} to quantify the sparsity of the coefficients.
#
# > Note: Coefficients are generally underestimated by QL irrespective of the choice of regularizer. Lasso overestimates the coefficients.  
#
#
#
# I found that the optimization depends on the choice of Gaussian mixture family.

# +
#collapse_hide

import numpy as np
np.set_printoptions(precision = 4, suppress=True)
import pandas as pd

from scipy import optimize
from scipy import stats
from scipy.special import digamma
from sklearn import linear_model

import matplotlib.pyplot as plt
from matplotlib import cm
from matplotlib import ticker as plticker
from mpl_toolkits.axes_grid1 import make_axes_locatable

import sys
sys.path.append("../../utils/")
import mpl_stylesheet
mpl_stylesheet.banskt_presentation(fontfamily = 'latex-clearsans', fontsize = 18, colors = 'banskt', dpi = 72)
# -

# ### Generate toy data
#
# Let us consider a logistic model with sparse coefficients, so that the number of causal variables `ncausal` is much less than the number of variables `nvar` in the model. To ensure sparsity of the coefficients, we use a Gaussian mixture prior of `nGcomp` components with variances given by $\sigma_k^2$ (`sigmak2`) and unknown mixture coefficients. The sparsity of the initialization is controlled by the variable `sparsity` that specifies the mixture coefficient of the zero-th component $\mathcal{N}(0, 0)$ (or the delta function). We are particularly interested in systems where the number of samples `nsample` is less than `nvar`. Validation is performed on separate test dataset with `ntest` samples.

# +
nsample = [50, 200]
nvar = [100, 100]
ncausal = 5
nGcomp = 10
sparsity = 0.95
prior_variance = 5
ntest = 2000
hg2 = 0.8

sigmak2 = np.array([prior_variance * np.square(np.power(2, (i)/nGcomp) - 1) for i in range(nGcomp)])


# -

# The data is generated from a liability threshold model such that the explanatory variables explain only a given fraction `hg2` of the response variable. $\mathbf{X}$ is centered and scaled such that for each variable $j$, the variance is $\mathrm{var}(\mathbf{x}_j) = 1$.

# +
# collapse-hide

def standardize(X):
    Xnorm = (X - np.mean(X, axis = 0)) 
    Xstd =  Xnorm / np.std(Xnorm, axis = 0)
    return Xstd

def logistic_data(X, beta):
    Xbeta = np.dot(X, beta)
    pred = 1 / (1 + np.exp(-Xbeta))
    Y = np.random.binomial(1, pred)
    return Y, pred

def liability_model(X, beta, prevalence = 0.5):
    hg2 = np.sum(np.square(beta))
    eff = np.dot(X, beta)
    rand_std = np.sqrt(1 - hg2)
    rand_eff = np.random.normal(0, rand_std, X.shape[0])
    liability = eff + rand_eff
    cutoff = stats.norm.ppf(1 - prevalence)
    cases = np.where(liability >= cutoff)[0]
    ctrls = np.where(liability <  cutoff)[0]
    Y = np.zeros(X.shape[0])
    Y[cases] = 1
    return Y, liability

def tjur_R2(X, Y, beta):
    Xbeta = np.dot(X, beta)
    pred = 1 / (1 + np.exp(-Xbeta))
    R2 = np.mean(pred[Y == 1]) - np.mean(pred[Y == 0])
    return R2

def gini(x):
    n = len(x)
    xabs = abs(x)
    xsort = xabs[np.argsort(xabs)]
    t1 = xsort / np.sum(xabs)
    t2 = (n - np.arange(n) - 1 + 0.5) / n
    gini = 1 - 2 * np.sum(t1 * t2)
    return gini

def simulate_data(nsample, nvar, ncausal, ntest, hg2):
    X = np.random.rand(nsample * nvar).reshape(nsample, nvar)
    X = standardize(X)
    Xtest = np.random.rand(ntest * nvar).reshape(ntest, nvar)
    Xtest = standardize(Xtest)

    beta = np.zeros(nvar)
    causal_idx = np.random.choice(nvar, ncausal, replace = False)
    beta[causal_idx] = np.random.normal(loc = 0., scale = 1., size = ncausal)
    beta = beta * np.sqrt(hg2 / np.sum(np.square(beta)))

    Y, Y_liab = liability_model(X, beta)
    Ytest, Ytest_liab = liability_model(Xtest, beta)
    #Y, Y_liab = logistic_data(X, beta)
    #Ytest, Ytest_liab = logistic_data(Xtest, beta)
    
    return X, Y, beta, Xtest, Ytest, Y_liab, Ytest_liab


# -

# Let us have a look at the generated data.

# +
# collapse-hide

X =          [None for n in nsample]
Xtest =      [None for n in nsample]
Y =          [None for n in nsample]
Ytest =      [None for n in nsample]
Y_liab =     [None for n in nsample]
Ytest_liab = [None for n in nsample]
beta =       [None for n in nsample]

for i, n in enumerate(nsample):
    X[i], Y[i], beta[i], Xtest[i], Ytest[i], Y_liab[i], Ytest_liab[i] \
        = simulate_data(n, nvar[i], ncausal, ntest, hg2)

print(f'There are {ncausal} non-zero coefficients with heritability {hg2:.4f}')

fig = plt.figure(figsize = (12,6))
ax = [fig.add_subplot(121),
      fig.add_subplot(122)]

for i, n in enumerate(nsample):
    Xbeta = np.dot(X[i], beta[i])
    pred = 1 / (1 + np.exp(-Xbeta))
    ax[i].scatter(Xbeta, Y[i], s = 10)
    ax[i].scatter(Xbeta, pred, s = 10)

    ax[i].set_xlabel(r'$\sum_i X_{ni} \beta_i$')
    ax[i].set_ylabel(r'$Y_n$')
    
plt.tight_layout()
plt.show()


# -

# ### Implementation of quasi-Laplace with variational approximation
#
# See Latex notes for details of the theory. Current implementation focus on readability and debug. Computational efficiency has been ignored.
#
# > Warning: ELBO decreasing with parameter update.
#
#
#
# There must be an error in theory or implementation. Debug!

# +
# collapse-hide

def regularized_log_likelihood(beta, X, Y, L):
    nvar = beta.shape[0]
    Xbeta = np.dot(X, beta)

    ## Function
    llb = np.sum(np.dot(Y, Xbeta) - np.log(1 + np.exp(Xbeta)))
    reg = - 0.5 * np.einsum('i,i->i', np.square(beta), L)
    loglik = llb + reg
    
    ## Gradient
    pred = 1 / (1 + np.exp(-Xbeta))
    der = np.einsum('i,ij->j', (Y - pred), X) - np.multiply(beta, L)
    
    return -loglik, -der

def precisionLL(X, beta, L):
    nvar = X.shape[1]
    Xbeta = np.einsum('i,ji->j', beta, X)
    pred = 1 / (1 + np.exp(-Xbeta))
    Sinv = np.einsum('i,i,ij,ik->jk', pred, (1 - pred), X, X)
    Sinv[np.diag_indices(nvar)] += L
    return Sinv

def get_regularized_mode(X, Y, beta0, L):
    nsample, nvar = X.shape
    args  = X, Y, L
    gmode = optimize.minimize(regularized_log_likelihood,
                              beta0,
                              args=args,
                              method='L-BFGS-B',
                              jac=True,
                              bounds=None,
                              options={'maxiter': 20000000,
                                       'maxfun': 20000000,
                                       'ftol': 1e-9,
                                       'gtol': 1e-9
                                       #'disp': True
                                      })
    return gmode.x

def get_elbo_qL(alpha, mu, s2, pi, sigmak2, XPX, SinvM, M, L, logdetS, debug = False):
    Ebeta = np.sum(np.multiply(alpha, mu)[:, 1:], axis = 1)
    mu2 = np.square(mu)
    t1 = - 0.5 * np.dot(Ebeta.T, np.dot(XPX, Ebeta))
    t2 = np.dot(Ebeta.T, SinvM)
    t3 = 0
    t4 = 0
    for j in range(nvar):
        Ebeta2j = 0
        for k in range(nGcomp):
            mom2jk = mu2[j, k] + s2[j, k]
            Ebeta2j += alpha[j, k] * mom2jk
            t4jk = 0
            if alpha[j, k] != 0:
                t4jk += np.log(pi[k])
                t4jk += - np.log(alpha[j, k])
                if k > 0:
                    t4jk += 0.5 * np.log(s2[j, k])
                    t4jk += - 0.5 * np.log(sigmak2[k])
                    t4jk += - 0.5 * mom2jk / sigmak2[k]
                    t4jk += 0.5
                t4jk *= alpha[j, k]
            t4 += t4jk
        t3 += 0.5 * XPX[j, j] * (Ebeta[j] * Ebeta[j] - Ebeta2j)
    logdetL = np.sum(np.log(L))
    t0 = - 0.5 * (logdetL + logdetS + np.dot(M.T, SinvM))
    elbo = t0 + t1 + t2 + t3 + t4
    if debug:
        print(t0, t1 + t2, t3, t4)
    return elbo

#def get_elbo(alpha, mu, s2, pi, sigmak2, XPX, SinvM, M, L, logdetS):
def get_elbo(alpha, mu, s2, logpi, sigmak2, XPX, SinvM, M, L, logdetS):
    nvar, nGcomp = alpha.shape
    Ebeta = np.sum(np.multiply(alpha, mu)[:, 1:], axis = 1)
    mu2 = np.square(mu)
    logdetL = np.sum(np.log(L))
    t0 = - 0.5 * (logdetL + logdetS + np.dot(M.T, SinvM))
    t1 = - 0.5 * np.dot(Ebeta.T, np.dot(XPX, Ebeta))
    t2 = np.dot(Ebeta.T, SinvM)
    t3 = 0
    t4 = 0
    for j in range(nvar):
        Ebeta2j = 0
        for k in range(nGcomp):
            mom2jk = mu2[j, k] + s2[j, k]
            Ebeta2j += alpha[j, k] * mom2jk
            t4jk = 0
            if alpha[j, k] != 0:
                #t4jk += np.log(pi[k])
                t4jk += logpi[k]
                t4jk += - np.log(alpha[j, k])
                if k > 0:
                    t4jk += 0.5 * np.log(s2[j, k])
                    t4jk += - 0.5 * np.log(sigmak2[k])
                    t4jk += - 0.5 * mom2jk / sigmak2[k]
                    t4jk += 0.5
                t4jk *= alpha[j, k]
            t4 += t4jk
        t3 += 0.5 * XPX[j, j] * (Ebeta[j] * Ebeta[j] - Ebeta2j)
    t5 = np.sum(np.multiply(- np.sum(alpha, axis = 0), logpi))
    elbo = t0 + t1 + t2 + t3 + t4 + t5
    return elbo

def ash_qL_regularizer(X, Y, binit, L, mode_optim = True):
    nsample, nvar = X.shape
    if mode_optim:
        M = get_regularized_mode(X, Y, binit, L)
    else:
        M = binit.copy()
    Sinv = precisionLL(X, M, L)
    S = np.linalg.inv(Sinv)
    sgndetS, logdetS = np.linalg.slogdet(S)
    XPX = Sinv.copy()
    XPX[np.diag_indices(nvar)] -= L
    return M, S, Sinv, logdetS, XPX


def ash_qL_pi_init(nGcomp, pi0, sparse = True):
    '''
    By default, pi is initialized with sparsity, i.e. pi[0] > pi[j], for j > 0
    '''
    pi = np.zeros(nGcomp)
    pi[0] = pi0
    pi[1:(nGcomp - 1)] = (1 - pi0) / (nGcomp - 1)
    pi[nGcomp - 1] = 1 - np.sum(pi)
    if not sparse:
        pi = np.repeat(1 / nGcomp, nGcomp)
    return pi


def ash_logistic_qL(X, Y, nGcomp, sigmak2, binit, Linit, 
                    a_dirich = None,
                    maxiter = 1000, tol = 1e-10, pi0 = 0.8, sparse_init = True,
                    reg_step = 1, reg_tol = 0.1,
                    debug = True, debug_step = 10):
    
    '''
    Current focus is on readability and debug.
    '''
    
    nsample, nvar = X.shape
    
    # default dirichlet prior on pi
    if a_dirich is None:
        a_dirich = np.repeat(1., nGcomp)
        a_dirich[0] = 1. # soft max at pi[0]
    
    '''
    Initialize
    '''
    L = Linit.copy()
    M, S, Sinv, logdetS, XPX = ash_qL_regularizer(X, Y, binit, L)
    SinvM = np.dot(Sinv, M)
    SinvRtj = np.dot(Sinv, M) - np.multiply(np.diag(Sinv), M)
    
    pi = ash_qL_pi_init(nGcomp, pi0, sparse = sparse_init)
    logpi = np.log(pi)
    alpha = np.repeat(np.array([pi]), nvar, axis = 0)
    logalpha = np.log(alpha)
    s2 = np.zeros((nvar, nGcomp))
    s2[:, 1:] = 1 / (np.diag(XPX).reshape(-1, 1) + 1 / sigmak2[1:])
    mu = np.multiply(s2, (SinvM - SinvRtj).reshape(-1,1))
    R = np.sum(np.multiply(alpha, mu)[:, 1:], axis = 1)
    
    '''
    Iteration
    '''
    elbo_old = get_elbo(alpha, mu, s2, logpi, sigmak2, XPX, SinvM, M, L, logdetS)
    pi_old = pi.copy()
    pi_reg = pi.copy()
    
    if debug:
        print(f"Starting ELBO: {elbo_old:.4f} Pi0: {pi[0]:.4f}")
        

    for it in range(maxiter):
        
        # E-step
        for j in range(nvar):
            SinvRtj = np.sum(np.multiply(Sinv[:, j], R[:])) - Sinv[j, j] * R[j]
            
            s2[j, 0] = 0.
            mu[j, 0] = 0.
            #logalpha[j, 0] = np.log(pi[0])
            logalpha[j, 0] = logpi[0]
            for k in range(1, nGcomp):
                s2[j, k] = 1 / (XPX[j, j] + (1 / sigmak2[k]))
                mu[j, k] = s2[j, k] * (SinvM[j] - SinvRtj)
                ssf = np.log(s2[j, k] / sigmak2[k])
                ssr = mu[j, k] * mu[j, k] / s2[j, k]
                logalpha[j, k] = logpi[k] + 0.5 * ssf + 0.5 * ssr
                #logalpha[j, k] = np.log(pi[k]) + 0.5 * ssf + 0.5 * ssr
            
            maxL = np.max(logalpha[j, :])
            alpha[j, :]  = np.exp(logalpha[j, :] - maxL)
            alpha[j, :] /= np.sum(alpha[j, :])
            R[j] = np.sum(np.multiply(alpha[j, 1:], mu[j, 1:]))
            
        
        # M-step
        bk = np.sum(alpha, axis = 0) + a_dirich
        pi = bk / np.sum(bk)
        logpi = digamma(bk) - digamma(np.sum(bk))
        #for k in range(nGcomp):
        #    pi[k] = np.sum(alpha[:, k]) / nvar
            
        # Conditional M-step
        if (it + 1) % reg_step == 0:
            eps_reg = max([abs(x - y) for x, y in zip(pi, pi_reg)])
            if eps_reg > reg_tol:
                print(f"Recalculating regularizer, step {it}")            
                L = 1 / np.sum(alpha * sigmak2, axis = 1)
                M, S, Sinv, logdetS, XPX = ash_qL_regularizer(X, Y, R, L)
                SinvM = np.dot(Sinv, M)
                pi_reg = pi.copy()

        elbo = get_elbo(alpha, mu, s2, logpi, sigmak2, XPX, SinvM, M, L, logdetS)
        eps_elbo = elbo - elbo_old
        eps_pi = max([abs(x - y) for x, y in zip(pi, pi_old)])
        
        if debug and (it % debug_step == 0):
            print(f"Iteration {it}. ELBO: {elbo:.4f} Pi0: {pi[0]:.4f} epsilon: {eps_pi:.4f}, {eps_elbo:.4f}")
        
        '''
        Check termination, continue otherwise.
        '''
        if eps_pi < tol:
            print(eps_pi)
            break
 
        '''
        Keep this step in memory.
        Use deep copy to avoid in place substitutions.
        '''
        pi_old    = pi.copy()
        alpha_old = alpha.copy()
        mu_old    = mu.copy()
        s2_old    = s2.copy()
        R_old     = R.copy()
        elbo_old  = elbo
        
    return pi, alpha, mu, s2, R, L
# -

# Hold the result summary statistic in a data frame

# +
# collapse-hide 

res_df = [pd.DataFrame(columns = ['Method', 'RMSE', 'Tjur R2', 'Gini Index']) for n in nsample]

data_Xbeta = [np.dot(Xtest[i], beta[i])                              for i in range(len(nsample))]
data_tjur  = [tjur_R2(Xtest[i], Ytest[i], beta[i])                   for i in range(len(nsample))]
data_rmse  = [np.sqrt(np.mean(np.square(data_Xbeta[i] - 
                                        np.dot(Xtest[i], beta[i])))) for i in range(len(nsample))]
data_gini  = [gini(beta[i])                                          for i in range(len(nsample))]

for i in range(len(nsample)):
    res_df[i] = res_df[i].append(pd.DataFrame([['Data', data_rmse[i], data_tjur[i], data_gini[i]]], 
                                              columns = res_df[i].columns))
# -

# ### Fitting
#
# - Lasso
#
#

# +
# collapse-hide

clf = [linear_model.LogisticRegressionCV(cv = 5, penalty='l1', solver='saga', 
                                        fit_intercept = False, max_iter = 5000) for n in nsample]

for i, n in enumerate(nsample):
    clf[i].fit(X[i], Y[i])

clf_Xbeta = [np.dot(Xtest[i], clf[i].coef_[0])                      for i in range(len(nsample))]
clf_tjur  = [tjur_R2(Xtest[i], Ytest[i], clf[i].coef_[0])           for i in range(len(nsample))]
clf_rmse  = [np.sqrt(np.mean(np.square(clf_Xbeta[i] - 
                                       np.dot(Xtest[i], beta[i])))) for i in range(len(nsample))]
clf_gini  = [gini(clf[i].coef_[0])                                  for i in range(len(nsample))]

for i in range(len(nsample)):
    res_df[i] = res_df[i].append(pd.DataFrame([['LASSO', clf_rmse[i], clf_tjur[i], clf_gini[i]]], 
                                              columns = res_df[i].columns))
# -

# - QL Ash
#
#

# +
# collapse-hide

pi =    [None for n in nsample]
alpha = [None for n in nsample]
mu =    [None for n in nsample]
s2 =    [None for n in nsample]
R =     [None for n in nsample]
L =     [None for n in nsample]

#probk = ash_qL_pi_init(nGcomp, sparsity)
for i in range(len(nsample)):
    gold_reg = np.repeat(1e4, nvar[i])
    for j, b in enumerate(beta[i]):
        if b != 0:
            gold_reg[j] = 1 / prior_variance
    #fix_reg = np.repeat(1 / sigmak2[nGcomp - 1], nvar[i])
    #fix_reg = np.repeat( 1 / np.dot(probk, sigmak2), nvar[i])
    fix_reg = np.repeat(1 / prior_variance, nvar[i])

    binit = np.zeros(nvar[i])
    Linit = fix_reg

    pi[i], alpha[i], mu[i], s2[i], R[i], L[i]  \
                             = ash_logistic_qL(X[i], Y[i], nGcomp, sigmak2, binit, Linit,
                                               tol = 1e-4, debug_step = 1, maxiter = 1000,
                                               reg_step = 10,
                                               pi0 = sparsity, sparse_init = True)

qLash_Xbeta = [np.dot(Xtest[i], R[i])            for i in range(len(nsample))]
qLash_rmse  = [np.sqrt(np.mean(np.square(qLash_Xbeta[i] - np.dot(Xtest[i], beta[i])))) 
                                                 for i in range(len(nsample))]
qLash_r2    = [tjur_R2(Xtest[i], Ytest[i], R[i]) for i in range(len(nsample))]
qLash_gini  = [gini(R[i])                        for i in range(len(nsample))]
for i in range(len(nsample)):
    res_df[i] = res_df[i].append(pd.DataFrame([['qLash', qLash_rmse[i], qLash_r2[i], qLash_gini[i]]], 
                                              columns = res_df[i].columns))
# -

# ### Results

i = 0
print (f'N = {nsample[i]}, P = {nvar[i]}')
res_df[i]

i = 1
print (f'N = {nsample[i]}, P = {nvar[i]}')
res_df[i]

# +
# collapse-hide

fig = plt.figure(figsize = (12, 7))
ax = [fig.add_subplot(121),
      fig.add_subplot(122)]

for i, n in enumerate(nsample):
    whichX = Xtest[i].copy()
    whichY = Ytest[i].copy()

    Xbeta = np.dot(whichX, beta[i])
    pred_data = 1 / (1 + np.exp(-Xbeta))
    ax[i].scatter(Xbeta, whichY, s = 1)

    Xbeta_clf = np.dot(whichX, clf[i].coef_[0])
    pred_clf = 1 / (1 + np.exp(-Xbeta_clf))
    ax[i].scatter(Xbeta, pred_clf, s = 1, label = 'Lasso')

    Xbeta_qL = np.dot(whichX, R[i])
    pred_qL = 1 / (1 + np.exp(-Xbeta_qL))
    ax[i].scatter(Xbeta, pred_qL, s = 1, label = 'qLash')

    leg = ax[i].legend(markerscale = 10)

    ax[i].set_xlabel(r'$\sum_i X_{ni} \beta_i$')
    ax[i].set_ylabel(r'$p_n$')
    ax[i].set_title(f'N = {nsample[i]}, P = {nvar[i]}', pad = 20)
    
plt.tight_layout()
plt.show()
# -

# {{ '<NAME>. (2009). Coefficients of Determination in Logistic Regression Models - A New Proposal: The Coefficient of Discrimination. The American Statistician, 63(4), 366-372. DOI: [10.1198/tast.2009.08210](https://doi.org/10.1198/tast.2009.08210)' | fndetail: 1 }}
#
# {{ 'For a review of sparsity indices, see <NAME> Rickard (2009). Comparing Measures of Sparsity. arXiv. DOI: [arXiv:0811.4706v2](https://arxiv.org/abs/0811.4706v2)' | fndetail: 2 }}



StarcoderdataJupyterScriptsDedupFiltered