# Source code for bayespy.inference.vmp.nodes.gamma

################################################################################
# Copyright (C) 2011-2012,2014 Jaakko Luttinen
#
################################################################################

"""
Module for the gamma distribution node.
"""

import numpy as np
import scipy.special as special

from .node import Node, Moments, ensureparents
from .deterministic import Deterministic
from .stochastic import Stochastic
from .expfamily import ExponentialFamily, ExponentialFamilyDistribution
from .constant import Constant

from bayespy.utils import misc
from bayespy.utils import random

def diagonal(alpha):
"""
Create a diagonal Wishart node from a Gamma node.
"""
return _GammaToDiagonalWishart(alpha,
name=alpha.name + " as Wishart")

class GammaPriorMoments(Moments):
"""
Class for the moments of the shape parameter in gamma distributions.
"""

dims = ( (), () )

def compute_fixed_moments(self, a):
"""
Compute the moments for a fixed value
"""
a = np.asanyarray(a)
if np.any(a <= 0):
raise ValueError("Shape parameter must be positive")
u0 = a
u1 = special.gammaln(a)
return [u0, u1]

@classmethod
def from_values(cls, a):
"""
Return the shape of the moments for a fixed value.
"""
return cls()

[docs]class GammaMoments(Moments):
"""
Class for the moments of gamma variables.
"""

dims = ( (), () )

[docs]    def compute_fixed_moments(self, x):
"""
Compute the moments for a fixed value
"""
x = np.asanyarray(x)
if np.any(x < 0):
raise ValueError("Values must be positive")
u0 = x
u1 = np.log(x)
return [u0, u1]

[docs]    @classmethod
def from_values(cls, x):
"""
Return the shape of the moments for a fixed value.
"""
return cls()

"""
Class for the VMP formulas of gamma variables.
"""

[docs]    def compute_message_to_parent(self, parent, index, u_self, u_a, u_b):
r"""
Compute the message to a parent node.
"""
x = u_self[0]
logx = u_self[1]

if index == 0:
b = u_b[0]
logb = u_b[1]
return [logx + logb,
-1]
elif index == 1:
a = u_a[0]
return [-x,
a]
else:
raise ValueError("Index out of bounds")

r"""
Compute the natural parameter vector given parent moments.
"""
return [-u_parents[1][0],
1*u_parents[0][0]]

r"""
Compute the moments and :math:g(\phi).

.. math::

\overline{\mathbf{u}}  (\boldsymbol{\phi})
&=
\begin{bmatrix}
- \frac{\phi_2} {\phi_1}
\\
\psi(\phi_2) - \log(-\phi_1)
\end{bmatrix}
\\
g_{\boldsymbol{\phi}} (\boldsymbol{\phi})
&=
TODO
"""
with np.errstate(invalid='raise', divide='raise'):
log_b = np.log(-phi[0])
u0 = phi[1] / (-phi[0])
u1 = special.digamma(phi[1]) - log_b
u = [u0, u1]
g = phi[1] * log_b - special.gammaln(phi[1])
return (u, g)

[docs]    def compute_cgf_from_parents(self, *u_parents):
r"""
Compute :math:\mathrm{E}_{q(p)}[g(p)]
"""
a = u_parents[0][0]
gammaln_a = u_parents[0][1] #special.gammaln(a)
b = u_parents[1][0]
log_b = u_parents[1][1]
g = a * log_b - gammaln_a
return g

r"""
Compute the moments and :math:f(x) for a fixed value.
"""
x = np.asanyarray(x)
if np.any(x < 0):
raise ValueError("Values must be positive")
logx = np.log(x)
u = [x, logx]
f = -logx
return (u, f)

[docs]    def random(self, *phi, plates=None):
r"""
Draw a random sample from the distribution.
"""
return random.gamma(phi[1], -1/phi[0], size=plates)

[docs]    def compute_gradient(self, g, u, phi):
r"""
Compute the moments and :math:g(\phi).

.. math::

\mathrm{d}\overline{\mathbf{u}} &=
\begin{bmatrix}
- \frac{\mathrm{d}\phi_2} {phi_1} + \frac{\phi_2}{\phi_1^2} \mathrm{d}\phi_1
\\
\psi^{(1)}(\phi_2) \mathrm{d}\phi_2 - \frac{1}{\phi_1} \mathrm{d}\phi_1
\end{bmatrix}

is, given the Riemannian gradient :math:\tilde{\nabla}:

.. math::

\nabla =
\begin{bmatrix}
\nabla_1 \frac{\phi_2}{\phi_1^2} - \nabla_2 \frac{1}{\phi_1}
\\
\nabla_2 \psi^{(1)}(\phi_2) - \nabla_1 \frac {1} {\phi_1}
\end{bmatrix}
"""
d0 = g[0] * phi[1] / phi[0]**2 - g[1] / phi[0]
d1 = g[1] * special.polygamma(1, phi[1]) - g[0] / phi[0]
return [d0, d1]

[docs]class Gamma(ExponentialFamily):
"""
Node for gamma random variables.

Parameters
----------

a : scalar or array

Shape parameter

b : gamma-like node or scalar or array

Rate parameter
"""

dims = ( (), () )
_moments = GammaMoments()
_parent_moments = (GammaPriorMoments(),
GammaMoments())

[docs]    def __init__(self, a, b, **kwargs):
"""
Create gamma random variable node
"""
super().__init__(a, b, **kwargs)

def __str__(self):
"""
Print the distribution using standard parameterization.
"""
a = self.phi[1]
b = -self.phi[0]
return ("%s ~ Gamma(a, b)\n"
"  a =\n"
"%s\n"
"  b =\n"
"%s\n"
% (self.name, a, b))

[docs]    def as_wishart(self, ndim=0):
if ndim != 0:
raise NotImplementedError()
return _GammaToScalarWishart(self, name=self.name + " as Wishart")

[docs]    def as_diagonal_wishart(self):
return _GammaToDiagonalWishart(self,
name=self.name + " as Wishart")

[docs]    def diag(self):
return self.as_diagonal_wishart()

class GammaShape(Stochastic):
"""
ML point estimator for the shape parameter of the gamma distribution
"""

dims = ( (), () )
_moments = GammaPriorMoments()
_parent_moments = ()

def __init__(self, m0=0, m1=0, **kwargs):
"""
Create gamma random variable node
"""
super().__init__(dims=self.dims, initialize=False, **kwargs)
self.u = self._moments.compute_fixed_moments(1)
self._m0 = m0
self._m1 = m1
return

def _update_distribution_and_lowerbound(self, m):
r"""
Find maximum likelihood estimate for the shape parameter

Messages from children appear in the lower bound as

.. math::

m_0 \cdot x +  m_1 \cdot \log(\Gamma(x))

Take derivative, put it zero and solve:

.. math::

m_0 + m_1 \cdot d\log(\Gamma(x)) &= 0
\\
m_0 + m_1 \cdot \psi(x) &= 0
\\
x &= \psi^{-1}(-\frac{m_0}{m_1})

where :math:\psi^{-1} is the inverse digamma function.
"""

# Maximum likelihood estimate
m0 = self._m0 + m[0]
m1 = self._m1 + m[1]
x = misc.invpsi(-m0 / m1)

# Compute moments
self.u = self._moments.compute_fixed_moments(x)

return

def initialize_from_value(self, x):
self.u = self._moments.compute_fixed_moments(x)
return

def lower_bound_contribution(self):
return 0

class _GammaToDiagonalWishart(Deterministic):
"""
Transform a set of gamma scalars into a diagonal Wishart matrix.

The last plate is used as the diagonal dimension.
"""

_parent_moments = [GammaMoments()]

@ensureparents
def __init__(self, alpha, **kwargs):

# Check for constant
if misc.is_numeric(alpha):
alpha = Constant(Gamma)(alpha)

if len(alpha.plates) == 0:
raise Exception("Gamma variable needs to have plates in "
"order to be used as a diagonal Wishart.")
D = alpha.plates[-1]

# FIXME: Put import here to avoid circular dependency import
from .wishart import WishartMoments
self._moments = WishartMoments((D,))
dims = ( (D,D), () )

# Construct the node
super().__init__(alpha,
dims=self._moments.dims,
**kwargs)

def _plates_to_parent(self, index):
D = self.dims[0][0]
return self.plates + (D,)

def _plates_from_parent(self, index):
return self.parents[index].plates[:-1]

@staticmethod
def _compute_weights_to_parent(index, weights):
return weights[..., np.newaxis]

def get_moments(self):
u = self.parents[0].get_moments()

# Form a diagonal matrix from the gamma variables
return [np.identity(self.dims[0][0]) * u[0][...,np.newaxis],
np.sum(u[1], axis=(-1))]

@staticmethod
def _compute_message_to_parent(index, m_children, *u_parents):

# Take the diagonal
m0 = np.einsum('...ii->...i', m_children[0])
m1 = np.reshape(m_children[1], np.shape(m_children[1]) + (1,))

return [m0, m1]

class _GammaToScalarWishart(Deterministic):
"""
Transform gamma scalar moments to ndim=0 scalar Wishart moments
"""

_parent_moments = [GammaMoments()]

@ensureparents
def __init__(self, alpha, **kwargs):

# Check for constant
if misc.is_numeric(alpha):
alpha = Constant(Gamma)(alpha)

# FIXME: Put import here to avoid circular dependency import
from .wishart import WishartMoments
self._moments = WishartMoments(())
dims = ( (), () )

# Construct the node
super().__init__(alpha,
dims=self._moments.dims,
**kwargs)

def get_moments(self):
return self.parents[0].get_moments()

@staticmethod
def _compute_message_to_parent(index, m_children, *u_parents):
return m_children