vbvarsel.elbo — vbvarsel 0.0.1 documentation

# Source code for vbvarsel.elbo

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

    import math
    import numpy as np

    # Evidence Lower Bound (ELBO) COMPUTATION

    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.
        """

        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))

        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?)

        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)

        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)

        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

        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)

        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