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)