Exemple #1
0
def _pdf_gamma(a, b, axes=None, scale=4, color="k"):
    """
    """
    if axes is None:
        axes = plt.gca()

    if np.size(a) != 1 or np.size(b) != 1:
        raise ValueError("Parameters must be scalars")
    mean = a / b
    v = scale * np.sqrt(a / b ** 2)
    m = max(0, mean - v)
    n = mean + v
    x = np.linspace(m, n, num=100)
    logx = np.log(x)
    lpdf = random.gamma_logpdf(b * x, logx, a * logx, a * np.log(b), special.gammaln(a))
    p = np.exp(lpdf)
    return axes.plot(x, p, color=color)
Exemple #2
0
def _pdf_gamma(a, b, axes=None, scale=4, color='k'):
    """
    """
    if axes is None:
        axes = plt.gca()

    if np.size(a) != 1 or np.size(b) != 1:
        raise ValueError("Parameters must be scalars")
    mean = a / b
    v = scale * np.sqrt(a / b**2)
    m = max(0, mean - v)
    n = mean + v
    x = np.linspace(m, n, num=100)
    logx = np.log(x)
    lpdf = random.gamma_logpdf(b * x, logx, a * logx, a * np.log(b),
                               special.gammaln(a))
    p = np.exp(lpdf)
    return axes.plot(x, p, color=color)
Exemple #3
0
    def _compute_bound(self, R, logdet=None, inv=None, Q=None, gradient=False, terms=False):
        """
        Rotate q(X) and q(alpha).

        Assume:
        p(X|alpha) = prod_m N(x_m|0,diag(alpha))
        p(alpha) = prod_d G(a_d,b_d)
        """

        ## R = self._full_rotation_matrix(R)
        ## if inv is not None:
        ##     inv = self._full_rotation_matrix(inv)

        #
        # Transform the distributions and moments
        #

        plates_alpha = self.plates_alpha
        plates_X = self.plates_X

        # Compute rotated second moment
        if self.plate_axis is not None:
            # The plate axis has been moved to be the last plate axis

            if Q is None:
                raise ValueError("Plates should be rotated but no Q give")

            # Transform covariance
            sumQ = np.sum(Q, axis=0)
            QCovQ = sumQ[:, None, None] ** 2 * self.CovX

            # Rotate plates
            if self.precompute:
                QX_QX = np.einsum("...kalb,...ik,...il->...iab", self.X_X, Q, Q)
                XX = QX_QX + QCovQ
                XX = sum_to_plates(XX, plates_alpha[:-1], ndim=2)
                Xmu = np.einsum("...kaib,...ik->...iab", self.X_mu, Q)
                Xmu = sum_to_plates(Xmu, plates_alpha[:-1], ndim=2)
            else:
                X = self.X
                mu = self.mu
                QX = np.einsum("...ik,...kj->...ij", Q, X)
                XX = sum_to_plates(QCovQ, plates_alpha[:-1], ndim=2) + sum_to_plates(
                    linalg.outer(QX, QX), plates_alpha[:-1], ndim=2, plates_from=plates_X
                )
                Xmu = sum_to_plates(linalg.outer(QX, self.mu), plates_alpha[:-1], ndim=2, plates_from=plates_X)

            mu2 = self.mu2
            D = np.shape(XX)[-1]
            logdet_Q = D * np.log(np.abs(sumQ))

        else:
            XX = self.XX
            mu2 = self.mu2
            Xmu = self.Xmu
            logdet_Q = 0

        # Compute transformed moments
        # mu2 = np.einsum('...ii->...i', mu2)
        RXmu = np.einsum("...ik,...ki->...i", R, Xmu)
        RXX = np.einsum("...ik,...kj->...ij", R, XX)
        RXXR = np.einsum("...ik,...ik->...i", RXX, R)

        # <(X-mu) * (X-mu)'>_R
        XmuXmu = RXXR - 2 * RXmu + mu2

        D = np.shape(R)[0]

        # Compute q(alpha)
        if self.update_alpha:
            # Parameters
            a0 = self.a0
            b0 = self.b0
            a = self.a
            b = b0 + 0.5 * sum_to_plates(XmuXmu, plates_alpha, plates_from=None, ndim=0)
            # Some expectations
            alpha = a / b
            logb = np.log(b)
            logalpha = -logb  # + const
            b0_alpha = b0 * alpha
            a0_logalpha = a0 * logalpha
        else:
            alpha = self.alpha
            logalpha = 0

        #
        # Compute the cost
        #

        def sum_plates(V, *plates):
            full_plates = misc.broadcasted_shape(*plates)

            r = self.node_X.broadcasting_multiplier(full_plates, np.shape(V))
            return r * np.sum(V)

        XmuXmu_alpha = XmuXmu * alpha

        if logdet is None:
            logdet_R = np.linalg.slogdet(R)[1]
            inv_R = np.linalg.inv(R)
        else:
            logdet_R = logdet
            inv_R = inv

        # Compute entropy H(X)
        logH_X = random.gaussian_entropy(-2 * sum_plates(logdet_R + logdet_Q, plates_X), 0)

        # Compute <log p(X|alpha)>
        logp_X = random.gaussian_logpdf(
            sum_plates(XmuXmu_alpha, plates_alpha[:-1] + [D]), 0, 0, sum_plates(logalpha, plates_X + [D]), 0
        )

        if self.update_alpha:

            # Compute entropy H(alpha)
            # This cancels out with the log(alpha) term in log(p(alpha))
            logH_alpha = 0

            # Compute <log p(alpha)>
            logp_alpha = random.gamma_logpdf(
                sum_plates(b0_alpha, plates_alpha), 0, sum_plates(a0_logalpha, plates_alpha), 0, 0
            )
        else:
            logH_alpha = 0
            logp_alpha = 0

        # Compute the bound
        if terms:
            bound = {self.node_X: logp_X + logH_X}
            if self.update_alpha:
                bound.update({self.node_alpha: logp_alpha + logH_alpha})
        else:
            bound = 0 + logp_X + logp_alpha + logH_X + logH_alpha

        if not gradient:
            return bound

        #
        # Compute the gradient with respect R
        #

        broadcasting_multiplier = self.node_X.broadcasting_multiplier

        def sum_plates(V, plates):
            ones = np.ones(np.shape(R))
            r = broadcasting_multiplier(plates, np.shape(V)[:-2])
            return r * misc.sum_multiply(V, ones, axis=(-1, -2), sumaxis=False, keepdims=False)

        D_XmuXmu = 2 * RXX - 2 * gaussian.transpose_covariance(Xmu)

        DXmuXmu_alpha = np.einsum("...i,...ij->...ij", alpha, D_XmuXmu)
        if self.update_alpha:
            D_b = 0.5 * D_XmuXmu
            XmuXmu_Dalpha = np.einsum(
                "...i,...i,...i,...ij->...ij",
                sum_to_plates(XmuXmu, plates_alpha, plates_from=None, ndim=0),
                alpha,
                -1 / b,
                D_b,
            )
            D_b0_alpha = np.einsum("...i,...i,...i,...ij->...ij", b0, alpha, -1 / b, D_b)
            D_logb = np.einsum("...i,...ij->...ij", 1 / b, D_b)
            D_logalpha = -D_logb
            D_a0_logalpha = a0 * D_logalpha
        else:
            XmuXmu_Dalpha = 0
            D_logalpha = 0

        D_XmuXmu_alpha = DXmuXmu_alpha + XmuXmu_Dalpha
        D_logR = inv_R.T

        # Compute dH(X)
        dlogH_X = random.gaussian_entropy(-2 * sum_plates(D_logR, plates_X), 0)

        # Compute d<log p(X|alpha)>
        dlogp_X = random.gaussian_logpdf(
            sum_plates(D_XmuXmu_alpha, plates_alpha[:-1]),
            0,
            0,
            (sum_plates(D_logalpha, plates_X) * broadcasting_multiplier((D,), plates_alpha[-1:])),
            0,
        )

        if self.update_alpha:

            # Compute dH(alpha)
            # This cancels out with the log(alpha) term in log(p(alpha))
            dlogH_alpha = 0

            # Compute d<log p(alpha)>
            dlogp_alpha = random.gamma_logpdf(
                sum_plates(D_b0_alpha, plates_alpha[:-1]), 0, sum_plates(D_a0_logalpha, plates_alpha[:-1]), 0, 0
            )
        else:
            dlogH_alpha = 0
            dlogp_alpha = 0

        if terms:
            raise NotImplementedError()
            dR_bound = {self.node_X: dlogp_X + dlogH_X}
            if self.update_alpha:
                dR_bound.update({self.node_alpha: dlogp_alpha + dlogH_alpha})
        else:
            dR_bound = 0 * dlogp_X + dlogp_X + dlogp_alpha + dlogH_X + dlogH_alpha

        if self.subset:
            indices = np.ix_(self.subset, self.subset)
            dR_bound = dR_bound[indices]

        if self.plate_axis is None:
            return (bound, dR_bound)

        #
        # Compute the gradient with respect to Q (if Q given)
        #

        # Some pre-computations
        Q_RCovR = np.einsum("...ik,...kl,...il,...->...i", R, self.CovX, R, sumQ)
        if self.precompute:
            Xr_rX = np.einsum("...abcd,...jb,...jd->...jac", self.X_X, R, R)
            QXr_rX = np.einsum("...akj,...ik->...aij", Xr_rX, Q)
            RX_mu = np.einsum("...jk,...akbj->...jab", R, self.X_mu)

        else:
            RX = np.einsum("...ik,...k->...i", R, X)
            QXR = np.einsum("...ik,...kj->...ij", Q, RX)
            QXr_rX = np.einsum("...ik,...jk->...kij", QXR, RX)
            RX_mu = np.einsum("...ik,...jk->...kij", RX, mu)

            QXr_rX = sum_to_plates(QXr_rX, plates_alpha[:-2], ndim=3, plates_from=plates_X[:-1])
            RX_mu = sum_to_plates(RX_mu, plates_alpha[:-2], ndim=3, plates_from=plates_X[:-1])

        def psi(v):
            """
            Compute: d/dQ 1/2*trace(diag(v)*<(X-mu)*(X-mu)>)

            = Q*<X>'*R'*diag(v)*R*<X> + ones * Q diag( tr(R'*diag(v)*R*Cov) ) 
              + mu*diag(v)*R*<X>
            """

            # Precompute all terms to plates_alpha because v has shape
            # plates_alpha.

            # Gradient of 0.5*v*<x>*<x>
            v_QXrrX = np.einsum("...kij,...ik->...ij", QXr_rX, v)

            # Gradient of 0.5*v*Cov
            Q_tr_R_v_R_Cov = np.einsum("...k,...k->...", Q_RCovR, v)[..., None, :]

            # Gradient of mu*v*x
            mu_v_R_X = np.einsum("...ik,...kji->...ij", v, RX_mu)

            return v_QXrrX + Q_tr_R_v_R_Cov - mu_v_R_X

        def sum_plates(V, plates):
            ones = np.ones(np.shape(Q))
            r = self.node_X.broadcasting_multiplier(plates, np.shape(V)[:-2])

            return r * misc.sum_multiply(V, ones, axis=(-1, -2), sumaxis=False, keepdims=False)

        if self.update_alpha:
            D_logb = psi(1 / b)
            XX_Dalpha = -psi(alpha / b * sum_to_plates(XmuXmu, plates_alpha))
            D_logalpha = -D_logb
        else:
            XX_Dalpha = 0
            D_logalpha = 0
        DXX_alpha = 2 * psi(alpha)
        D_XX_alpha = DXX_alpha + XX_Dalpha
        D_logdetQ = D / sumQ
        N = np.shape(Q)[-1]

        # Compute dH(X)
        dQ_logHX = random.gaussian_entropy(-2 * sum_plates(D_logdetQ, plates_X[:-1]), 0)

        # Compute d<log p(X|alpha)>
        dQ_logpX = random.gaussian_logpdf(
            sum_plates(D_XX_alpha, plates_alpha[:-2]),
            0,
            0,
            (sum_plates(D_logalpha, plates_X[:-1]) * broadcasting_multiplier((N, D), plates_alpha[-2:])),
            0,
        )

        if self.update_alpha:

            D_alpha = -psi(alpha / b)
            D_b0_alpha = b0 * D_alpha
            D_a0_logalpha = a0 * D_logalpha

            # Compute dH(alpha)
            # This cancels out with the log(alpha) term in log(p(alpha))
            dQ_logHalpha = 0

            # Compute d<log p(alpha)>
            dQ_logpalpha = random.gamma_logpdf(
                sum_plates(D_b0_alpha, plates_alpha[:-2]), 0, sum_plates(D_a0_logalpha, plates_alpha[:-2]), 0, 0
            )
        else:

            dQ_logHalpha = 0
            dQ_logpalpha = 0

        if terms:
            raise NotImplementedError()
            dQ_bound = {self.node_X: dQ_logpX + dQ_logHX}
            if self.update_alpha:
                dQ_bound.update({self.node_alpha: dQ_logpalpha + dQ_logHalpha})
        else:
            dQ_bound = 0 * dQ_logpX + dQ_logpX + dQ_logpalpha + dQ_logHX + dQ_logHalpha
        return (bound, dR_bound, dQ_bound)