def check(shape, plates, einsum_x, einsum_xx, axis=-1):
# TODO/FIXME: Improve by having non-diagonal precision/covariance
# parameter for the Gaussian X
D = shape[axis]
X = GaussianARD(np.random.randn(*(plates + shape)),
np.random.rand(*(plates + shape)),
shape=shape,
plates=plates)
(x, xx) = X.get_moments()
R = np.random.randn(D, D)
X.rotate(R, axis=axis)
(rx, rxxr) = X.get_moments()
self.assertAllClose(rx, np.einsum(einsum_x, R, x))
self.assertAllClose(rxxr, np.einsum(einsum_xx, R, xx, R))
pass
def check(shape, plates, einsum_x, einsum_xx, axis=-1):
# TODO/FIXME: Improve by having non-diagonal precision/covariance
# parameter for the Gaussian X
D = shape[axis]
X = GaussianARD(np.random.randn(*(plates+shape)),
np.random.rand(*(plates+shape)),
shape=shape,
plates=plates)
(x, xx) = X.get_moments()
R = np.random.randn(D,D)
X.rotate(R, axis=axis)
(rx, rxxr) = X.get_moments()
self.assertAllClose(rx,
np.einsum(einsum_x, R, x))
self.assertAllClose(rxxr,
np.einsum(einsum_xx, R, xx, R))
pass
def test_rotate_plates(self):
# Basic test for Gaussian vectors
X = GaussianARD(np.random.randn(3,2),
np.random.rand(3,2),
shape=(2,),
plates=(3,))
(u0, u1) = X.get_moments()
Cov = u1 - linalg.outer(u0, u0, ndim=1)
Q = np.random.randn(3,3)
Qu0 = np.einsum('ik,kj->ij', Q, u0)
QCov = np.einsum('k,kij->kij', np.sum(Q, axis=0)**2, Cov)
Qu1 = QCov + linalg.outer(Qu0, Qu0, ndim=1)
X.rotate_plates(Q, plate_axis=-1)
(u0, u1) = X.get_moments()
self.assertAllClose(u0, Qu0)
self.assertAllClose(u1, Qu1)
# Test full covariance, that is, with observations
X = GaussianARD(np.random.randn(3,2),
np.random.rand(3,2),
shape=(2,),
plates=(3,))
Y = Gaussian(X, [[2.0, 1.5], [1.5, 3.0]],
plates=(3,))
Y.observe(np.random.randn(3,2))
X.update()
(u0, u1) = X.get_moments()
Cov = u1 - linalg.outer(u0, u0, ndim=1)
Q = np.random.randn(3,3)
Qu0 = np.einsum('ik,kj->ij', Q, u0)
QCov = np.einsum('k,kij->kij', np.sum(Q, axis=0)**2, Cov)
Qu1 = QCov + linalg.outer(Qu0, Qu0, ndim=1)
X.rotate_plates(Q, plate_axis=-1)
(u0, u1) = X.get_moments()
self.assertAllClose(u0, Qu0)
self.assertAllClose(u1, Qu1)
pass
def test_message_to_child(self):
"""
Test the message to child of Concatenate node.
"""
# Two parents without shapes
X1 = GaussianARD(0, 1, plates=(2, ), shape=())
X2 = GaussianARD(0, 1, plates=(3, ), shape=())
Y = Concatenate(X1, X2)
u1 = X1.get_moments()
u2 = X2.get_moments()
u = Y.get_moments()
self.assertAllClose((u[0] * np.ones((5, )))[:2], u1[0] * np.ones(
(2, )))
self.assertAllClose((u[1] * np.ones((5, )))[:2], u1[1] * np.ones(
(2, )))
self.assertAllClose((u[0] * np.ones((5, )))[2:], u2[0] * np.ones(
(3, )))
self.assertAllClose((u[1] * np.ones((5, )))[2:], u2[1] * np.ones(
(3, )))
# Two parents with shapes
X1 = GaussianARD(0, 1, plates=(2, ), shape=(4, ))
X2 = GaussianARD(0, 1, plates=(3, ), shape=(4, ))
Y = Concatenate(X1, X2)
u1 = X1.get_moments()
u2 = X2.get_moments()
u = Y.get_moments()
self.assertAllClose((u[0] * np.ones((5, 4)))[:2], u1[0] * np.ones(
(2, 4)))
self.assertAllClose((u[1] * np.ones((5, 4, 4)))[:2], u1[1] * np.ones(
(2, 4, 4)))
self.assertAllClose((u[0] * np.ones((5, 4)))[2:], u2[0] * np.ones(
(3, 4)))
self.assertAllClose((u[1] * np.ones((5, 4, 4)))[2:], u2[1] * np.ones(
(3, 4, 4)))
# Test with non-constant axis
X1 = GaussianARD(0, 1, plates=(2, 4), shape=())
X2 = GaussianARD(0, 1, plates=(3, 4), shape=())
Y = Concatenate(X1, X2, axis=-2)
u1 = X1.get_moments()
u2 = X2.get_moments()
u = Y.get_moments()
self.assertAllClose((u[0] * np.ones((5, 4)))[:2], u1[0] * np.ones(
(2, 4)))
self.assertAllClose((u[1] * np.ones((5, 4)))[:2], u1[1] * np.ones(
(2, 4)))
self.assertAllClose((u[0] * np.ones((5, 4)))[2:], u2[0] * np.ones(
(3, 4)))
self.assertAllClose((u[1] * np.ones((5, 4)))[2:], u2[1] * np.ones(
(3, 4)))
# Test with constant parent
X1 = np.random.randn(2, 4)
X2 = GaussianARD(0, 1, plates=(3, ), shape=(4, ))
Y = Concatenate(X1, X2)
u1 = Y.parents[0].get_moments()
u2 = X2.get_moments()
u = Y.get_moments()
self.assertAllClose((u[0] * np.ones((5, 4)))[:2], u1[0] * np.ones(
(2, 4)))
self.assertAllClose((u[1] * np.ones((5, 4, 4)))[:2], u1[1] * np.ones(
(2, 4, 4)))
self.assertAllClose((u[0] * np.ones((5, 4)))[2:], u2[0] * np.ones(
(3, 4)))
self.assertAllClose((u[1] * np.ones((5, 4, 4)))[2:], u2[1] * np.ones(
(3, 4, 4)))
pass
def check(indices, plates, shape, axis=-1, use_mask=False):
mu = np.random.rand(*(plates+shape))
alpha = np.random.rand(*(plates+shape))
X = GaussianARD(mu, alpha, shape=shape, plates=plates)
Y = Take(X, indices, plate_axis=axis)
Z = GaussianARD(Y, 1, shape=shape)
z = np.random.randn(*(Z.get_shape(0)))
if use_mask:
mask = np.mod(np.reshape(np.arange(np.prod(Z.plates)), Z.plates), 2) != 0
else:
mask = True
Z.observe(z, mask=mask)
X.update()
(x0, x1) = X.get_moments()
# For comparison, build the same model brute force
X = GaussianARD(mu, alpha, shape=shape, plates=plates)
# Number of trailing plate axes before the take axis
N = len(X.plates) + axis
# Reshape the take axes into a single axis
z_shape = X.plates[:axis] + (-1,)
if axis < -1:
z_shape = z_shape + X.plates[(axis+1):]
z_shape = z_shape + shape
z = np.reshape(z, z_shape)
# Reshape the take axes into a single axis
if use_mask:
mask_shape = X.plates[:axis] + (-1,)
if axis < -1:
mask_shape = mask_shape + X.plates[(axis+1):]
mask = np.reshape(mask, mask_shape)
for (j, i) in enumerate(range(np.size(indices))):
ind = np.array(indices).flatten()[i]
index_x = N*(slice(None),) + (ind,)
index_z = N*(slice(None),) + (j,)
# print(index)
Xi = X[index_x]
zi = z[index_z]
Zi = GaussianARD(Xi, 1, ndim=len(shape))
if use_mask:
maski = mask[index_z]
else:
maski = True
Zi.observe(zi, mask=maski)
X.update()
self.assertAllClose(
x0,
X.get_moments()[0],
)
self.assertAllClose(
x1,
X.get_moments()[1],
)
return
def test_message_to_child(self):
"""
Test the message to child of Concatenate node.
"""
# Two parents without shapes
X1 = GaussianARD(0, 1, plates=(2,), shape=())
X2 = GaussianARD(0, 1, plates=(3,), shape=())
Y = Concatenate(X1, X2)
u1 = X1.get_moments()
u2 = X2.get_moments()
u = Y.get_moments()
self.assertAllClose((u[0]*np.ones((5,)))[:2],
u1[0]*np.ones((2,)))
self.assertAllClose((u[1]*np.ones((5,)))[:2],
u1[1]*np.ones((2,)))
self.assertAllClose((u[0]*np.ones((5,)))[2:],
u2[0]*np.ones((3,)))
self.assertAllClose((u[1]*np.ones((5,)))[2:],
u2[1]*np.ones((3,)))
# Two parents with shapes
X1 = GaussianARD(0, 1, plates=(2,), shape=(4,))
X2 = GaussianARD(0, 1, plates=(3,), shape=(4,))
Y = Concatenate(X1, X2)
u1 = X1.get_moments()
u2 = X2.get_moments()
u = Y.get_moments()
self.assertAllClose((u[0]*np.ones((5,4)))[:2],
u1[0]*np.ones((2,4)))
self.assertAllClose((u[1]*np.ones((5,4,4)))[:2],
u1[1]*np.ones((2,4,4)))
self.assertAllClose((u[0]*np.ones((5,4)))[2:],
u2[0]*np.ones((3,4)))
self.assertAllClose((u[1]*np.ones((5,4,4)))[2:],
u2[1]*np.ones((3,4,4)))
# Test with non-constant axis
X1 = GaussianARD(0, 1, plates=(2,4), shape=())
X2 = GaussianARD(0, 1, plates=(3,4), shape=())
Y = Concatenate(X1, X2, axis=-2)
u1 = X1.get_moments()
u2 = X2.get_moments()
u = Y.get_moments()
self.assertAllClose((u[0]*np.ones((5,4)))[:2],
u1[0]*np.ones((2,4)))
self.assertAllClose((u[1]*np.ones((5,4)))[:2],
u1[1]*np.ones((2,4)))
self.assertAllClose((u[0]*np.ones((5,4)))[2:],
u2[0]*np.ones((3,4)))
self.assertAllClose((u[1]*np.ones((5,4)))[2:],
u2[1]*np.ones((3,4)))
# Test with constant parent
X1 = np.random.randn(2, 4)
X2 = GaussianARD(0, 1, plates=(3,), shape=(4,))
Y = Concatenate(X1, X2)
u1 = Y.parents[0].get_moments()
u2 = X2.get_moments()
u = Y.get_moments()
self.assertAllClose((u[0]*np.ones((5,4)))[:2],
u1[0]*np.ones((2,4)))
self.assertAllClose((u[1]*np.ones((5,4,4)))[:2],
u1[1]*np.ones((2,4,4)))
self.assertAllClose((u[0]*np.ones((5,4)))[2:],
u2[0]*np.ones((3,4)))
self.assertAllClose((u[1]*np.ones((5,4,4)))[2:],
u2[1]*np.ones((3,4,4)))
pass
class BayesianRegression(object):
"""Bayesian linear regression."""
def __init__(self,
n_feature,
prior_mean=0,
prior_precision=1e-6,
prior_a=10,
prior_b=1):
super().__init__()
self.n_feature = n_feature
if np.shape(prior_mean) != (n_feature, ):
prior_mean = prior_mean * np.ones(n_feature)
if np.shape(prior_precision) != (n_feature, n_feature):
prior_precision = prior_precision * np.ones(n_feature)
self.prior_mean = prior_mean
self.prior_precision = prior_precision
self.prior_a = prior_a
self.prior_b = prior_b
self._init_weights()
# print("Intialize regression")
# self.print()
def _init_weights(self):
self.weights = GaussianARD(self.prior_mean,
self.prior_precision,
shape=(self.n_feature, ))
def fit(self, X, y):
self._init_weights()
# self.cost,
# self.myopic_voc(action, state),
# self.vpi_action(action, state),
# self.vpi(state),
# self.expected_term_reward(state)
self.tau = Gamma(self.prior_a, self.prior_b)
F = SumMultiply('i,i', self.weights, X)
y_obs = GaussianARD(F, self.tau)
y_obs.observe(y)
Q = VB(y_obs, self.weights)
Q.update(repeat=10, tol=1e-4, verbose=False)
def predict(self, x, return_var=False):
y = SumMultiply('i,i', self.weights, x)
y_hat, var, *_ = y.get_moments()
if return_var:
return y_hat, var
else:
return y_hat
def sample(self, x):
w = self.weights.random()
return x @ w
def print(self, diagonal=True):
mean, m2 = self.weights.get_moments()[:2]
var = m2 - mean**2
if diagonal:
var = np.diagonal(var)
bar = '_' * 40
print(f'{bar}\n{mean.round(3)}\n{var.round(3)}\n{bar}')
