def __init__(
self,
train_X: Tensor,
train_Y: Tensor,
train_Yvar: Tensor,
outcome_transform: Optional[OutcomeTransform] = None,
) -> None:
r"""A single-task exact GP model using a heteroskedastic noise model.
Args:
train_X: A `batch_shape x n x d` tensor of training features.
train_Y: A `batch_shape x n x m` tensor of training observations.
train_Yvar: A `batch_shape x n x m` tensor of observed measurement
noise.
outcome_transform: An outcome transform that is applied to the
training data during instantiation and to the posterior during
inference (that is, the `Posterior` obtained by calling
`.posterior` on the model will be on the original scale).
Note that the noise model internally log-transforms the
variances, which will happen after this transform is applied.
Example:
>>> train_X = torch.rand(20, 2)
>>> train_Y = torch.sin(train_X).sum(dim=1, keepdim=True)
>>> se = torch.norm(train_X, dim=1, keepdim=True)
>>> train_Yvar = 0.1 + se * torch.rand_like(train_Y)
>>> model = HeteroskedasticSingleTaskGP(train_X, train_Y, train_Yvar)
"""
if outcome_transform is not None:
train_Y, train_Yvar = outcome_transform(train_Y, train_Yvar)
self._validate_tensor_args(X=train_X, Y=train_Y, Yvar=train_Yvar)
validate_input_scaling(train_X=train_X,
train_Y=train_Y,
train_Yvar=train_Yvar)
self._set_dimensions(train_X=train_X, train_Y=train_Y)
noise_likelihood = GaussianLikelihood(
noise_prior=SmoothedBoxPrior(-3, 5, 0.5, transform=torch.log),
batch_shape=self._aug_batch_shape,
noise_constraint=GreaterThan(MIN_INFERRED_NOISE_LEVEL,
transform=None,
initial_value=1.0),
)
noise_model = SingleTaskGP(
train_X=train_X,
train_Y=train_Yvar,
likelihood=noise_likelihood,
outcome_transform=Log(),
)
likelihood = _GaussianLikelihoodBase(HeteroskedasticNoise(noise_model))
super().__init__(train_X=train_X,
train_Y=train_Y,
likelihood=likelihood)
self.register_added_loss_term("noise_added_loss")
self.update_added_loss_term("noise_added_loss",
NoiseModelAddedLossTerm(noise_model))
if outcome_transform is not None:
self.outcome_transform = outcome_transform
self.to(train_X)
def test_initializations(self):
train_X = torch.rand(15, 1, device=self.device)
train_Y = torch.rand(15, 1, device=self.device)
stacked_train_X = torch.cat((train_X, train_X), dim=0)
for X, num_ind in [[train_X, 5], [stacked_train_X, 20], [stacked_train_X, 5]]:
model = SingleTaskVariationalGP(train_X=X, inducing_points=num_ind)
if num_ind == 5:
self.assertLessEqual(
model.model.variational_strategy.inducing_points.shape,
torch.Size((5, 1)),
)
else:
# should not have 20 inducing points when 15 singular dimensions
# are passed
self.assertLess(
model.model.variational_strategy.inducing_points.shape[-2], num_ind
)
test_X = torch.rand(5, 1, device=self.device)
# test transforms
for inp_trans, out_trans in itertools.product(
[None, Normalize(d=1)], [None, Log()]
):
model = SingleTaskVariationalGP(
train_X=train_X,
train_Y=train_Y,
outcome_transform=out_trans,
input_transform=inp_trans,
)
if inp_trans is not None:
self.assertIsInstance(model.input_transform, Normalize)
else:
self.assertFalse(hasattr(model, "input_transform"))
if out_trans is not None:
self.assertIsInstance(model.outcome_transform, Log)
posterior = model.posterior(test_X)
self.assertIsInstance(posterior, TransformedPosterior)
else:
self.assertFalse(hasattr(model, "outcome_transform"))
def test_chained_outcome_transform(self):
ms = (1, 2)
batch_shapes = (torch.Size(), torch.Size([2]))
dtypes = (torch.float, torch.double)
# test transform and untransform
for m, batch_shape, dtype in itertools.product(ms, batch_shapes, dtypes):
# test init
tf1 = Log()
tf2 = Standardize(m=m, batch_shape=batch_shape)
tf = ChainedOutcomeTransform(b=tf1, a=tf2)
self.assertTrue(tf.training)
self.assertEqual(list(tf.keys()), ["b", "a"])
self.assertEqual(tf["b"], tf1)
self.assertEqual(tf["a"], tf2)
# make copies for validation below
tf1_, tf2_ = deepcopy(tf1), deepcopy(tf2)
# no observation noise
Y = torch.rand(*batch_shape, 3, m, device=self.device, dtype=dtype)
Y_tf, Yvar_tf = tf(Y, None)
Y_tf_, Yvar_tf_ = tf2_(*tf1_(Y, None))
self.assertTrue(tf.training)
self.assertIsNone(Yvar_tf_)
self.assertTrue(torch.allclose(Y_tf, Y_tf_))
tf.eval()
self.assertFalse(tf.training)
Y_utf, Yvar_utf = tf.untransform(Y_tf, Yvar_tf)
torch.allclose(Y_utf, Y)
self.assertIsNone(Yvar_utf)
# subset_output
tf_subset = tf.subset_output(idcs=[0])
Y_tf_subset, Yvar_tf_subset = tf_subset(Y[..., [0]])
self.assertTrue(torch.equal(Y_tf[..., [0]], Y_tf_subset))
self.assertIsNone(Yvar_tf_subset)
with self.assertRaises(RuntimeError):
tf.subset_output(idcs=[0, 1, 2])
# test error if observation noise present
Y = torch.rand(*batch_shape, 3, m, device=self.device, dtype=dtype)
Yvar = 1e-8 + torch.rand(
*batch_shape, 3, m, device=self.device, dtype=dtype
)
with self.assertRaises(NotImplementedError):
tf(Y, Yvar)
# untransform_posterior
tf1 = Log()
tf2 = Standardize(m=m, batch_shape=batch_shape)
tf = ChainedOutcomeTransform(log=tf1, standardize=tf2)
Y_tf, Yvar_tf = tf(Y, None)
tf.eval()
shape = batch_shape + torch.Size([3, m])
posterior = _get_test_posterior(shape, device=self.device, dtype=dtype)
p_utf = tf.untransform_posterior(posterior)
self.assertIsInstance(p_utf, TransformedPosterior)
self.assertEqual(p_utf.device.type, self.device.type)
self.assertTrue(p_utf.dtype == dtype)
samples = p_utf.rsample()
self.assertEqual(samples.shape, torch.Size([1]) + shape)
samples = p_utf.rsample(sample_shape=torch.Size([4]))
self.assertEqual(samples.shape, torch.Size([4]) + shape)
samples2 = p_utf.rsample(sample_shape=torch.Size([4, 2]))
self.assertEqual(samples2.shape, torch.Size([4, 2]) + shape)
# test transforming a subset of outcomes
for batch_shape, dtype in itertools.product(batch_shapes, dtypes):
m = 2
outputs = [-1]
# test init
tf1 = Log(outputs=outputs)
tf2 = Standardize(m=m, outputs=outputs, batch_shape=batch_shape)
tf = ChainedOutcomeTransform(log=tf1, standardize=tf2)
self.assertTrue(tf.training)
self.assertEqual(sorted(tf.keys()), ["log", "standardize"])
self.assertEqual(tf["log"], tf1)
self.assertEqual(tf["standardize"], tf2)
self.assertEqual(tf["log"]._outputs, [-1]) # don't know dimension yet
self.assertEqual(tf["standardize"]._outputs, [1])
# make copies for validation below
tf1_, tf2_ = deepcopy(tf1), deepcopy(tf2)
# no observation noise
Y = torch.rand(*batch_shape, 3, m, device=self.device, dtype=dtype)
Y_tf, Yvar_tf = tf(Y, None)
Y_tf_, Yvar_tf_ = tf2_(*tf1_(Y, None))
self.assertTrue(tf.training)
self.assertIsNone(Yvar_tf_)
self.assertTrue(torch.allclose(Y_tf, Y_tf_))
tf.eval()
self.assertFalse(tf.training)
Y_utf, Yvar_utf = tf.untransform(Y_tf, Yvar_tf)
torch.allclose(Y_utf, Y)
self.assertIsNone(Yvar_utf)
# with observation noise
Y = torch.rand(*batch_shape, 3, m, device=self.device, dtype=dtype)
Yvar = 1e-8 + torch.rand(
*batch_shape, 3, m, device=self.device, dtype=dtype
)
with self.assertRaises(NotImplementedError):
tf(Y, Yvar)
# error on untransform_posterior
with self.assertRaises(NotImplementedError):
tf.untransform_posterior(None)
def test_log(self):
ms = (1, 2)
batch_shapes = (torch.Size(), torch.Size([2]))
dtypes = (torch.float, torch.double)
# test transform and untransform
for m, batch_shape, dtype in itertools.product(ms, batch_shapes, dtypes):
# test init
tf = Log()
self.assertTrue(tf.training)
self.assertIsNone(tf._outputs)
# no observation noise
Y = torch.rand(*batch_shape, 3, m, device=self.device, dtype=dtype)
Y_tf, Yvar_tf = tf(Y, None)
self.assertTrue(tf.training)
self.assertTrue(torch.allclose(Y_tf, torch.log(Y)))
self.assertIsNone(Yvar_tf)
tf.eval()
self.assertFalse(tf.training)
Y_utf, Yvar_utf = tf.untransform(Y_tf, Yvar_tf)
torch.allclose(Y_utf, Y)
self.assertIsNone(Yvar_utf)
# subset_output
tf_subset = tf.subset_output(idcs=[0])
Y_tf_subset, Yvar_tf_subset = tf_subset(Y[..., [0]])
self.assertTrue(torch.equal(Y_tf[..., [0]], Y_tf_subset))
self.assertIsNone(Yvar_tf_subset)
# test error if observation noise present
tf = Log()
Y = torch.rand(*batch_shape, 3, m, device=self.device, dtype=dtype)
Yvar = 1e-8 + torch.rand(
*batch_shape, 3, m, device=self.device, dtype=dtype
)
with self.assertRaises(NotImplementedError):
tf(Y, Yvar)
tf.eval()
with self.assertRaises(NotImplementedError):
tf.untransform(Y, Yvar)
# untransform_posterior
tf = Log()
Y_tf, Yvar_tf = tf(Y, None)
tf.eval()
shape = batch_shape + torch.Size([3, m])
posterior = _get_test_posterior(shape, device=self.device, dtype=dtype)
p_utf = tf.untransform_posterior(posterior)
self.assertIsInstance(p_utf, TransformedPosterior)
self.assertEqual(p_utf.device.type, self.device.type)
self.assertTrue(p_utf.dtype == dtype)
self.assertTrue(p_utf._sample_transform == torch.exp)
mean_expected = norm_to_lognorm_mean(posterior.mean, posterior.variance)
variance_expected = norm_to_lognorm_variance(
posterior.mean, posterior.variance
)
self.assertTrue(torch.allclose(p_utf.mean, mean_expected))
self.assertTrue(torch.allclose(p_utf.variance, variance_expected))
samples = p_utf.rsample()
self.assertEqual(samples.shape, torch.Size([1]) + shape)
samples = p_utf.rsample(sample_shape=torch.Size([4]))
self.assertEqual(samples.shape, torch.Size([4]) + shape)
samples2 = p_utf.rsample(sample_shape=torch.Size([4, 2]))
self.assertEqual(samples2.shape, torch.Size([4, 2]) + shape)
# test transforming a subset of outcomes
for batch_shape, dtype in itertools.product(batch_shapes, dtypes):
m = 2
outputs = [-1]
# test init
tf = Log(outputs=outputs)
self.assertTrue(tf.training)
# cannot normalize indices b/c we don't know dimension yet
self.assertEqual(tf._outputs, [-1])
# no observation noise
Y = torch.rand(*batch_shape, 3, m, device=self.device, dtype=dtype)
Y_tf, Yvar_tf = tf(Y, None)
self.assertTrue(tf.training)
self.assertTrue(torch.allclose(Y_tf[..., 1], torch.log(Y[..., 1])))
self.assertTrue(torch.allclose(Y_tf[..., 0], Y[..., 0]))
self.assertIsNone(Yvar_tf)
tf.eval()
self.assertFalse(tf.training)
Y_utf, Yvar_utf = tf.untransform(Y_tf, Yvar_tf)
torch.allclose(Y_utf, Y)
self.assertIsNone(Yvar_utf)
# subset_output
with self.assertRaises(NotImplementedError):
tf_subset = tf.subset_output(idcs=[0])
# with observation noise
tf = Log(outputs=outputs)
Y = torch.rand(*batch_shape, 3, m, device=self.device, dtype=dtype)
Yvar = 1e-8 + torch.rand(
*batch_shape, 3, m, device=self.device, dtype=dtype
)
with self.assertRaises(NotImplementedError):
tf(Y, Yvar)
# error on untransform_posterior
with self.assertRaises(NotImplementedError):
tf.untransform_posterior(None)
# test subset_output with positive on subset of outcomes (pos. index)
tf = Log(outputs=[0])
Y_tf, Yvar_tf = tf(Y, None)
tf_subset = tf.subset_output(idcs=[0])
Y_tf_subset, Yvar_tf_subset = tf_subset(Y[..., [0]], None)
self.assertTrue(torch.equal(Y_tf_subset, Y_tf[..., [0]]))
self.assertIsNone(Yvar_tf_subset)