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)