def test_lkj_covariance_prior_validate_args(self): sd_prior = SmoothedBoxPrior(exp(-1), exp(1), validate_args=True) LKJCovariancePrior(2, 1.0, sd_prior) with self.assertRaises(ValueError): LKJCovariancePrior(1.5, 1.0, sd_prior, validate_args=True) with self.assertRaises(ValueError): LKJCovariancePrior(2, -1.0, sd_prior, validate_args=True)
def test_lkj_covariance_prior_log_prob_hetsd(self, cuda=False): device = torch.device("cuda") if cuda else torch.device("cpu") a = torch.tensor([exp(-1), exp(-2)], device=device) b = torch.tensor([exp(1), exp(2)], device=device) sd_prior = SmoothedBoxPrior(a, b, log_transform=True) prior = LKJCovariancePrior(2, torch.tensor(0.5, device=device), sd_prior) self.assertFalse(prior.log_transform) S = torch.eye(2, device=device) self.assertAlmostEqual(prior.log_prob(S).item(), -4.71958, places=4) S = torch.stack( [S, torch.tensor([[1.0, 0.5], [0.5, 1]], device=S.device)]) self.assertTrue( approx_equal(prior.log_prob(S), torch.tensor([-4.71958, -4.57574], device=S.device))) with self.assertRaises(ValueError): prior.log_prob(torch.eye(3, device=device)) # For eta=1.0 log_prob is flat over all covariance matrices prior = LKJCovariancePrior(2, torch.tensor(1.0, device=device), sd_prior) marginal_sd = torch.diagonal(S, dim1=-2, dim2=-1).sqrt() log_prob_expected = prior.correlation_prior.C + prior.sd_prior.log_prob( marginal_sd) self.assertTrue(approx_equal(prior.log_prob(S), log_prob_expected))
def test_lkj_prior_sample(self): prior = LKJCovariancePrior(2, 0.5, sd_prior=SmoothedBoxPrior(exp(-1), exp(1))) random_samples = prior.sample(torch.Size((6, ))) # need to check that these are positive semi-sefinite min_eval = torch.linalg.eigh(random_samples)[0].min() self.assertTrue(min_eval >= 0) # and that they are symmetric max_non_symm = (random_samples - random_samples.transpose(-1, -2)).abs().max() self.assertLess(max_non_symm, 1e-4) self.assertEqual(random_samples.shape, torch.Size((6, 2, 2)))
def test_lkj_covariance_prior_to_gpu(self): if torch.cuda.is_available(): sd_prior = SmoothedBoxPrior(exp(-1), exp(1)) prior = LKJCovariancePrior(2, 1.0, sd_prior).cuda() self.assertEqual(prior.correlation_prior.eta.device.type, "cuda") self.assertEqual(prior.correlation_prior.C.device.type, "cuda") self.assertEqual(prior.sd_prior.a.device.type, "cuda")
def test_lkj_covariance_prior_batch_log_prob(self, cuda=False): device = torch.device("cuda") if cuda else torch.device("cpu") v = torch.ones(2, 1, device=device) sd_prior = SmoothedBoxPrior(exp(-1) * v, exp(1) * v) prior = LKJCovariancePrior(2, torch.tensor([0.5, 1.5], device=device), sd_prior) corr_dist = LKJCholesky(2, torch.tensor([0.5, 1.5], device=device)) S = torch.eye(2, device=device) dist_log_prob = corr_dist.log_prob(S) + sd_prior.log_prob(S.diag()) self.assertLessEqual((prior.log_prob(S) - dist_log_prob).abs().sum(), 1e-4) S = torch.stack( [S, torch.tensor([[1.0, 0.5], [0.5, 1]], device=S.device)]) S_chol = torch.linalg.cholesky(S) dist_log_prob = corr_dist.log_prob(S_chol) + sd_prior.log_prob( torch.diagonal(S, dim1=-2, dim2=-1)) self.assertLessEqual((prior.log_prob(S) - dist_log_prob).abs().sum(), 1e-4)
def test_lkj_covariance_prior_log_prob_hetsd(self, cuda=False): device = torch.device("cuda") if cuda else torch.device("cpu") a = torch.tensor([exp(-1), exp(-2)], device=device) b = torch.tensor([exp(1), exp(2)], device=device) sd_prior = SmoothedBoxPrior(a, b) prior = LKJCovariancePrior(2, torch.tensor(0.5, device=device), sd_prior) corr_dist = LKJCholesky(2, torch.tensor(0.5, device=device)) S = torch.eye(2, device=device) dist_log_prob = corr_dist.log_prob(S) + sd_prior.log_prob( S.diag()).sum() self.assertAlmostEqual(prior.log_prob(S), dist_log_prob, places=4) S = torch.stack( [S, torch.tensor([[1.0, 0.5], [0.5, 1]], device=S.device)]) S_chol = torch.linalg.cholesky(S) dist_log_prob = corr_dist.log_prob(S_chol) + sd_prior.log_prob( torch.diagonal(S, dim1=-2, dim2=-1)) self.assertTrue(approx_equal(prior.log_prob(S), dist_log_prob))
def __init__(self, train_x, train_y, likelihood): super(HadamardMultitaskGPModel, self).__init__(train_x, train_y, likelihood) # Default bounds on mean are (-1e10, 1e10) self.mean_module = ConstantMean() # We use the very common RBF kernel self.covar_module = RBFKernel() # We learn an IndexKernel for 2 tasks # (so we'll actually learn 2x2=4 tasks with correlations) sd_prior = SmoothedBoxPrior(exp(-4), exp(4), log_transform=True) cov_prior = LKJCovariancePrior(n=2, eta=1, sd_prior=sd_prior) self.task_covar_module = IndexKernel(num_tasks=2, rank=1, prior=cov_prior)
def test_lkj_covariance_prior_batch_log_prob(self, cuda=False): device = torch.device("cuda") if cuda else torch.device("cpu") v = torch.ones(2, 1, device=device) sd_prior = SmoothedBoxPrior(exp(-1) * v, exp(1) * v) prior = LKJCovariancePrior(2, torch.tensor([0.5, 1.5], device=device), sd_prior) S = torch.eye(2, device=device) self.assertTrue(approx_equal(prior.log_prob(S), torch.tensor([-3.59981, -2.21351], device=S.device))) S = torch.stack([S, torch.tensor([[1.0, 0.5], [0.5, 1]], device=S.device)]) self.assertTrue(approx_equal(prior.log_prob(S), torch.tensor([-3.59981, -2.35735], device=S.device))) with self.assertRaises(ValueError): prior.log_prob(torch.eye(3, device=device))