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