def untransform_posterior(self, posterior: Posterior) -> Posterior:
r"""Un-standardize the posterior.
Args:
posterior: A posterior in the standardized space.
Returns:
The un-standardized posterior. If the input posterior is a MVN,
the transformed posterior is again an MVN.
"""
if self._outputs is not None:
raise NotImplementedError(
"Standardize does not yet support output selection for "
"untransform_posterior"
)
if not self._m == posterior.event_shape[-1]:
raise RuntimeError(
"Incompatible output dimensions encountered for transform "
f"{self._m} and posterior {posterior.event_shape[-1]}"
)
if not isinstance(posterior, GPyTorchPosterior):
# fall back to TransformedPosterior
return TransformedPosterior(
posterior=posterior,
sample_transform=lambda s: self.means + self.stdvs * s,
mean_transform=lambda m, v: self.means + self.stdvs * m,
variance_transform=lambda m, v: self._stdvs_sq * v,
)
# GPyTorchPosterior (TODO: Should we Lazy-evaluate the mean here as well?)
mvn = posterior.mvn
offset = self.means
scale_fac = self.stdvs
if not posterior._is_mt:
mean_tf = offset.squeeze(-1) + scale_fac.squeeze(-1) * mvn.mean
scale_fac = scale_fac.squeeze(-1).expand_as(mean_tf)
else:
mean_tf = offset + scale_fac * mvn.mean
reps = mean_tf.shape[-2:].numel() // scale_fac.size(-1)
scale_fac = scale_fac.squeeze(-2)
if mvn._interleaved:
scale_fac = scale_fac.repeat(*[1 for _ in scale_fac.shape[:-1]], reps)
else:
scale_fac = torch.repeat_interleave(scale_fac, reps, dim=-1)
if (
not mvn.islazy
# TODO: Figure out attribute namming weirdness here
or mvn._MultivariateNormal__unbroadcasted_scale_tril is not None
):
# if already computed, we can save a lot of time using scale_tril
covar_tf = CholLazyTensor(mvn.scale_tril * scale_fac.unsqueeze(-1))
else:
lcv = mvn.lazy_covariance_matrix
# allow batch-evaluation of the model
scale_mat = DiagLazyTensor(scale_fac.expand(lcv.shape[:-1]))
covar_tf = scale_mat @ lcv @ scale_mat
kwargs = {"interleaved": mvn._interleaved} if posterior._is_mt else {}
mvn_tf = mvn.__class__(mean=mean_tf, covariance_matrix=covar_tf, **kwargs)
return GPyTorchPosterior(mvn_tf)
def untransform_posterior(
self, posterior: HigherOrderGPPosterior) -> TransformedPosterior:
# TODO: return a HigherOrderGPPosterior once rescaling constant
# muls * LazyTensors won't force a dense decomposition rather than a
# Kronecker structured one.
return TransformedPosterior(
posterior=posterior,
sample_transform=lambda s: self._return_to_output_shape(
self.means + self.stdvs * self._squeeze_to_single_output(s)),
mean_transform=lambda m, v: self._return_to_output_shape(
self.means + self.stdvs * self._squeeze_to_single_output(m)),
variance_transform=lambda m, v: self._return_to_output_shape(
self._stdvs_sq * self._squeeze_to_single_output(v)),
)
def untransform_posterior(self, posterior: Posterior) -> Posterior:
r"""Un-transform the power-transformed posterior.
Args:
posterior: A posterior in the power-transformed space.
Returns:
The un-transformed posterior.
"""
if self._outputs is not None:
raise NotImplementedError(
"Power does not yet support output selection for untransform_posterior"
)
return TransformedPosterior(
posterior=posterior,
sample_transform=lambda x: x.pow(1.0 / self.power),
)
def untransform_posterior(self, posterior: Posterior) -> Posterior:
r"""Un-transform the log-transformed posterior.
Args:
posterior: A posterior in the log-transformed space.
Returns:
The un-transformed posterior.
"""
if self._outputs is not None:
raise NotImplementedError(
"Log does not yet support output selection for untransform_posterior"
)
return TransformedPosterior(
posterior=posterior,
sample_transform=torch.exp,
mean_transform=norm_to_lognorm_mean,
variance_transform=norm_to_lognorm_variance,
)
def test_standardize(self):
# test error on incompatible dim
tf = Standardize(m=1)
with self.assertRaises(RuntimeError):
tf(torch.zeros(3, 2, device=self.device), None)
# test error on incompatible batch shape
with self.assertRaises(RuntimeError):
tf(torch.zeros(2, 3, 1, device=self.device), None)
ms = (1, 2)
batch_shapes = (torch.Size(), torch.Size([2]))
dtypes = (torch.float, torch.double)
# test transform, untransform, untransform_posterior
for m, batch_shape, dtype in itertools.product(ms, batch_shapes, dtypes):
# test init
tf = Standardize(m=m, batch_shape=batch_shape)
self.assertTrue(tf.training)
self.assertEqual(tf._m, m)
self.assertIsNone(tf._outputs)
self.assertEqual(tf._batch_shape, batch_shape)
self.assertEqual(tf._min_stdv, 1e-8)
# 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.all(Y_tf.mean(dim=-2).abs() < 1e-4))
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)
with self.assertRaises(RuntimeError):
tf.subset_output(idcs=[0, 1, 2])
# with observation noise
tf = Standardize(m=m, batch_shape=batch_shape)
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
)
Y_tf, Yvar_tf = tf(Y, Yvar)
self.assertTrue(tf.training)
self.assertTrue(torch.all(Y_tf.mean(dim=-2).abs() < 1e-4))
Yvar_tf_expected = Yvar / Y.std(dim=-2, keepdim=True) ** 2
self.assertTrue(torch.allclose(Yvar_tf, Yvar_tf_expected))
tf.eval()
self.assertFalse(tf.training)
Y_utf, Yvar_utf = tf.untransform(Y_tf, Yvar_tf)
torch.allclose(Y_utf, Y)
torch.allclose(Yvar_utf, Yvar)
# untransform_posterior
for interleaved, lazy in itertools.product((True, False), (True, False)):
if m == 1 and interleaved: # interleave has no meaning for m=1
continue
shape = batch_shape + torch.Size([3, m])
posterior = _get_test_posterior(
shape,
device=self.device,
dtype=dtype,
interleaved=interleaved,
lazy=lazy,
)
p_utf = tf.untransform_posterior(posterior)
self.assertEqual(p_utf.device.type, self.device.type)
self.assertTrue(p_utf.dtype == dtype)
mean_expected = tf.means + tf.stdvs * posterior.mean
variance_expected = tf.stdvs ** 2 * 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)
# TODO: Test expected covar (both interleaved and non-interleaved)
# untransform_posterior for non-GPyTorch posterior
posterior2 = TransformedPosterior(
posterior=posterior,
sample_transform=lambda s: s,
mean_transform=lambda m, v: m,
variance_transform=lambda m, v: v,
)
p_utf2 = tf.untransform_posterior(posterior2)
self.assertEqual(p_utf2.device.type, self.device.type)
self.assertTrue(p_utf2.dtype == dtype)
mean_expected = tf.means + tf.stdvs * posterior.mean
variance_expected = tf.stdvs ** 2 * posterior.variance
self.assertTrue(torch.allclose(p_utf2.mean, mean_expected))
self.assertTrue(torch.allclose(p_utf2.variance, variance_expected))
# TODO: Test expected covar (both interleaved and non-interleaved)
samples = p_utf2.rsample()
self.assertEqual(samples.shape, torch.Size([1]) + shape)
samples = p_utf2.rsample(sample_shape=torch.Size([4]))
self.assertEqual(samples.shape, torch.Size([4]) + shape)
samples2 = p_utf2.rsample(sample_shape=torch.Size([4, 2]))
self.assertEqual(samples2.shape, torch.Size([4, 2]) + shape)
# test error on incompatible output dimension
tf_big = Standardize(m=4).eval()
with self.assertRaises(RuntimeError):
tf_big.untransform_posterior(posterior2)
# test transforming a subset of outcomes
for batch_shape, dtype in itertools.product(batch_shapes, dtypes):
m = 2
outputs = [-1]
# test init
tf = Standardize(m=m, outputs=outputs, batch_shape=batch_shape)
self.assertTrue(tf.training)
self.assertEqual(tf._m, m)
self.assertEqual(tf._outputs, [1])
self.assertEqual(tf._batch_shape, batch_shape)
self.assertEqual(tf._min_stdv, 1e-8)
# 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)
Y_tf_mean = Y_tf.mean(dim=-2)
self.assertTrue(torch.all(Y_tf_mean[..., 1].abs() < 1e-4))
self.assertTrue(torch.allclose(Y_tf_mean[..., 0], Y.mean(dim=-2)[..., 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
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])
# with observation noise
tf = Standardize(m=m, outputs=outputs, batch_shape=batch_shape)
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
)
Y_tf, Yvar_tf = tf(Y, Yvar)
self.assertTrue(tf.training)
Y_tf_mean = Y_tf.mean(dim=-2)
self.assertTrue(torch.all(Y_tf_mean[..., 1].abs() < 1e-4))
self.assertTrue(torch.allclose(Y_tf_mean[..., 0], Y.mean(dim=-2)[..., 0]))
Yvar_tf_expected = Yvar / Y.std(dim=-2, keepdim=True) ** 2
self.assertTrue(torch.allclose(Yvar_tf[..., 1], Yvar_tf_expected[..., 1]))
self.assertTrue(torch.allclose(Yvar_tf[..., 0], Yvar[..., 0]))
tf.eval()
self.assertFalse(tf.training)
Y_utf, Yvar_utf = tf.untransform(Y_tf, Yvar_tf)
torch.allclose(Y_utf, Y)
torch.allclose(Yvar_utf, Yvar)
# error on untransform_posterior
with self.assertRaises(NotImplementedError):
tf.untransform_posterior(None)
