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