################################################################################
# Copyright (C) 2011-2014 Jaakko Luttinen
#
# This file is licensed under the MIT License.
################################################################################
import numpy as np
from bayespy.utils import misc
from .node import Node
from .deterministic import Deterministic
from .gaussian import Gaussian, GaussianMoments
from .gaussian import GaussianGammaMoments
from .ml import DeltaMoments
[docs]class SumMultiply(Deterministic):
r"""
Node for computing general products and sums of Gaussian nodes.
The node is similar to `numpy.einsum`, which is a very general
function for computing dot products, sums, products and other sums
of products of arrays.
For instance, consider the following arrays:
>>> import numpy as np
>>> X = np.random.randn(2, 3, 4)
>>> Y = np.random.randn(3, 5)
>>> Z = np.random.randn(4, 2)
Then, the Einstein summation can be used as:
>>> np.einsum('abc,bd,ca->da', X, Y, Z)
array([[...]])
SumMultiply node can be used similarly for Gaussian nodes. For instance,
consider the following Gaussian nodes:
>>> from bayespy.nodes import GaussianARD
>>> X = GaussianARD(0, 1, shape=(2, 3, 4))
>>> Y = GaussianARD(0, 1, shape=(3, 5))
>>> Z = GaussianARD(0, 1, shape=(4, 2))
Then, similarly to `numpy.einsum`, SumMultiply could be used as:
>>> from bayespy.nodes import SumMultiply
>>> SumMultiply('abc,bd,ca->da', X, Y, Z)
<bayespy.inference.vmp.nodes.dot.SumMultiply object at 0x...>
or
>>> SumMultiply(X, [0,1,2], Y, [1,3], Z, [2,0], [3,0])
<bayespy.inference.vmp.nodes.dot.SumMultiply object at 0x...>
which is similar to the alternative syntax of numpy.einsum.
This node operates similarly as numpy.einsum. However, you must use all the
elements of each node, that is, an operation like np.einsum('ii->i',X) is
not allowed. Thus, for each node, each axis must be given unique id. The id
identifies which axes correspond to which axes between the different
nodes. Also, Ellipsis ('...') is not yet supported for simplicity. It would
also have some problems with constant inputs (because how to determine
ndim), so let us just forget it for now.
Each output axis must appear in the input mappings.
The keys must refer to variable dimension axes only, not plate axes.
The input nodes may be Gaussian-gamma (isotropic) nodes.
The output message is Gaussian-gamma (isotropic) if any of the input nodes
is Gaussian-gamma.
Examples
--------
Sum over the rows:
'ij->j'
Inner product of three vectors:
'i,i,i'
Matrix-vector product:
'ij,j->i'
Matrix-matrix product:
'ik,kj->ij'
Outer product:
'i,j->ij'
Vector-matrix-vector product:
'i,ij,j'
Notes
-----
This operation can be extremely slow if not used wisely. For large and
complex operations, it is sometimes more efficient to split the operation
into multiple nodes. For instance, the example above could probably be
computed faster by
>>> XZ = SumMultiply(X, [0,1,2], Z, [2,0], [0,1])
>>> F = SumMultiply(XZ, [0,1], Y, [1,2], [2,0])
because the third axis ('c') could be summed out already in the first
operation. This same effect applies also to numpy.einsum in general.
"""
[docs] def __init__(self, *args, iterator_axis=None, **kwargs):
"""
SumMultiply(Node1, map1, Node2, map2, ..., NodeN, mapN [, map_out])
"""
args = list(args)
if len(args) < 2:
raise ValueError("Not enough inputs")
if iterator_axis is not None:
raise NotImplementedError("Iterator axis not implemented yet")
if iterator_axis is not None and not isinstance(iterator_axis, int):
raise ValueError("Iterator axis must be integer")
# Two different parsing methods, depends on how the arguments are given
if misc.is_string(args[0]):
# This is the format:
# SumMultiply('ik,k,kj->ij', X, Y, Z)
strings = args[0]
nodes = args[1:]
# Remove whitespace
strings = misc.remove_whitespace(strings)
# Split on '->' (should contain only one '->' or none)
strings = strings.split('->')
if len(strings) > 2:
raise ValueError('The string contains too many ->')
strings_in = strings[0]
if len(strings) == 2:
string_out = strings[1]
else:
string_out = ''
# Split former part on ',' (the number of parts should be equal to
# nodes)
strings_in = strings_in.split(',')
if len(strings_in) != len(nodes):
raise ValueError('Number of given input nodes is different '
'from the input keys in the string')
# Split strings into key lists using single character keys
keysets = [list(string_in) for string_in in strings_in]
keys_out = list(string_out)
else:
# This is the format:
# SumMultiply(X, [0,2], Y, [2], Z, [2,1], [0,1])
# If given, the output mapping is the last argument
if len(args) % 2 == 0:
keys_out = []
else:
keys_out = args.pop(-1)
# Node and axis mapping are given in turns
nodes = args[::2]
keysets = args[1::2]
# Find all the keys (store only once each)
full_keyset = []
for keyset in keysets:
full_keyset += keyset
#full_keyset += list(keyset.keys())
full_keyset = list(set(full_keyset))
# Input and output messages are Gaussian unless there is at least one
# Gaussian-gamma message from the parents
self.gaussian_gamma = False
self.is_constant = []
for i in range(len(nodes)):
try:
# First, try to handle the node as a "constant" with so-called
# delta moments. These consume much less memory as only the
# "value" (or the first moment) is used.
nodes[i] = self._ensure_moments(
nodes[i],
DeltaMoments,
ndim=len(keysets[i])
)
except DeltaMoments.NoConverterError:
self.is_constant.append(False)
try:
# Second, try to convert to gaussian moments. These moments
# will consume more memory as it calculates the second
# moment too.
nodes[i] = self._ensure_moments(
nodes[i],
GaussianMoments,
ndim=len(keysets[i])
)
except GaussianMoments.NoConverterError:
self.gaussian_gamma = True
else:
self.is_constant.append(True)
if self.gaussian_gamma:
nodes = [
(
self._ensure_moments(
node,
GaussianGammaMoments,
ndim=len(keyset)
)
if not is_const else
node
)
for (node, keyset, is_const) in zip(nodes, keysets, self.is_constant)
]
self._parent_moments = tuple(node._moments for node in nodes)
#
# Check the validity of each node
#
for n in range(len(nodes)):
# Check that the maps and the size of the variable are consistent
if len(nodes[n].dims[0]) != len(keysets[n]):
raise ValueError("Wrong number of keys (%d) for the node "
"number %d with %d dimensions"
% (len(keysets[n]),
n,
len(nodes[n].dims[0])))
# Check that the keys are unique
if len(set(keysets[n])) != len(keysets[n]):
raise ValueError("Axis keys for node number %d are not unique"
% n)
# Check the validity of output keys: each output key must be included in
# the input keys
if len(keys_out) != len(set(keys_out)):
raise ValueError("Output keys are not unique")
for key in keys_out:
if key not in full_keyset:
raise ValueError("Output key %s does not appear in any input"
% key)
# Check the validity of the nodes with respect to the key mapping.
# Check that the node dimensions map and broadcast properly, that is,
# all the nodes using the same key for axes must have equal size for
# those axes (or size 1).
broadcasted_size = {}
for key in full_keyset:
broadcasted_size[key] = 1
for (node, keyset) in zip(nodes, keysets):
try:
# Find the axis for the key
index = keyset.index(key)
except ValueError:
# OK, this node doesn't use this key for any axis
pass
else:
# Length of the axis for that key
node_size = node.dims[0][index]
if node_size != broadcasted_size[key]:
if broadcasted_size[key] == 1:
# Apply broadcasting
broadcasted_size[key] = node_size
elif node_size != 1:
# Different sizes and neither has size 1
raise ValueError("Axes using key %s do not "
"broadcast properly"
% key)
# Compute the shape of the output
shape = tuple([broadcasted_size[key] for key in keys_out])
if self.gaussian_gamma:
self._moments = GaussianGammaMoments(shape)
else:
self._moments = GaussianMoments(shape)
# Rename the keys to [0,1,...,N-1] where N is the total number of keys
self.N_keys = len(full_keyset)
self.out_keys = [full_keyset.index(key) for key in keys_out]
self.in_keys = [ [full_keyset.index(key) for key in keyset]
for keyset in keysets ]
super().__init__(*nodes,
dims=self._moments.dims,
**kwargs)
def _compute_function(self, *x_parents):
# TODO: Add unit tests for this function
(xs, alphas) = (
(x_parents, 1) if not self.gaussian_gamma else
zip(*x_parents)
)
# Add Ellipsis for the plates
in_keys = [[Ellipsis] + k for k in self.in_keys]
out_keys = [Ellipsis] + self.out_keys
samples_and_keys = misc.zipper_merge(xs, in_keys)
y = np.einsum(*(samples_and_keys + [out_keys]))
return (
y if not self.gaussian_gamma else
(y, misc.multiply(*alphas))
)
def _compute_moments(self, *u_parents):
# Compute the number of plate axes for each node
plate_counts0 = [
(np.ndim(u_parent[0]) - len(keys))
for (keys,u_parent) in zip(self.in_keys, u_parents)
]
plate_counts1 = [
(
# Gaussian moments: Use second moments "matrix"
(np.ndim(u_parent[1]) - 2*len(keys))
if not is_const else
# Delta moments: Use first moment "vector"
(np.ndim(u_parent[0]) - len(keys))
)
for (keys, u_parent, is_const) in zip(
self.in_keys,
u_parents,
self.is_constant
)
]
# The number of plate axes for the output
N0 = max(plate_counts0)
N1 = max(plate_counts1)
# The total number of unique keys used (keys are 0,1,...,N_keys-1)
D = self.N_keys
#
# Compute the mean
#
out_all_keys = list(range(D+N0-1, D-1, -1)) + self.out_keys
#nodes_dim_keys = self.nodes_dim_keys
in_all_keys = [list(range(D+plate_count-1, D-1, -1)) + keys
for (plate_count, keys) in zip(plate_counts0,
self.in_keys)]
u0 = [u[0] for u in u_parents]
args = misc.zipper_merge(u0, in_all_keys) + [out_all_keys]
x0 = np.einsum(*args)
#
# Compute the covariance
#
out_all_keys = (list(range(2*D+N1-1, 2*D-1, -1))
+ [D+key for key in self.out_keys]
+ self.out_keys)
in_all_keys = [
x
for (plate_count, node_keys, is_const) in zip(
plate_counts1,
self.in_keys,
self.is_constant,
)
for x in (
# Gaussian moments: Use the second moment
[
list(range(2*D+plate_count-1, 2*D-1, -1))
+ [D+key for key in node_keys]
+ node_keys
]
if not is_const else
# Delta moments: Use the first moment tiwce
[
(
list(range(2*D+plate_count-1, 2*D-1, -1))
+ [D+key for key in node_keys]
),
(
list(range(2*D+plate_count-1, 2*D-1, -1))
+ node_keys
),
]
)
]
u1 = [
x
for (u, is_const) in zip(u_parents, self.is_constant)
for x in (
# Gaussian moments: Use the second moment
[u[1]]
if not is_const else
# Delta moments: Use the first moment twice
[u[0], u[0]]
)
]
args = misc.zipper_merge(u1, in_all_keys) + [out_all_keys]
x1 = np.einsum(*args)
if not self.gaussian_gamma:
return [x0, x1]
# Compute Gaussian-gamma specific moments
x2 = 1
x3 = 0
for i in range(len(u_parents)):
x2 = x2 * (1 if self.is_constant[i] else u_parents[i][2])
x3 = x3 + (0 if self.is_constant[i] else u_parents[i][3])
return [x0, x1, x2, x3]
[docs] def get_parameters(self):
# Compute mean and variance
u = self.get_moments()
u[1] -= u[0]**2
return u
def _message_to_parent(self, index, u_parent=None):
"""
Compute the message and mask to a parent node.
"""
if self.is_constant[index]:
raise NotImplementedError(
"Message to DeltaMoments parent not yet implemented."
)
# Check index
if index >= len(self.parents):
raise ValueError("Parent index larger than the number of parents")
# Get messages from other parents and children
u_parents = self._message_from_parents(exclude=index)
m = self._message_from_children()
mask = self.mask
# Normally we don't need to care about masks when computing the
# message. However, in this node we want to avoid computing huge message
# arrays so we sum some axes already here. Thus, we need to apply the
# mask.
#
# Actually, we don't need to care about masks because the message from
# children has already been masked.
parent = self.parents[index]
#
# Compute the first message
#
msg = [None, None]
# Compute the two messages
for ind in range(2):
# The total number of keys for the non-plate dimensions
N = (ind+1) * self.N_keys
parent_num_dims = len(parent.dims[ind])
parent_num_plates = len(parent.plates)
parent_plate_keys = list(range(N + parent_num_plates,
N,
-1))
parent_dim_keys = self.in_keys[index]
if ind == 1:
parent_dim_keys = ([key + self.N_keys
for key in self.in_keys[index]]
+ parent_dim_keys)
args = []
# This variable counts the maximum number of plates of the
# arguments, thus it will tell the number of plates in the result
# (if the artificially added plates above were ignored).
result_num_plates = 0
result_plates = ()
# Mask and its keysr
mask_num_plates = np.ndim(mask)
mask_plates = np.shape(mask)
mask_plate_keys = list(range(N + mask_num_plates,
N,
-1))
result_num_plates = max(result_num_plates,
mask_num_plates)
result_plates = misc.broadcasted_shape(result_plates,
mask_plates)
# Moments and keys of other parents
for (k, u) in enumerate(u_parents):
if k != index:
num_dims = (
(ind+1) * len(self.in_keys[k])
if not self.is_constant[k] else
len(self.in_keys[k])
)
ui = (
u[ind] if not self.is_constant[k] else
u[0]
)
num_plates = np.ndim(ui) - num_dims
plates = np.shape(ui)[:num_plates]
plate_keys = list(range(N + num_plates,
N,
-1))
if ind == 0:
args.append(ui)
args.append(plate_keys + self.in_keys[k])
else:
in_keys2 = [key + self.N_keys for key in self.in_keys[k]]
if not self.is_constant[k]:
# Gaussian moments: Use second moment once
args.append(ui)
args.append(plate_keys + in_keys2 + self.in_keys[k])
else:
# Delta moments: Use first moment twice
args.append(ui)
args.append(plate_keys + self.in_keys[k])
args.append(ui)
args.append(plate_keys + in_keys2)
result_num_plates = max(result_num_plates, num_plates)
result_plates = misc.broadcasted_shape(result_plates,
plates)
# Message and keys from children
child_num_dims = (ind+1) * len(self.out_keys)
child_num_plates = np.ndim(m[ind]) - child_num_dims
child_plates = np.shape(m[ind])[:child_num_plates]
child_plate_keys = list(range(N + child_num_plates,
N,
-1))
child_dim_keys = self.out_keys
if ind == 1:
child_dim_keys = ([key + self.N_keys
for key in self.out_keys]
+ child_dim_keys)
args.append(m[ind])
args.append(child_plate_keys + child_dim_keys)
result_num_plates = max(result_num_plates, child_num_plates)
result_plates = misc.broadcasted_shape(result_plates,
child_plates)
# Output keys, that is, the keys of the parent[index]
parent_keys = parent_plate_keys + parent_dim_keys
# Performance trick: Check which axes can be summed because they
# have length 1 or are non-existing in parent[index]. Thus, remove
# keys corresponding to unit length axes in parent[index] so that
# einsum sums over those axes. After computations, these axes must
# be added back in order to get the correct shape for the message.
# Also, remove axes/keys that are in output (parent[index]) but not in
# any inputs (children and other parents).
parent_shape = parent.get_shape(ind)
removed_axes = []
for j in range(len(parent_keys)):
if parent_shape[j] == 1:
# Remove the key (take into account the number of keys that
# have already been removed)
del parent_keys[j-len(removed_axes)]
removed_axes.append(j)
else:
# Remove the key if it doesn't appear in any of the
# messages from children or other parents.
if not np.any([parent_keys[j-len(removed_axes)] in keys
for keys in args[1::2]]):
del parent_keys[j-len(removed_axes)]
removed_axes.append(j)
args.append(parent_keys)
# THE BEEF: Compute the message
msg[ind] = np.einsum(*args)
# Find the correct shape for the message array
message_shape = list(np.shape(msg[ind]))
# First, add back the axes with length 1
for ax in removed_axes:
message_shape.insert(ax, 1)
# Second, remove leading axes for plates that were not present in
# the child nor other parents' messages. This is not really
# necessary, but it is just elegant to remove the leading unit
# length axes that we added artificially at the beginning just
# because we wanted the key mapping to be simple.
if parent_num_plates > result_num_plates:
del message_shape[:(parent_num_plates-result_num_plates)]
# Then, the actual reshaping
msg[ind] = np.reshape(msg[ind], message_shape)
# Broadcasting is not supported for variable dimensions, thus force
# explicit correct shape for variable dimensions
var_dims = parent.dims[ind]
msg[ind] = msg[ind] * np.ones(var_dims)
# Apply plate multiplier: If this node has non-unit plates that are
# unit plates in the parent, those plates are summed. However, if
# the message has unit axis for that plate, it should be first
# broadcasted to the plates of this node and then summed to the
# plates of the parent. In order to avoid this broadcasting and
# summing, it is more efficient to just multiply by the correct
# factor.
r = self.broadcasting_multiplier(self.plates,
result_plates,
parent.plates)
if r != 1:
msg[ind] *= r
if self.gaussian_gamma:
alphas = [
(u_parents[i][2] if not is_const else 1.0)
for (i, is_const) in zip(range(len(u_parents)), self.is_constant)
if i != index
]
m2 = self._compute_message(m[2], mask, *alphas,
ndim=0,
plates_from=self.plates,
plates_to=parent.plates)
m3 = self._compute_message(m[3], mask,
ndim=0,
plates_from=self.plates,
plates_to=parent.plates)
msg = msg + [m2, m3]
return msg
[docs]def Dot(*args, **kwargs):
"""
Node for computing inner product of several Gaussian vectors.
This is a simple wrapper of the much more general SumMultiply. For now, it
is here for backward compatibility.
"""
einsum = 'i' + ',i'*(len(args)-1)
return SumMultiply(einsum, *args, **kwargs)