Source code for vbvarsel.elbo

from scipy.special import digamma, psi, gammaln, gamma

import math
import numpy as np


# Evidence Lower Bound (ELBO) COMPUTATION
[docs] class ELBO_Computation: """ A class to contain the Evidence Lower Bound (ELBO) computation. This class has no attributes, aside from Python's built-ins, and contains only methods, the majority of which are private. """
[docs] def _ln_pi(self, alphaK, k): """Private method to calculate Pi natural log. Params alphaK: np.ndarray Calculated variational Paramter alphaK k: int The Kth target cluster Returns: float Calculated ln(π) """ return digamma(alphaK[k]) - digamma(sum(alphaK))
[docs] def _ln_precision(self, akj, B): """Private method for calculating precision natural log. Params akj: float Calculated variational parameter Alphakj B: np.ndarray Calculated variational parameter B Returns: np.ndarray Calculated precision natural log """ if B > 1e-30: ld = np.log(B) else: ld = 0.0 return psi(akj) - ld
#(A27?)
[docs] def _log_resp_annealed(self, exp_ln_tau, exp_ln_sigma, f0, N, K, C, T): """Private function to calculate log resp annealed. Params exp_ln_tau: list Expected ln(tau) array exp_ln_sigma: np.ndarray Expected ln(sigma) array f0: np.ndarray Calculated f0 value N: int The Nth observation K: int The Kth cluster C: np.ndarray Calculated variational parameter C, the Covariate selection indicators T: float The annealing temperature Returns log_resp: np.ndarray array of calculated values """ log_resp = np.zeros((N, K)) # ln Z for k in range(K): log_resp[:, k] = ( 0.5 * exp_ln_tau[k] - 0.5 * sum(C) * np.log(2 * math.pi) - 0.5 * exp_ln_sigma[:, k] + f0 ) / T return log_resp
#(A62)
[docs] def _ln_delta_annealed(self, C, j, d, T): """Private function to calculate the annealed value of delta natural log. Params C: np.ndarray Covariate selection indicators j: int Iteration count d: np.ndarray Shape parameter of the Beta distribution on the probability. T: float The annealing temperature Returns ln_delta_ann: np.ndarray calculated array of natural log of annealed delta values """ ln_delta_ann = digamma((C[j] + d + T - 1.0) / T) - digamma( (2 * d + 2 * T - 1.0) / T ) return ln_delta_ann
#(A56)
[docs] def _ln_delta_minus_annealed(self, C, j, d, T): """Private function to calculate the minus of annealed delta natural log. Params C: np.ndarray Covariate selection indicators j: int Iteration count d: np.ndarray Shape parameter of the Beta distribution on the probability. T: float The annealing temperature Returns ln_delta_ann_minus: np.ndarray calculated array of natural log of annealed delta values """ ln_delta_ann_minus = digamma((T - C[j] + d) / T) - digamma( (2 * d + 2 * T - 1.0) / T ) return ln_delta_ann_minus
[docs] def _entropy_wishart(self, k, j, b, a): """Private function to calculate wishart entropy. Params K: int The Kth cluster of the observation j: int Iteration count b: np.ndarray calculated value for variational parameter, betakj a: float calculated value for variational parameter, alphakj Returns e_w: np.ndarray array of entropy wishart values """ if b[k][j, j] > 1e-30: ld = np.log(b[k][j, j]) else: ld = 0.0 e_w = gammaln(a[k][j]) - (a[k][j] - 1) * digamma(a[k][j]) - ld + a[k][j] return e_w
#(A51)
[docs] def compute( self, XDim, K, N, C, Z, d, delta, beta, beta0, alpha, alpha0, a, a0, b, b0, m, m0, exp_ln_tau, exp_ln_sigma, f0, T=1, ) -> float: """Function to compute the Evidence Lower Bound (ELBO). The ELBO is the useful lower-bound on the log-likelihood of observed data. Params XDim: int Number of variables (columns) K: int The Kth cluster of the observation N: int The Nth observation C: np.ndarray Covariate selection indicators Z: np.ndarray cluster assignment matrix d: np.ndarray Shape parameter of the Beta distribution on the probability. delta: int Calculated variational parameter, delta beta: np.ndarray Calculated variational parameter, betakj beta0: float Shrinkage parameter of the Gaussian conditional prior on the cluster means alpha: np.ndarray Calculated variational parameter, alphak alpha0: float Prior coefficient count a: float Calculated variational parameter, alphakj a0: np.ndarray Degrees of freedom for the Gamma prior on the cluster precision b: np.ndarray Calculated variational parameter, B b0: np.ndarray Prior covariance m: np.ndarray Calculated variational parameter, m m0: np.ndarray Prior mean exp_ln_tau: list expected ln(tau) array exp_ln_sigma: np.ndarray expected ln(sigma) array f0: np.ndarray Calculated variational parameter, f0 T: float The annealing temperature Returns: elbo: float Calculated ELBO value """ # nifty way to find out what the attributes are of the class I am using when # I am unfamiliar with how the mathematics of these functions works. # import inspect # import pprint # sig, elbo_locals = inspect.signature(self.compute), locals() # pprint.pprint([f"{str(_key)} : {type(elbo_locals[param.name])}" for _key, param in sig.parameters.items()]) # # [<class 'int'>, <class 'int'>, <class 'int'>, <class 'numpy.ndarray'>, # # <class 'numpy.ndarray'>, <class 'int'>, <class 'numpy.ndarray'>, # # <class 'numpy.ndarray'>, <class 'float'>, <class 'numpy.ndarray'>, # # <class 'float'>, <class 'numpy.ndarray'>, <class 'float'>, # # <class 'numpy.ndarray'>, <class 'numpy.ndarray'>, <class 'numpy.ndarray'>, # # <class 'numpy.ndarray'>, <class 'numpy.ndarray'>, <class 'list'>, # # <class 'numpy.ndarray'>, <class 'numpy.ndarray'>, <class 'float'>] # E[ln p(X|Z, μ, Λ)] def _first_term(N, K, Z, C, exp_ln_tau, exp_ln_sigma, f0, T): """Private internal function to calculate the 1st term of the ELBO Params N: int The Nth observation K: int The Kth cluster of the observation Z: np.ndarray cluster assignment matrix C: np.ndarray[int] Covariate selection indicators exp_ln_tau: list expected ln(tau) array exp_ln_sigma: np.ndarray expected ln(mu) array f0: np.ndarray calculated f0 value T: float The Annealing temperature Returns F2: float first ELBO algorithm term """ ln_resp = self._log_resp_annealed(exp_ln_tau, exp_ln_sigma, f0, N, K, C, T) F2 = 0 for n in range(N): for k in range(K): F2 += Z[n, k] * (ln_resp[n, k]) return F2 # E[ln p(Z|π)] def _second_term(N, K, Z, alpha): """Private internal function to calculate ELBO 2nd term Params N: int The Nth observation K: int The Kth cluster of the observation Z: np.ndarray cluster assignment matrix alpha: np.ndarray calculated alphak value Returns s: float calculated second ELBO algorithm term """ s = 0 for n in range(N): for k in range(K): s += Z[n, k] * self._ln_pi(alpha, k) return s # E[ln p(π)] def _third_term(alpha0, K, alpha): """Private internal function to calculate the 3rd term of the ELBO Params alpha0: float Degrees of freedom for the Gamma prior on the cluster precision K: int The Kth cluster of the observation alpha: float Prior coefficient count Return: a + b: float calculated third term for ELBO algorithm """ a = gammaln(alpha0 * K) - K * gammaln(alpha0) b = (alpha0 - 1) * sum([self._ln_pi(alpha, k) for k in range(K)]) return a + b # E[ln p(μ, Λ)] def _fourth_term(K, XDim, beta0, beta, a0, a, b0, b, m, m0): """Private internal function to calculate the 4th term of the ELBO Params K: int The Kth cluster of the observation XDim: int Number of variables (columns) beta0: float Shrinkage parameter of the Gaussian conditional prior on the cluster beta: np.ndarray calculated betakj values a0: int Degrees of freedom for the Gamma prior on the cluster precision a: np.ndarray calculated akj values b0: np.ndarray prior covariance b: np.ndarray calculated bkj value m: np.ndarray calculated m value m0: np.ndarray prior mean Returns t: float calculated fourth term of the ELBO algorithm """ t = 0 for k in range(K): for j in range(XDim): F0 = 0.5 * np.log(beta0 / (2 * math.pi)) +0.5 * self._ln_precision(a[k, j], b[k][j, j]) if beta[k][j] > 0: F1 = ( (beta0 * a[k][j] / beta[k][j]) * ((m[k][j] - m0[j]) ** 2) + beta0 / beta[k][j] + b0[j, j] * a[k][j] / beta[k][j] ) else: F1 = ( (beta0 * a[k][j]) * ((m[k][j] - m0[j]) ** 2) + beta0 + b0[j, j] * a[k][j] ) F2 = ( -np.log(gamma(a0)) + a0 * np.log(b0[j, j]) + (a0 - 2) * self._ln_precision(a[k, j], b[k][j, j]) ) t += F0 - F1 + F2 return t # E[ln p(𝛾,𝛿)] def _fifth_term(XDim, d, C, T): """Private internal function to calculate the 5th term of the ELBO Params XDim: int Number of variables (columns) d: np.ndarray Shape parameter of the Beta distribution on the probability. C: np.ndarray[int] Covariate selection indicators matrix T: float The Annealing temperature Returns a: ndarray the calculated 6th term of the ELBO algorithm """ a = 0 for j in range(XDim): F1 = (d + C[j] - 1) * self._ln_delta_annealed(C, j, d, T) F2 = (d + C[j]) * self._ln_delta_minus_annealed(C, j, d, T) F3 = np.log(gamma(2 * d)) - 2 * np.log(gamma(d)) a += F1 + F2 + F3 return a # E[ln q(Z)] def _sixth_term(Z:np.ndarray, N:int, K:int): """Private internal function to calculate the 6th term of the ELBO Params Z: np.ndarray cluster assignment matrix N: int The Nth observation K: int The Kth cluster of the observation Returns a: ndarray the calculated 6th term of the ELBO algorithm """ a = 0 for n in range(N): for k in range(K): if Z[n, k] > 1e-30: ld = np.log(Z[n, k]) else: ld = 0.0 a += Z[n, k] * ld return a # E[ln q(π)] def _seventh_term(alpha:np.ndarray): """Private internal function to calculate the 7th term of the ELBO Params alpha: np.ndarray Calculated alpha value Returns a + b """ a = sum([(alpha[k] - 1) * self._ln_pi(alpha, k) for k in range(K)]) b = gammaln(sum(alpha)) - sum([gammaln(alpha[k]) for k in range(K)]) return a + b # E[ln q(μ, Λ)] def _eighth_term(K, XDim, beta, a, b): """Private internal function to calculate the 8th term of the ELBO Params K: int The Kth cluster of the observation XDim: int Number of variables (columns) beta: ndarray[float] claculated betakj values a: ndarray[float] calculated alphakj values b: ndarray[float] calculated B values Returns t: float calculated 8th term for the ELBO algorithm """ t = 0 for k in range(K): for j in range(XDim): t += ( 0.5 * self._ln_precision(a[k, j], b[k][j, j]) + 0.5 * np.log(beta[k][j] / (2 * np.pi)) - 0.5 - self._entropy_wishart(k, j, b, a) ) return t # E[ln q(𝛾,𝛿)] def _ninth_term(XDim, d, delta, C, T): """Private internal function to calculate the 9th term of the ELBO Params XDim: int Number of variables (columns) d: np.ndarray Shape parameter of the Beta distribution on the probability. delta: int calculated delta value C: np.ndarray[int] Covariate selection indicators matrix T: float The annealing temperature Returns F0 + F1: float calculated 9th term of the ELBO algorithm """ F0 = ( 2 * XDim * ( (d - 1) * (digamma(d) - digamma(2 * d)) + np.log(gamma(2 * d)) - 2 * np.log(gamma(d)) ) ) F1 = 0 for j in range(XDim): F1 += delta[j] * self._ln_delta_annealed(C, j, d, T) + ( 1 - delta[j] ) * self._ln_delta_minus_annealed(C, j, d, T) return F0 + F1 a_t = _first_term(N, K, Z, C, exp_ln_tau, exp_ln_sigma, f0, T) b_t = _second_term(N, K, Z, alpha) c_t = _third_term(alpha0, K, alpha) d_t = _fourth_term(K, XDim, beta0, beta, a0, a, b0, b, m, m0) e_t = _fifth_term(XDim, d, C, T) f_t = _sixth_term(Z, N, K) g_t = _seventh_term(alpha) h_t = _eighth_term(K, XDim, beta, a, b) i_t = _ninth_term(XDim, d, delta, C, T) return a_t + b_t + c_t + d_t + e_t - (f_t + g_t + h_t + i_t) * T