Project: s41467-024-53818-w | Function: get_elbo (python)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** Computes the Evidence Lower BOund (ELBO) for the variational Bayesian inference in SharePro-GxE, used to monitor convergence and optimize the model

**KEY TERMS:** {
  "ll1": "likelihood term 1",
  "ll2": "likelihood term 2",
  "ll3": "likelihood term 3",
  "mklbeta": "KL divergence for beta",
  "mkldelta": "KL divergence for delta",
  "WtRW": "quadratic form for correlation"
}

**CONTEXT CODE:** import argparse
import logging
import pandas as pd
import numpy as np
from scipy.special import softmax, expit
from scipy.stats import chi2

np.set_printoptions(precision=4, linewidth=200)

def title():
    print('**********************************************************************')
    print('* SharePro for joint fine-mapping and GxE analysis                   *')
    print('* Version 2.0.0                                                      *')
    print('* (C) Wenmin Zhang (wenmin.zhang@mail.mcgill.ca)                     *')
    print('**********************************************************************')


class SparseReg(object):
    def __init__(self, P, K, N, sigma, varb):
        """initialize and set parameters"""
        self.p = P
        self.k = K
        self.n = N
        self.sigma = sigma
        self.y_tau = 1.0 / (1 - varb)
        self.priorb = 1.0 / varb
        self.postn = self.n * self.y_tau + self.priorb
        self.lamb = self.priorb / self.postn
        self.beta_mu = np.zeros([self.p, self.k])
        self.delta = np.zeros((self.p, self.k))
        self.u1 = np.log(self.sigma / (1 - self.sigma)) + 0.5 * np.log(self.lamb)

    def infer_q_beta(self, b_hat_s, R_s, GAM, k):
        """perform variational updates for the k-th effect with beta and c"""
        raise NotImplementedError("Function implementation has been masked")

    def get_elbo(self, b_hat_s, R_s, GAM):
        """get elbo with beta and c"""
        raise NotImplementedError("Function implementation has been masked")


class SharePro(object):
    def __init__(self, P, K, num, sigma, Nlst, varblst):
        """initialize and set parameters"""
        self.p = P
        self.k = K
        self.num = num
        self.sigma = sigma
        self.prior_pi = np.ones((self.p,)) * (1 / self.p)
        self.gamma = np.zeros((self.p, self.k))
        self.usum1 = self.num * np.log(1 - self.sigma**(1/self.num)) + np.log(self.prior_pi)
        self.SR = [SparseReg(P, K, Nlst[i], self.sigma**(1/self.num), varblst[i]) for i in range(self.num)]

    def infer_q_s(self, bhatlst, ld):
        """perform variational updates for s"""
        for k in range(self.k):
            idxall = list(range(self.k))
            idxall.remove(k)
            unum = np.array([j.infer_q_beta(bhatlst[i], ld, self.gamma, k) for (i,j) in enumerate(self.SR)])
            unumexp = np.log(1 + np.exp(-unum))
            uall = unum.sum(axis=0) + unumexp.sum(axis=0) + self.usum1
            self.gamma[:, k] = softmax(uall)

    def get_PIP(self):
        PIP_overall = self.gamma.max(axis=1).round(4) # causal in either
        return PIP_overall

    def get_effect(self, ld, cthres, pthres):
        vidx = np.argsort(-self.gamma, axis=1)
        matidx = np.argsort(-self.gamma, axis=0)
        mat_eff = np.zeros((self.p, self.k))  # effective gamma
        for p in range(self.p):
            mat_eff[p, vidx[p, 0]] = self.gamma[p, vidx[p, 0]]
        csum = mat_eff.sum(axis=0).round(2)
        logging.info(f"Attainable coverage for effect groups: {csum}")
        eff = {}
        eff_gamma = {}
        eff_purity = {}
        for k in range(self.k):
            if csum[k] > cthres:
                p = 0
                while np.sum(mat_eff[matidx[0:p, k], k]) < cthres * csum[k]:
                    p = p + 1
                cidx = matidx[0:p, k].tolist()
                purity = abs(ld[np.ix_(cidx, cidx)]).min()
                if purity > pthres:
                    eff[k] = cidx
                    eff_gamma[k] = mat_eff[eff[k], k].round(4).tolist()
                    eff_purity[k] = purity
        return eff, eff_gamma, eff_purity

    def update_pi(self, new_pi):
        self.prior_pi = new_pi

    def get_elbo(self, bhatlst, ld):
        """get elbo with s"""
        llsum, mklbetasum, mkldeltasum = zip(*[j.get_elbo(bhatlst[i], ld, self.gamma) for (i,j) in enumerate(self.SR)])
        gammaterm1 = (self.gamma * np.tile(self.prior_pi.reshape(-1, 1), (1, self.k))).sum()
        gammaterm2 = (self.gamma[self.gamma != 0] * np.log(self.gamma[self.gamma != 0])).sum()
        mklgamma = gammaterm1 - gammaterm2
        ll = sum(llsum)
        mklbeta = sum(mklbetasum)
        mkldelta = sum(mkldeltasum)
        elbo = ll + mklbeta + mkldelta + mklgamma
        return ll, mklbeta, mkldelta, mklgamma, elbo

    def train(self, bhatlst, ld, maxite, eps, ubound):
        loss = 100000.0
        for ite in range(maxite):
            self.infer_q_s(bhatlst, ld)
            ll, mklbeta, mkldelta, mklgamma, elbo = self.get_elbo(bhatlst, ld)
            logging.info('*' * 70)
            logging.info('Iteration-->{} . Likelihood: {:.1f} . KL_b: {:.1f} . KL_c: {:.1f} . KL_s: {:.1f} . ELBO: {:.1f}'.format(ite, ll, mklbeta, mkldelta, mklgamma, elbo))
            if abs(elbo - loss) < eps:
                converged = True
                break
            if ite == (maxite - 1) or elbo > ubound:
                print("Detected mismatch between summary statistics and LD matrix.")
                converged = False
                break
            loss = elbo
        return converged

def get_eff_maxld(eff, ld):
    idx = [i[0] for i in eff.values()]
    if len(eff)>1:
        maxld = np.abs(np.tril(ld[np.ix_(idx,idx)],-1)).max()
    else:
        maxld = 0
    return maxld

def adaptive_train(bhatlst, ld, Nlst, sigma, K, maxite, eps, ubound, cthres, pthres):
    P = bhatlst[0].shape[0]
    num = len(Nlst)
    varb = [(i**2).max() for i in bhatlst]
    Nmaxlst = [i.max() for i in Nlst]


Complete the function implementation replacing "raise NotImplementedError".