def draw_samples(self,
num_samples,
means=None,
cov=None,
base_rvs=None,
parallelism=None,
chol=None,
model=None,
name=None,
dtype=None,
**kwargs):
if (model is None): model = self.model
with tf.name_scope('draw_samples') as scope:
# Compute Cholesky factor LL^{T} = K
if (chol is None):
jitter = getattr(model, 'jitter', None)
chol = utils.jitter_cholesky(cov, jitter=jitter)
elif (dtype is not None):
chol = tf.cast(chol, dtype)
# Generate i.i.d. standard normal random variables
if (base_rvs is None):
if (parallelism is None): parallelism = chol.get_shape()[-1]
rvs_shape = tf.TensorShape((num_samples, parallelism))
base_rvs = tf.random_normal(rvs_shape, dtype=chol.dtype)
elif (dtype is not None):
base_rvs = tf.cast(base_rvs, dtype)
# Produce MVN samples
samples = tf.tensordot(chol, base_rvs, [[-1], [-1]])
if (means is not None):
samples += tf.cast(means, samples.dtype)
return utils.swap_axes(samples, -1, -2, name=name)
def outer_mapper(mapped_args): #iterate over pools
# Precompute precision matrix
pool, pool_samples = mapped_args
inputs_old_ext = append_inputs(pool)
K_ext = self.model.covar_fn(inputs_old_ext, inputs_old_ext)
chol_ext = utils.jitter_cholesky(K_ext + I_nz,
jitter=self.model.jitter)
def inner_mapper(
pool_sample): #iterate over fantasized pool outcomes
outputs_old_ext = append_outputs(pool_sample[..., None])
return self.model.predict\
(
discretization,
inputs_old_ext,
outputs_old_ext,
chol=chol_ext,
full_cov=True,
)
means, covs = tf.map_fn\
(
inner_mapper,
pool_samples,
(tf.float64, tf.float64),
)
return means, covs[:1] #identical slices along first axis
def compute_cholesky(self, inputs, cov=None, noisy=False, **kwargs):
with tf.name_scope('compute_cholesky') as scope:
if (cov is None):
cov = self.covar_fn(inputs, inputs, noisy=noisy, **kwargs)
elif noisy:
cov = self.noise_fn(cov, **kwargs)
chol = utils.jitter_cholesky(cov, jitter=self.jitter)
return chol
def mapped_fn(*sharded_args, kwargs=kwargs):
pools = sharded_args[0]
if use_rand_features:
theta_mu = tf.transpose(self.mapping.mu)
means = self.mapping(pools, **{**rf_kwargs, 'theta': theta_mu})
samplesT = self.mapping(pools, **rf_kwargs)
samples = utils.swap_axes(samplesT, -1, -2)
return fn(means,
None,
samples=samples,
pools=pools,
inputs_old=inputs_old,
outputs_old=outputs_old,
**kwargs)
if (posteriors is not None):
# Reuse precomputed posteriors
means, cov, chol = sharded_args[-1]
sharded_args = sharded_args[:-1]
else:
# Compute marginal/joint posterior individual pools
means, cov = model.get_or_create_node\
(
group='predict',
fn=model.predict,
args=sharded_args + constant_args,
kwargs=predict_kwargs,
stateful=True,
)
chol = utils.jitter_cholesky(cov, jitter=model.jitter)
losses = fn(means,
cov,
pools=pools,
inputs_old=inputs_old,
outputs_old=outputs_old,
chol=chol,
**kwargs)
return losses, (means, cov, chol)
def _sample_theta(num_samples, mu, cov=None, eigen_decomp=None, \
Z=None, base_rvs=None, noise=None, jitter=None, dtype=None):
'''
Subroutine for actually drawing samples from a belief over
the parameters $\theta$ used to combine random features.
'''
if (dtype is None): dtype = mu.dtype
# Draw i.i.d. standard normal r.v.'s
if (base_rvs is None):
rvs_shape = (num_samples, tf.shape(mu)[0])
base_rvs = tf.random_normal(rvs_shape, dtype=dtype)
# Compute sample residuals for \theta
if (eigen_decomp is not None and Z is not None):
# Eigendecomposition terms
sigma2 = tf.cast(jitter if (noise is None) else noise + jitter, dtype)
D, U = eigen_decomp[0], eigen_decomp[1]
test_eigenvalues = tf.assert_non_negative(D, name='test_eigenvalues')
with tf.control_dependencies([test_eigenvalues]):
R = tf.reciprocal(D + tf.sqrt(D) * tf.sqrt(sigma2))
# See appendix B.2 of [Seeger, 2008]
# 'Bayesian Inference and Optimal Design for the Sparse Linear Model'
q = base_rvs #here for notational clarity
Zq = tf.matmul(Z, q, transpose_b=True)
RUZq = tf.matmul(R * U, Zq, transpose_a=True)
resid = q - tf.matmul(tf.matmul(U, RUZq), Z, transpose_a=True)
else:
assert cov is not None, 'Either a covariance matrix or eigendecomposition'\
' and features Z must be provided'
chol = utils.jitter_cholesky(cov)
resid = tf.matmul(base_rvs, chol, transpose_b=True)
# Add in mean and rescale
samples = tf.add(tf.squeeze(mu), resid)
return samples
log_lenscales=np.log(0.1),
log_amplitude=np.log(1.0),
log_noise=np.log(1e-3),
mean=100,
dtype='float64',
)
predict_op = gp.predict(X1_ref, X0_ref, Y0_ref)
gp_init_op = tf.variables_initializer(list(gp.state.values()))
# Create TensorFlow session, initialize GP variables,
sess = tf.InteractiveSession()
sess.run(gp_init_op)
# Sample Y0 from the true prior
K_xx = gp.covar_fn(X0_ref, X0_ref)
L_xx = utils.jitter_cholesky(K_xx, jitter=gp.jitter)
rvs = tf.random_normal([tf.shape(L_xx)[0], 1], dtype=L_xx.dtype)
sample_y0 = gp.mean + tf.matmul(L_xx, rvs)
Y0_src = sess.run(sample_y0, {X0_ref: X0_src})
feed_dict[Y0_ref] = Y0_src
# Calculate closed-form GP posterior
mu_exact, var_exact = sess.run(predict_op, feed_dict)
# Generate random feature-based sample paths
paths = _build_mappings(sess,
gp,
X0_src,
Y0_src,
num_paths=1024,
num_features=1024)
def _monte_carlo(self,
means_pool,
cov_pool,
xov_pd,
means_disc,
cov_disc,
chol_pool=None,
future_rvs=None,
fantasy_rvs=None,
parallelism=None,
num_futures=None,
num_fantasies=None,
**kwargs):
'''
Generic Monte Carlo subroutine for 1-step lookahead integrands.
Computes updated discretization posteriors conditioned on each
sample future.
'''
if (num_futures is None): num_futures = self.num_futures
if (num_fantasies is None): num_fantasies = self.num_fantasies
with tf.name_scope('monte_carlo'.format(self.name)) as scope:
if (chol_pool is None): chol_pool = utils.jitter_cholesky(cov_pool)
# Fantasize residuals (Y - \mu) at pool locations
resid_pool = self.draw_samples\
(
num_futures,
chol=chol_pool,
base_rvs=future_rvs,
parallelism=parallelism,
)
# Sample pool outcomes to condition upon
futures = resid_pool + utils.swap_axes(means_pool, -1, -2)
# Solve for beta := K(disc, pool) K_nz(pool, pool)^{-1}
beta = utils.swap_axes(tf.cholesky_solve(chol_pool, xov_pd), -1,
-2)
# Compute changes in discretization posterior
d_cov = -tf.einsum('ijk,ikl->ijl', beta, xov_pd)
ds_means = tf.einsum('ijk,ilk->ilj', beta, resid_pool)
# Posteriors conditional on each sampled future
means_future = means_disc + tf.expand_dims(ds_means, -1)
covs_future = cov_disc[None, None] + d_cov[:, None]
if self.run_unit_test:
means_future, covs_future = self.test_conditional_posterior\
(
means_future,
covs_future,
futures
)
# Monte Carlo integration subroutine
estimate = self.integrand\
(
means_future,
covs_future,
num_samples=num_fantasies,
base_rvs=fantasy_rvs,
futures=futures,
means_pool=means_pool,
means_disc=means_disc,
cov_disc=cov_disc,
**kwargs
)
return estimate