def main(argv):
del argv # Unused.
n_examples = 10
a = 1.3
b = 2.2
kappa = 1.5
mu0 = 0.3
tau = np.random.gamma(a, 1. / b)
mu = np.random.normal(mu0, 1. / np.sqrt(tau * kappa))
x = np.random.normal(mu, 1. / np.sqrt(tau), n_examples)
all_args = [a, b, kappa, mu0, tau, mu, x]
all_args_ex_mu = [a, b, kappa, mu0, tau, x]
log_joint = ph.make_log_joint_fn(model)
mu_conditional_factory = conjugacy.complete_conditional(
log_joint, 5, conjugacy.SupportTypes.REAL, *all_args)
mu_conditional = mu_conditional_factory(*all_args_ex_mu)
log_p_tau = conjugacy.marginalize(log_joint, 5, conjugacy.SupportTypes.REAL,
*all_args)
tau_conditional_factory = conjugacy.complete_conditional(
log_p_tau, 4, conjugacy.SupportTypes.NONNEGATIVE, *all_args_ex_mu)
tau_conditional = tau_conditional_factory(*[a, b, kappa, mu0, x])
print('True tau: {}'.format(tau))
print('tau posterior is gamma({}, {}). Mean is {}, std. dev. is {}.'.format(
tau_conditional.args[0], 1. / tau_conditional.args[2],
tau_conditional.args[0] * tau_conditional.args[2],
np.sqrt(tau_conditional.args[0]) * tau_conditional.args[2]))
print()
print('True mu: {}'.format(mu))
print('mu posterior given tau is normal({}, {})'.format(
mu_conditional.args[0], mu_conditional.args[1]))
def testFactorAnalysis(self):
def log_joint(x, w, epsilon, tau, alpha, beta):
log_p_epsilon = log_probs.norm_gen_log_prob(epsilon, 0, 1)
log_p_w = log_probs.norm_gen_log_prob(w, 0, 1)
log_p_tau = log_probs.gamma_gen_log_prob(tau, alpha, beta)
# TODO(mhoffman): The transposed version below should work.
# log_p_x = log_probs.norm_gen_log_prob(x, np.dot(epsilon, w), 1. / np.sqrt(tau))
log_p_x = log_probs.norm_gen_log_prob(x, np.einsum('ik,jk->ij', epsilon, w),
1. / np.sqrt(tau))
return log_p_epsilon + log_p_w + log_p_tau + log_p_x
n_examples = 20
D = 10
K = 5
alpha = 2.
beta = 8.
tau = np.random.gamma(alpha, beta)
w = np.random.normal(loc=0, scale=1, size=[D, K])
epsilon = np.random.normal(loc=0, scale=1, size=[n_examples, K])
x = np.random.normal(loc=epsilon.dot(w.T), scale=np.sqrt(tau))
all_args = [x, w, epsilon, tau, alpha, beta]
w_conditional_factory = complete_conditional(log_joint, 1,
SupportTypes.REAL, *all_args)
conditional = w_conditional_factory(x, epsilon, tau, alpha, beta)
true_cov = np.linalg.inv(tau * np.einsum('nk,nl->kl', epsilon, epsilon) +
np.eye(K))
true_mean = tau * np.einsum('nk,nd,kl->dl', epsilon, x, true_cov)
for d in range(D):
self.assertTrue(np.allclose(conditional[d].cov, true_cov))
self.assertTrue(np.allclose(conditional[d].mean, true_mean[d]))
epsilon_conditional_factory = complete_conditional(log_joint, 2,
SupportTypes.REAL,
*all_args)
conditional = epsilon_conditional_factory(x, w, tau, alpha, beta)
true_cov = np.linalg.inv(tau * np.einsum('dk,dl->kl', w, w) + np.eye(K))
true_mean = tau * np.einsum('dk,nd,kl->nl', w, x, true_cov)
for n in range(n_examples):
self.assertTrue(np.allclose(conditional[n].cov, true_cov))
self.assertTrue(np.allclose(conditional[n].mean, true_mean[n]))
tau_conditional_factory = complete_conditional(log_joint, 3,
SupportTypes.NONNEGATIVE,
*all_args)
conditional = tau_conditional_factory(x, w, epsilon, alpha, beta)
true_a = alpha + 0.5 * n_examples * D
true_b = beta + 0.5 * np.sum(np.square(x - epsilon.dot(w.T)))
self.assertAlmostEqual(true_a, conditional.args[0])
self.assertAlmostEqual(true_b, 1. / conditional.args[2])
def _condition_and_marginalize(log_joint, argnum, support, *args): sub_args = args[:argnum] + args[argnum + 1:] marginalized = marginalize(log_joint, argnum, support, *args) marginalized_value = marginalized(*sub_args) conditional_factory = complete_conditional(log_joint, argnum, support, *args) conditional = conditional_factory(*sub_args) return conditional, marginalized_value
def main(argv):
del argv
n_clusters = 5
n_dimensions = 2
n_observations = 200
alpha = 3.3 * np.ones(n_clusters)
sigma_sq_mu = 1.5**2
sigma_sq = 0.5**2
np.random.seed(10001)
pi = np.random.gamma(alpha)
pi /= pi.sum()
mu = np.random.normal(0, np.sqrt(sigma_sq_mu), [n_clusters, n_dimensions])
z = np.random.choice(np.arange(n_clusters), size=n_observations, p=pi)
x = np.random.normal(mu[z, :], sigma_sq)
pi_est = np.ones(n_clusters) / n_clusters
z_est = np.random.choice(np.arange(n_clusters),
size=n_observations,
p=pi_est)
mu_est = np.random.normal(0., 0.01, [n_clusters, n_dimensions])
all_args = [sigma_sq, alpha, sigma_sq_mu, pi_est, mu_est, z_est, x]
log_joint = ph.make_log_joint_fn(model)
pi_posterior = conjugacy.complete_conditional(
log_joint, 3, conjugacy.SupportTypes.SIMPLEX, *all_args)
z_posterior = conjugacy.complete_conditional(
log_joint, 5, conjugacy.SupportTypes.INTEGER, *all_args)
mu_posterior = conjugacy.complete_conditional(log_joint, 4,
conjugacy.SupportTypes.REAL,
*all_args)
print('iteration\tlog_joint')
for iteration in range(100):
z_est[:] = z_posterior(*remove_arg(5, all_args)).rvs()
pi_est[:] = pi_posterior(*remove_arg(3, all_args)).rvs()
mu_est[:] = mu_posterior(*remove_arg(4, all_args)).rvs()
print('{}\t\t{}'.format(iteration, log_joint(*all_args)))
def run_gibbs(log_joint_fn, all_args, num_iterations):
"""Train model with Gibbs sampling."""
alpha, beta, epsilon, w, tau, x = all_args
# Form complete conditionals for Gibbs sampling.
epsilon_conditional_factory = conjugacy.complete_conditional(
log_joint_fn, 2, conjugacy.SupportTypes.REAL, *all_args)
w_conditional_factory = conjugacy.complete_conditional(
log_joint_fn, 3, conjugacy.SupportTypes.REAL, *all_args)
tau_conditional_factory = conjugacy.complete_conditional(
log_joint_fn, 4, conjugacy.SupportTypes.NONNEGATIVE, *all_args)
epsilon_conditional = lambda w, tau: epsilon_conditional_factory( # pylint: disable=g-long-lambda
alpha, beta, w, tau, x)
w_conditional = lambda epsilon, tau: w_conditional_factory( # pylint: disable=g-long-lambda
alpha, beta, epsilon, tau, x)
tau_conditional = lambda epsilon, w: tau_conditional_factory( # pylint: disable=g-long-lambda
alpha, beta, epsilon, w, x)
log_posterior = lambda epsilon, w, tau: log_joint_fn( # pylint: disable=g-long-lambda
alpha, beta, epsilon, w, tau, x)
# Run training loop. Track expected log joint probability, i.e.,
# E [ log p(xnew, params | xtrain) ]. It is estimated with 1 posterior sample.
print('Running Gibbs...')
epsilon = ph.norm.rvs(0, 1, size=epsilon.shape)
w = ph.norm.rvs(0, 1, size=w.shape)
tau = ph.gamma.rvs(alpha, scale=1. / beta)
log_joints = []
runtimes = []
start = time.time()
for t in range(num_iterations):
epsilon = epsilon_conditional(w, tau).rvs()
w = w_conditional(epsilon, tau).rvs()
tau = tau_conditional(epsilon, w).rvs()
if t % FLAGS.num_print == 0 or (t + 1) == num_iterations:
log_joint = log_posterior(epsilon, w, tau)
runtime = time.time() - start
print('Iteration: {:>3d} Log Joint: {:.3f} '
'Runtime (s): {:.3f}'.format(t, log_joint, runtime))
log_joints.append(log_joint)
runtimes.append(runtime)
return log_joints, runtimes
def make_marginal_fn():
x1_given_y1_factory = complete_conditional(
log_p_x1_y1, 0, SupportTypes.REAL, *([1.] * 4))
log_p_y1 = marginalize(log_p_x1_y1, 0, SupportTypes.REAL, *([1.] * 4))
log_p_xtt_ytt = marginalize(
log_p_xt_xtt_ytt, 0, SupportTypes.REAL, *([1.] * 7))
log_p_ytt = marginalize(
log_p_xtt_ytt, 0, SupportTypes.REAL, *([1.] * 6))
xt_conditional_factory = complete_conditional(
log_p_xtt_ytt, 0, SupportTypes.REAL, *([1.] * 6))
def marginal(y_list, x_scale, y_scale):
log_p_y = log_p_y1(y_list[0], x_scale, y_scale)
xt_conditional = x1_given_y1_factory(y_list[0], x_scale, y_scale)
for t in range(1, len(y_list)):
log_p_y += log_p_ytt(y_list[t], xt_conditional.args[0],
xt_conditional.args[1], x_scale, y_scale)
xt_conditional = xt_conditional_factory(
y_list[t], xt_conditional.args[0], xt_conditional.args[1], x_scale,
y_scale)
return log_p_y
return marginal
def testLinearRegression(self):
def log_joint(X, beta, y):
predictions = np.einsum('ij,j->i', X, beta)
errors = y - predictions
log_prior = np.einsum('i,i,i->', -0.5 * np.ones_like(beta), beta, beta)
log_likelihood = np.einsum(',k,k->', -0.5, errors, errors)
return log_prior + log_likelihood
n_examples = 10
n_predictors = 2
X = np.random.randn(n_examples, n_predictors)
beta = np.random.randn(n_predictors)
y = np.random.randn(n_examples)
graph = make_expr(log_joint, X, beta, y)
graph = canonicalize(graph)
args = graph.free_vars.keys()
sufficient_statistic_nodes = find_sufficient_statistic_nodes(graph, args[1])
sufficient_statistics = [eval_node(node, graph.free_vars,
{'X': X, 'beta': beta, 'y': y})
for node in sufficient_statistic_nodes]
correct_sufficient_statistics = [
-0.5 * beta.dot(beta), beta,
-0.5 * np.einsum('ij,ik,j,k', X, X, beta, beta)
]
self.assertTrue(_match_values(sufficient_statistics,
correct_sufficient_statistics))
_, natural_parameter_funs = _extract_conditional_factors(graph, 'beta')
self.assertTrue(_match_values(natural_parameter_funs.keys(),
['x', 'einsum(...a,...b->...ab, x, x)',
'einsum(...,...->..., x, x)'],
lambda x, y: x == y))
natural_parameter_vals = [f(X, beta, y) for f in
natural_parameter_funs.values()]
correct_parameter_vals = [-0.5 * np.ones(n_predictors), -0.5 * X.T.dot(X),
y.dot(X)]
self.assertTrue(_match_values(natural_parameter_vals,
correct_parameter_vals))
conditional_factory = complete_conditional(log_joint, 1, SupportTypes.REAL,
X, beta, y)
conditional = conditional_factory(X, y)
true_cov = np.linalg.inv(X.T.dot(X) + np.eye(n_predictors))
true_mean = true_cov.dot(y.dot(X))
self.assertTrue(np.allclose(true_cov, conditional.cov))
self.assertTrue(np.allclose(true_mean, conditional.mean))
def testLinearRegression(self):
def log_joint(X, beta, y):
predictions = np.einsum('ij,j->i', X, beta)
errors = y - predictions
log_prior = np.einsum('i,i,i->', -0.5 * np.ones_like(beta), beta, beta)
log_likelihood = np.einsum(',k,k->', -0.5, errors, errors)
return log_prior + log_likelihood
n_examples = 10
n_predictors = 2
X = np.random.randn(n_examples, n_predictors)
beta = np.random.randn(n_predictors)
y = np.random.randn(n_examples)
graph = make_expr(log_joint, X, beta, y)
graph = canonicalize(graph)
args = graph.free_vars.keys()
sufficient_statistic_nodes = find_sufficient_statistic_nodes(graph, args[1])
sufficient_statistics = [eval_node(node, graph.free_vars,
{'X': X, 'beta': beta, 'y': y})
for node in sufficient_statistic_nodes]
correct_sufficient_statistics = [
-0.5 * beta.dot(beta), beta,
-0.5 * np.einsum('ij,ik,j,k', X, X, beta, beta)
]
self.assertTrue(_match_values(sufficient_statistics,
correct_sufficient_statistics))
new_log_joint, _, stats_funs, _ = (
statistic_representation(log_joint, (X, beta, y),
(SupportTypes.REAL,), (1,)))
beta_stat_fun = stats_funs[0]
beta_natparam = grad_namedtuple(new_log_joint, 1)(X, beta_stat_fun(beta), y)
correct_beta_natparam = (-0.5 * X.T.dot(X), y.dot(X),
-0.5 * np.ones(n_predictors))
self.assertTrue(_match_values(beta_natparam, correct_beta_natparam))
conditional_factory = complete_conditional(log_joint, 1, SupportTypes.REAL,
X, beta, y)
conditional = conditional_factory(X, y)
true_cov = np.linalg.inv(X.T.dot(X) + np.eye(n_predictors))
true_mean = true_cov.dot(y.dot(X))
self.assertTrue(np.allclose(true_cov, conditional.cov))
self.assertTrue(np.allclose(true_mean, conditional.mean))
def testTwoGaussians(self):
def log_joint(x1, x2):
log_p_x1 = -0.5 * x1 * x1
x_diff = x2 - x1
log_p_x2 = -0.5 * x_diff * x_diff
return log_p_x1 + log_p_x2
x1 = np.random.randn()
x2 = x1 + np.random.randn()
all_args = [x1, x2]
marginal_p_x2 = marginalize(log_joint, 0, SupportTypes.REAL, *all_args)
correct_marginalized_value = (
-0.25 * x2 * x2 - 0.5 * np.log(2.) + 0.5 * np.log(2. * np.pi))
self.assertAlmostEqual(correct_marginalized_value, marginal_p_x2(x2))
x2_conditional = complete_conditional(marginal_p_x2, 0, SupportTypes.REAL,
x2)()
self.assertAlmostEqual(x2_conditional.args[0], 0.)
self.assertAlmostEqual(x2_conditional.args[1] ** 2, 2.)
def testMixtureOfGaussians(self):
def log_joint(x, pi, z, mu, sigma_sq, alpha, sigma_sq_mu):
log_p_pi = log_probs.dirichlet_gen_log_prob(pi, alpha)
log_p_mu = log_probs.norm_gen_log_prob(mu, 0, np.sqrt(sigma_sq_mu))
z_one_hot = one_hot(z, len(pi))
log_p_z = np.einsum('ij,j->', z_one_hot, np.log(pi))
mu_z = np.einsum('ij,jk->ik', z_one_hot, mu)
log_p_x = log_probs.norm_gen_log_prob(x, mu_z, np.sqrt(sigma_sq))
return log_p_pi + log_p_z + log_p_mu + log_p_x
n_clusters = 5
n_dimensions = 2
n_observations = 200
alpha = 3.3 * np.ones(n_clusters)
sigma_sq_mu = 1.5 ** 2
sigma_sq = 0.5 ** 2
np.random.seed(10001)
pi = np.random.gamma(alpha)
pi /= pi.sum()
mu = np.random.normal(0, np.sqrt(sigma_sq_mu), [n_clusters, n_dimensions])
z = np.random.choice(np.arange(n_clusters), size=n_observations, p=pi)
x = np.random.normal(mu[z, :], sigma_sq)
pi_est = np.ones(n_clusters) / n_clusters
z_est = np.random.choice(np.arange(n_clusters), size=n_observations,
p=pi_est)
mu_est = np.random.normal(0., 0.01, [n_clusters, n_dimensions])
all_args = [x, pi_est, z_est, mu_est, sigma_sq, alpha, sigma_sq_mu]
pi_posterior_args = all_args[:1] + all_args[2:]
z_posterior_args = all_args[:2] + all_args[3:]
mu_posterior_args = all_args[:3] + all_args[4:]
pi_posterior = complete_conditional(log_joint, 1, SupportTypes.SIMPLEX,
*all_args)
z_posterior = complete_conditional(log_joint, 2, SupportTypes.INTEGER,
*all_args)
mu_posterior = complete_conditional(log_joint, 3, SupportTypes.REAL,
*all_args)
self.assertTrue(np.allclose(
pi_posterior(*pi_posterior_args).alpha,
alpha + np.histogram(z_est, np.arange(n_clusters+1))[0]))
correct_z_logits = -0.5 / sigma_sq * np.square(x[:, :, None] -
mu_est.T[None, :, :]).sum(1)
correct_z_logits += np.log(pi_est)
correct_z_posterior = np.exp(correct_z_logits -
misc.logsumexp(correct_z_logits, 1,
keepdims=True))
self.assertTrue(np.allclose(correct_z_posterior,
z_posterior(*z_posterior_args).p))
correct_mu_posterior_mean = np.zeros_like(mu_est)
correct_mu_posterior_var = np.zeros_like(mu_est)
for k in range(n_clusters):
n_k = (z_est == k).sum()
correct_mu_posterior_var[k] = 1. / (1. / sigma_sq_mu + n_k / sigma_sq)
correct_mu_posterior_mean[k] = (
x[z_est == k].sum(0) / sigma_sq * correct_mu_posterior_var[k])
mu_posterior_val = mu_posterior(*mu_posterior_args)
self.assertTrue(np.allclose(correct_mu_posterior_mean,
mu_posterior_val.args[0]))
self.assertTrue(np.allclose(correct_mu_posterior_var,
mu_posterior_val.args[1] ** 2))
