def sample(self, query_points: TensorType, num_samples: int) -> TensorType: # Taken from GPflow implementation of `GPModel.predict_f_samples` in gpflow.models.model mean, cov = self.model_gpflux.predict_f(query_points, full_cov=True) mean_for_sample = tf.linalg.adjoint(mean) samples = sample_mvn(mean_for_sample, cov, True, num_samples=num_samples) samples = tf.linalg.adjoint(samples) return samples
def draw_conditional_sample(mean: TensorType, cov: TensorType, f_old: TensorType) -> tf.Tensor: r""" Draw a sample :math:`\tilde{f}_\text{new}` from the conditional multivariate Gaussian :math:`p(f_\text{new} | f_\text{old})`, where the parameters ``mean`` and ``cov`` are the mean and covariance matrix of the joint multivariate Gaussian over :math:`[f_\text{old}, f_\text{new}]`. :param mean: A tensor with the shape ``[..., D, N+M]`` with the mean of ``[f_old, f_new]``. For each ``[..., D]`` this is a stacked vector of the form: .. math:: \begin{pmatrix} \operatorname{mean}(f_\text{old}) \;[N] \\ \operatorname{mean}(f_\text{new}) \;[M] \end{pmatrix} :param cov: A tensor with the shape ``[..., D, N+M, N+M]`` with the covariance of ``[f_old, f_new]``. For each ``[..., D]``, there is a 2x2 block matrix of the form: .. math:: \begin{pmatrix} \operatorname{cov}(f_\text{old}, f_\text{old}) \;[N, N] & \operatorname{cov}(f_\text{old}, f_\text{new}) \;[N, M] \\ \operatorname{cov}(f_\text{new}, f_\text{old}) \;[M, N] & \operatorname{cov}(f_\text{new}, f_\text{new}) \;[M, M] \end{pmatrix} :param f_old: A tensor of observations with the shape ``[..., D, N]``, drawn from Normal distribution with mean :math:`\operatorname{mean}(f_\text{old}) \;[N]`, and covariance :math:`\operatorname{cov}(f_\text{old}, f_\text{old}) \;[N, N]` :return: A sample :math:`\tilde{f}_\text{new}` from the conditional normal :math:`p(f_\text{new} | f_\text{old})` with the shape ``[..., D, M]``. """ N, D = tf.shape(f_old)[-1], tf.shape(f_old)[-2] # noqa: F841 M = tf.shape(mean)[-1] - N cov_old = cov[..., :N, :N] # [..., D, N, N] cov_new = cov[..., -M:, -M:] # [..., D, M, M] cov_cross = cov[..., :N, -M:] # [..., D, N, M] L_old = _cholesky_with_jitter(cov_old) # [..., D, N, N] A = tf.linalg.triangular_solve(L_old, cov_cross, lower=True) # [..., D, N, M] var_new = cov_new - tf.matmul(A, A, transpose_a=True) # [..., D, M, M] mean_new = mean[..., -M:] # [..., D, M] mean_old = mean[..., :N] # [..., D, N] mean_old_diff = (f_old - mean_old)[..., None] # [..., D, N, 1] AM = tf.linalg.triangular_solve(L_old, mean_old_diff) # [..., D, N, 1] mean_new = mean_new + (tf.matmul(A, AM, transpose_a=True)[..., 0] ) # [..., D, M] return sample_mvn(mean_new, var_new, full_cov=True)
def test_sample_mvn(full_cov): """ Draws 10,000 samples from a distribution with known mean and covariance. The test checks if the mean and covariance of the samples is close to the true mean and covariance. """ N, D = 10000, 2 means = tf.ones((N, D), dtype=float_type) if full_cov: covs = tf.eye(D, batch_shape=[N], dtype=float_type) else: covs = tf.ones((N, D), dtype=float_type) samples = sample_mvn(means, covs, full_cov) samples_mean = np.mean(samples, axis=0) samples_cov = np.cov(samples, rowvar=False) np.testing.assert_array_almost_equal(samples_mean, [1.0, 1.0], decimal=1) np.testing.assert_array_almost_equal(samples_cov, [[1.0, 0.0], [0.0, 1.0]], decimal=1)
def test_sample_mvn_shapes(leading_dims, n, d, num_samples, full_cov): mean_shape = leading_dims + (n, d) means = tf.zeros(mean_shape, dtype=default_float()) if full_cov: covariance_shape = leading_dims + (n, d, d) sqrt_cov = tf.random.normal(covariance_shape, dtype=default_float()) covariances = tf.matmul(sqrt_cov, sqrt_cov, transpose_b=True) else: covariance_shape = leading_dims + (n, d) covariances = tf.random.normal(covariance_shape, dtype=default_float()) samples = sample_mvn(means, covariances, full_cov, num_samples) if num_samples: expected_shape = leading_dims + (num_samples, n, d) else: expected_shape = leading_dims + (n, d) assert_equal(samples.shape, expected_shape)
def all_layer_mean_var_samples( gp_layers: Sequence[GPLayer], X: TensorType ) -> Tuple[List[np.ndarray], List[np.ndarray], List[np.ndarray]]: S = 5 sample = X means, covs, samples = [], [], [] for layer in gp_layers: mean, cov = layer.predict(sample, full_output_cov=False, full_cov=True) # [N, D], [D, N, N] all_samples = sample_mvn(tf.linalg.adjoint(mean), cov, full_cov=True, num_samples=S) all_samples = tf.linalg.adjoint(all_samples) sample = all_samples[0] means.append(mean.numpy()) covs.append(cov.numpy()) samples.append(all_samples.numpy()) return means, covs, samples
def test_sample_mvn(cov_structure): """ Draws 10,000 samples from a distribution with known mean and covariance. The test checks if the mean and covariance of the samples is close to the true mean and covariance. """ N, D = 10000, 2 means = tf.ones((N, D), dtype=float_type) if cov_structure == "full": covs = tf.eye(D, batch_shape=[N], dtype=float_type) elif cov_structure == "diag": covs = tf.ones((N, D), dtype=float_type) else: raise (NotImplementedError) samples = sample_mvn(means, covs, cov_structure) samples_mean = np.mean(samples, axis=0) samples_cov = np.cov(samples, rowvar=False) np.testing.assert_array_almost_equal(samples_mean, [1., 1.], decimal=1) np.testing.assert_array_almost_equal(samples_cov, [[1., 0.], [0., 1.]], decimal=1)