################################################################################
# Copyright (C) 2013-2014 Jaakko Luttinen
#
# This file is licensed under the MIT License.
################################################################################
import warnings
import numpy as np
from bayespy.utils import misc
from .node import ensureparents
from .stochastic import Stochastic, Distribution
[docs]class ExponentialFamilyDistribution(Distribution):
"""
Sub-classes implement distribution specific computations.
"""
#
# The following methods are for ExponentialFamily distributions
#
[docs] def compute_message_to_parent(self, parent, index, u_self, *u_parents):
raise NotImplementedError()
[docs] def compute_phi_from_parents(self, *u_parents, mask=True):
raise NotImplementedError()
[docs] def compute_moments_and_cgf(self, phi, mask=True):
raise NotImplementedError()
#
# The following methods are for Mixture class
#
[docs] def compute_cgf_from_parents(self, *u_parents):
raise NotImplementedError()
[docs] def compute_fixed_moments_and_f(self, x, mask=True):
raise NotImplementedError()
[docs] def compute_logpdf(self, u, phi, g, f, ndims):
""" Compute E[log p(X)] given E[u], E[phi], E[g] and
E[f]. Does not sum over plates."""
# TODO/FIXME: Should I take into account what is latent or
# observed, or what is even totally ignored (by the mask).
L = g + f
for (phi_i, u_i, ndims_i) in zip(phi, u, ndims):
# Axes to sum (dimensions of the variable, not the plates)
axis_sum = tuple(range(-ndims_i,0))
# Compute the term
# TODO/FIXME: Use einsum!
L = L + np.sum(
np.where(u_i != 0, phi_i, 0) * u_i,
axis=axis_sum
)
return L
[docs] def compute_gradient(self, g, u, phi):
r"""
Compute the standard gradient with respect to the natural parameters.
"""
raise NotImplementedError("Standard gradient not yet implemented for %s"
% (self.__class__.__name__))
def useconstructor(__init__):
def constructor_decorator(self, *args, **kwargs):
if (self.dims is None or
self._distribution is None or
self._moments is None or
self._parent_moments is None):
(args, kwargs, dims, plates, dist, stats, pstats) = \
self._constructor(*args, **kwargs)
self.dims = dims
self._distribution = dist
self._moments = stats
self._parent_moments = pstats
self.plates = plates
__init__(self, *args, **kwargs)
return constructor_decorator
[docs]class ExponentialFamily(Stochastic):
"""
A base class for nodes using natural parameterization `phi`.
phi
Sub-classes must implement the following static methods:
_compute_message_to_parent(index, u_self, *u_parents)
_compute_phi_from_parents(*u_parents, mask)
_compute_moments_and_cgf(phi, mask)
_compute_fixed_moments_and_f(x, mask=True)
Sub-classes may need to re-implement:
1. If they manipulate plates:
_compute_weights_to_parent(index, weights)
_compute_plates_to_parent(self, index, plates)
_compute_plates_from_parent(self, index, plates)
"""
# Sub-classes should overwrite this (possibly using _constructor)
dims = None
# Sub-classes should overwrite this
_distribution = None
[docs] @useconstructor
def __init__(self, *parents, initialize=True, phi_bias=None, **kwargs):
self.annealing = 1.0
# Terms for the lower bound (G for latent and F for observed)
self.g = np.array(np.nan)
self.f = np.array(np.nan)
self._phi_bias = phi_bias if phi_bias is not None else len(self.dims) * [0.0]
super().__init__(*parents,
initialize=initialize,
dims=self.dims,
**kwargs)
if not initialize:
axes = len(self.plates)*(1,)
self.phi = [misc.nans(axes+dim) for dim in self.dims]
@classmethod
@ensureparents
def _constructor(cls, *parents, **kwargs):
"""
Constructs distribution and moments objects.
If __init__ uses useconstructor decorator, this method is called to
construct distribution and moments objects.
The method is given the same inputs as __init__. For some nodes, some of
these can't be "static" class attributes, then the node class must
overwrite this method to construct the objects manually.
The point of distribution class is to move general distribution but
not-node specific code. The point of moments class is to define the
messaging protocols.
"""
parent_plates = [cls._distribution.plates_from_parent(ind, parent.plates)
for (ind, parent) in enumerate(parents)]
return (parents,
kwargs,
cls.dims,
cls._total_plates(kwargs.get('plates'), *parent_plates),
cls._distribution,
cls._moments,
cls._parent_moments)
def _initialize_from_parent_moments(self, *u_parents):
if not np.all(self.observed):
# Update natural parameters using parents
self._update_phi_from_parents(*u_parents)
# Update moments
mask = np.logical_not(self.observed)
(u, g) = self._distribution.compute_moments_and_cgf(self.phi,
mask=mask)
# TODO/FIXME/BUG: You should use observation mask in order to not
# overwrite them!
self._set_moments_and_cgf(u, g, mask=mask)
[docs] def initialize_from_prior(self):
u_parents = self._message_from_parents()
self._initialize_from_parent_moments(*u_parents)
[docs] def initialize_from_parameters(self, *args):
u_parents = [p_mom.compute_fixed_moments(x)
for (p_mom, x) in zip(self._parent_moments, args)]
self._initialize_from_parent_moments(*u_parents)
[docs] def initialize_from_value(self, x, *args):
# Update moments from value
mask = np.logical_not(self.observed)
u = self._moments.compute_fixed_moments(x, *args)
# Check that the shape is correct
for i in range(len(u)):
ndim = len(self.dims[i])
if ndim > 0:
if np.shape(u[i])[-ndim:] != self.dims[i]:
raise ValueError("The initial value for node %s has invalid shape %s."
% (np.shape(x)))
self._set_moments_and_cgf(u, np.inf, mask=mask)
[docs] def initialize_from_random(self):
"""
Set the variable to a random sample from the current distribution.
"""
#self.initialize_from_prior()
X = self.random()
self.initialize_from_value(X)
def _update_phi_from_parents(self, *u_parents):
# TODO/FIXME: Could this be combined to the function
# _update_distribution_and_lowerbound ?
# No, because some initialization methods may want to use this.
# This makes correct broadcasting
self.phi = [
a + b
for (a, b) in zip(
self._distribution.compute_phi_from_parents(*u_parents),
self._phi_bias
)
]
# Make sure phi has the correct number of axes. It makes life
# a bit easier elsewhere.
for i in range(len(self.phi)):
axes = len(self.plates) + self.ndims[i] - np.ndim(self.phi[i])
if axes > 0:
# Add axes
self.phi[i] = misc.add_leading_axes(self.phi[i], axes)
elif axes < 0:
# Remove extra leading axes
first = -(len(self.plates)+self.ndims[i])
sh = np.shape(self.phi[i])[first:]
self.phi[i] = np.reshape(self.phi[i], sh)
# Check that the shape is correct
if not misc.is_shape_subset(np.shape(self.phi[i]),
self.get_shape(i)):
raise ValueError("Incorrect shape of phi[%d] in node class %s. "
"Shape is %s but it should be broadcastable "
"to shape %s."
% (i,
self.__class__.__name__,
np.shape(self.phi[i]),
self.get_shape(i)))
def _set_moments_and_cgf(self, u, g, mask=True):
self._set_moments(u, mask=mask)
self.g = np.where(mask, g, self.g)
return
[docs] def get_riemannian_gradient(self):
r"""
Computes the Riemannian/natural gradient.
"""
u_parents = self._message_from_parents()
m_children = self._message_from_children()
# TODO/FIXME: Put observed plates to zero?
# Compute the gradient
phi = [
a + b
for (a, b) in zip(
self._distribution.compute_phi_from_parents(*u_parents),
self._phi_bias
)
]
for i in range(len(self.phi)):
phi[i] = self.annealing * (phi[i] + m_children[i]) - self.phi[i]
phi[i] = phi[i] * np.ones(self.get_shape(i))
return phi
[docs] def get_gradient(self, rg):
r""" Computes gradient with respect to the natural parameters.
The function takes the Riemannian gradient as an input. This is for
three reasons: 1) You probably want to use the Riemannian gradient
anyway so this helps avoiding accidental use of this function. 2) The
gradient is computed by using the Riemannian gradient and chain rules.
3) Probably you need both Riemannian and normal gradients anyway so you
can provide it to this function to avoid re-computing it."""
g = self._distribution.compute_gradient(rg, self.u, self.phi)
for i in range(len(g)):
g[i] /= self.annealing
return g
## def update_parameters(self, d, scale=1.0):
## r"""
## Update the parameters of the VB distribution given a change.
## The parameters should be such that they can be used for
## optimization, that is, use log transformation for positive
## parameters.
## """
## phi = self.get_parameters()
## for i in range(len(phi)):
## phi[i] = phi[i] + scale*d[i]
## self.set_parameters(phi)
## return
[docs] def get_parameters(self):
r"""
Return parameters of the VB distribution.
The parameters should be such that they can be used for
optimization, that is, use log transformation for positive
parameters.
"""
return [np.copy(p) for p in self.phi]
def _decode_parameters(self, x):
return [np.copy(p) for p in x]
[docs] def set_parameters(self, x):
r"""
Set the parameters of the VB distribution.
The parameters should be such that they can be used for
optimization, that is, use log transformation for positive
parameters.
"""
self.phi = self._decode_parameters(x)
self._update_moments_and_cgf()
return
def _update_distribution_and_lowerbound(self, m_children, *u_parents):
# Update phi first from parents..
self._update_phi_from_parents(*u_parents)
# .. then just add children's message
self.phi = [self.annealing * (phi + m)
for (phi, m) in zip(self.phi, m_children)]
# Update u and g
self._update_moments_and_cgf()
def _update_moments_and_cgf(self):
"""
Update moments and cgf based on current phi.
"""
# Mask for plates to update (i.e., unobserved plates)
update_mask = np.logical_not(self.observed)
# Compute the moments (u) and CGF (g)...
(u, g) = self._distribution.compute_moments_and_cgf(self.phi,
mask=update_mask)
# ... and store them
self._set_moments_and_cgf(u, g, mask=update_mask)
[docs] def observe(self, x, *args, mask=True):
"""
Fix moments, compute f and propagate mask.
"""
# Compute fixed moments
(u, f) = self._distribution.compute_fixed_moments_and_f(x, *args,
mask=mask)
# # Check the dimensionality of the observations
# self._check_shape()
# for (i,v) in enumerate(u):
# # This is what the dimensionality "should" be
# s = self.plates + self.dims[i]
# t = np.shape(v)
# if s != t:
# msg = "Dimensionality of the observations incorrect."
# msg += "\nShape of input: " + str(t)
# msg += "\nExpected shape: " + str(s)
# msg += "\nCheck plates."
# raise Exception(msg)
# Set the moments. Shape checking is done there.
self._set_moments(u, mask=mask, broadcast=False)
self.f = np.where(mask, f, self.f)
# Observed nodes should not be ignored
self.observed = mask
self._update_mask()
[docs] def lower_bound_contribution(self, gradient=False, ignore_masked=True):
r"""Compute E[ log p(X|parents) - log q(X) ]
If deterministic annealing is used, the term E[ -log q(X) ] is
divided by the anneling coefficient. That is, phi and cgf of q
are multiplied by the temperature (inverse annealing
coefficient).
"""
# Annealing temperature
T = 1 / self.annealing
# Messages from parents
u_parents = self._message_from_parents()
phi = [
a # + b # TODO: Should the bias be added here or not?
for (a, b) in zip(
self._distribution.compute_phi_from_parents(*u_parents),
self._phi_bias
)
]
# G from parents
L = self._distribution.compute_cgf_from_parents(*u_parents)
# G for unobserved variables (ignored variables are handled properly
# automatically)
latent_mask = np.logical_not(self.observed)
# G and F
if np.all(self.observed):
z = np.nan
elif T == 1:
z = -self.g
else:
z = -T * self.g
## TRIED THIS BUT IT WAS WRONG:
## z = -T * self.g + (1-T) * self.f
## if np.any(np.isnan(self.f)):
## warnings.warn("F(x) not implemented for node %s. This "
## "is required for annealed lower bound "
## "computation." % self.__class__.__name__)
##
## It was wrong because the optimal q distribution has f which is
## weighted by 1/T and here the f of q is weighted by T so the
## total weight is 1, thus it cancels out with f of p.
L = L + np.where(self.observed, self.f, z)
for (phi_p, phi_q, u_q, dims) in zip(phi, self.phi, self.u, self.dims):
# Form a mask which puts observed variables to zero and
# broadcasts properly
latent_mask_i = misc.add_trailing_axes(
misc.add_leading_axes(
latent_mask,
len(self.plates) - np.ndim(latent_mask)),
len(dims))
axis_sum = tuple(range(-len(dims),0))
# Compute the term
phi_q = np.where(latent_mask_i, phi_q, 0)
# Apply annealing
phi_diff = phi_p - T * phi_q
# Handle 0 * -inf
phi_diff = np.where(u_q != 0, phi_diff, 0)
# TODO/FIXME: Use einsum here?
Z = np.sum(phi_diff * u_q, axis=axis_sum)
L = L + Z
if ignore_masked:
return (np.sum(np.where(self.mask, L, 0))
* self.broadcasting_multiplier(self.plates,
np.shape(L),
np.shape(self.mask))
* np.prod(self.plates_multiplier))
else:
return (np.sum(L)
* self.broadcasting_multiplier(self.plates,
np.shape(L))
* np.prod(self.plates_multiplier))
[docs] def logpdf(self, X, mask=True):
"""
Compute the log probability density function Q(X) of this node.
"""
if mask is not True:
raise NotImplementedError('Mask not yet implemented')
(u, f) = self._distribution.compute_fixed_moments_and_f(X, mask=mask)
Z = 0
for (phi_d, u_d, dims) in zip(self.phi, u, self.dims):
axis_sum = tuple(range(-len(dims),0))
# TODO/FIXME: Use einsum here?
Z = Z + np.sum(phi_d * u_d, axis=axis_sum)
#Z = Z + misc.sum_multiply(phi_d, u_d, axis=axis_sum)
return (self.g + f + Z)
[docs] def pdf(self, X, mask=True):
"""
Compute the probability density function of this node.
"""
return np.exp(self.logpdf(X, mask=mask))
def _save(self, group):
"""
Save the state of the node into a HDF5 file.
group can be the root
"""
## if name is None:
## name = self.name
## subgroup = group.create_group(name)
for i in range(len(self.phi)):
misc.write_to_hdf5(group, self.phi[i], 'phi%d' % i)
misc.write_to_hdf5(group, self.f, 'f')
misc.write_to_hdf5(group, self.g, 'g')
super()._save(group)
def _load(self, group):
"""
Load the state of the node from a HDF5 file.
"""
# TODO/FIXME: Check that the shapes are correct!
for i in range(len(self.phi)):
phii = group['phi%d' % i][...]
self.phi[i] = phii
self.f = group['f'][...]
self.g = group['g'][...]
super()._load(group)
[docs] def random(self):
"""
Draw a random sample from the distribution.
"""
return self._distribution.random(*(self.phi), plates=self.plates)