Add setting_closure to PositiveIndexKernel for prior support (meta-pytorch#3267)

Summary:
Pull Request resolved: meta-pytorch#3267

`_set_lower_triangle_corr` method wired as the `setting_closure` in `register_prior`, enabling `sample_from_prior` / `sample_all_priors` to work with task correlation priors (e.g. BetaPrior). Values are validated to be in [0, 1] since PositiveIndexKernel enforces positive correlations.

Reviewed By: sdaulton

Differential Revision: D99841562

fbshipit-source-id: bf15dfdadfa917aa34c7c86b0ff29436d2ab7717

Author: Carl Hvarfner

botorch/models/kernels/positive_index.py

@@ -30,7 +30,6 @@ def __init__(
         num_tasks: int,
         rank: int = 1,
         task_prior: Prior | None = None,
-        diag_prior: Prior | None = None,
         normalize_covar_matrix: bool = False,
         var_constraint: Interval | None = None,
         target_task_index: int = 0,
@@ -43,7 +42,6 @@ def __init__(
             num_tasks (int): Total number of indices.
             rank (int): Rank of the covariance matrix parameterization.
             task_prior (Prior, optional): Prior for the covariance matrix.
-            diag_prior (Prior, optional): Prior for the diagonal elements.
             normalize_covar_matrix (bool): Whether to normalize the covariance matrix.
             target_task_index (int): Index of the task whose diagonal element should be
                 normalized to 1. Defaults to 0 (first task).
@@ -88,11 +86,11 @@ def __init__(
                     f"{type(task_prior).__name__}"
                 )
             self.register_prior(
-                "IndexKernelPrior", task_prior, lambda m: m._lower_triangle_corr
+                "IndexKernelPrior",
+                task_prior,
+                lambda m: m._lower_triangle_corr,
+                lambda m, v: m._set_lower_triangle_corr(v),
             )
-        if diag_prior is not None:
-            self.register_prior("ScalePrior", diag_prior, lambda m: m._diagonal)
-
         self.register_constraint("raw_covar_factor", GreaterThan(0.0))
 
     def _covar_factor_params(self, m):
@@ -127,15 +125,49 @@ def _lower_triangle_corr(self):
 
         return low_tri
 
-    @property
-    def _diagonal(self):
-        return torch.diagonal(self.covar_matrix, dim1=-2, dim2=-1)
+    def _set_lower_triangle_corr(self, value):
+        """Set covar_factor to produce the given lower-triangle correlations.
+
+        Assembles a symmetric correlation matrix from the lower-triangle values,
+        then projects it to the nearest positive-definite correlation matrix via
+        eigenvalue clamping before Cholesky decomposition. This guarantees the
+        setter never fails even when independently sampled correlation values
+        do not form a PD matrix.
+
+        Args:
+            value: Tensor of lower-triangle correlation values.
+        """
+        n = self.num_tasks
+        eps = 1e-6
+        n_lower = n * (n - 1) // 2
+        lower_row, lower_col = torch.tril_indices(n, n, offset=-1)
+        # Expand under-batched input (e.g. scalar from sample_from_prior)
+        if value.dim() == 0 or (value.dim() == 1 and value.shape[0] != n_lower):
+            value = value.unsqueeze(-1).expand(*self.batch_shape, n_lower)
+        elif value.shape[:-1] != self.batch_shape:
+            value = value.expand(*self.batch_shape, n_lower)
+        batch_shape = value.shape[:-1]
+        corr = (
+            torch.eye(n, dtype=value.dtype, device=value.device)
+            .expand(*batch_shape, n, n)
+            .clone()
+        )
+        corr[..., lower_row, lower_col] = value.clamp(0.0, 1.0)
+        corr[..., lower_col, lower_row] = value.clamp(0.0, 1.0)
+        # Project to nearest PD correlation matrix via eigenvalue clamping
+        eigvals, eigvecs = torch.linalg.eigh(corr)
+        eigvals = eigvals.clamp(min=eps)
+        corr = eigvecs @ torch.diag_embed(eigvals) @ eigvecs.transpose(-1, -2)
+        # Re-normalize diagonals to 1
+        d = corr.diagonal(dim1=-1, dim2=-2).sqrt()
+        corr = corr / (d.unsqueeze(-1) * d.unsqueeze(-2))
+        chol = torch.linalg.cholesky(corr)
+        rank = self.raw_covar_factor.shape[-1]
+        self._set_covar_factor(chol[..., :, :rank].clamp(min=eps))
 
     def _eval_covar_matrix(self):
         cf = self.covar_factor
-        covar = cf @ cf.transpose(-1, -2) + self.var * torch.eye(
-            self.num_tasks, dtype=cf.dtype, device=cf.device
-        )
+        covar = cf @ cf.transpose(-1, -2) + torch.diag_embed(self.var)
         # Normalize by the target task's diagonal element
         if self.unit_scale_for_target:
             norm_factor = covar[..., self.target_task_index, self.target_task_index]

test/models/kernels/test_positive_index.py

@@ -6,8 +6,10 @@
 
 import torch
 from botorch.models.kernels.positive_index import PositiveIndexKernel
+from botorch.models.utils.priors import BetaPrior
+from botorch.optim.utils import sample_all_priors
 from botorch.utils.testing import BotorchTestCase
-from gpytorch.priors import NormalPrior
+from gpytorch.priors import NormalPrior, UniformPrior
 
 
 class TestPositiveIndexKernel(BotorchTestCase):
@@ -125,18 +127,15 @@ def test_positive_index_kernel(self):
             with self.subTest("with_priors", dtype=dtype):
                 num_tasks = 4
                 task_prior = NormalPrior(0, 1)
-                diag_prior = NormalPrior(1, 0.1)
 
                 kernel = PositiveIndexKernel(
                     num_tasks=num_tasks,
                     rank=2,
                     task_prior=task_prior,
-                    diag_prior=diag_prior,
                     initialize_to_mode=False,
                 ).to(dtype=dtype)
                 prior_names = [p[0] for p in kernel.named_priors()]
                 self.assertIn("IndexKernelPrior", prior_names)
-                self.assertIn("ScalePrior", prior_names)
 
             # Test batch forward
             with self.subTest("batch_forward", dtype=dtype):
@@ -154,25 +153,6 @@ def test_positive_index_kernel(self):
                 # Check that batch dimensions are preserved
                 self.assertEqual(result.shape[0], 2)
 
-            # Test diagonal property (default target_task_index=0)
-            with self.subTest("diagonal", dtype=dtype):
-                kernel = PositiveIndexKernel(num_tasks=4, rank=2).to(dtype=dtype)
-                diag = kernel._diagonal
-
-                self.assertEqual(diag.shape, torch.Size([4]))
-                # First diagonal element should be 1.0 (default target_task_index=0)
-                self.assertAllClose(diag[0], torch.tensor(1.0, dtype=dtype), atol=1e-4)
-
-                # Test diagonal property with custom target_task_index
-                kernel = PositiveIndexKernel(
-                    num_tasks=4, rank=2, target_task_index=1
-                ).to(dtype=dtype)
-                diag = kernel._diagonal
-
-                self.assertEqual(diag.shape, torch.Size([4]))
-                # Second diagonal element should be 1.0 (target_task_index=1)
-                self.assertAllClose(diag[1], torch.tensor(1.0, dtype=dtype), atol=1e-4)
-
             # Test lower triangle property
             with self.subTest("lower_triangle", dtype=dtype):
                 num_tasks = 5
@@ -222,3 +202,153 @@ def test_positive_index_kernel(self):
                 new_value = torch.ones(3, 2, dtype=dtype) * 3.0
                 kernel._covar_factor_closure(kernel, new_value)
                 self.assertAllClose(kernel.covar_factor, new_value, atol=1e-5)
+
+            # Test _set_lower_triangle_corr produces valid covariance
+            with self.subTest("set_lower_triangle_corr", dtype=dtype):
+                kernel = PositiveIndexKernel(num_tasks=3, rank=3).to(dtype=dtype)
+                target_corr = torch.tensor([0.8, 0.5, 0.6], dtype=dtype)
+                kernel._set_lower_triangle_corr(target_corr)
+
+                # Covariance matrix should be PD and symmetric
+                covar = kernel.covar_matrix
+                eigvals = torch.linalg.eigvalsh(covar)
+                self.assertTrue((eigvals > 0).all())
+                self.assertAllClose(covar, covar.T, atol=1e-5)
+
+                # Recovered correlations should be positive
+                recovered = kernel._lower_triangle_corr
+                self.assertTrue((recovered >= 0).all())
+                self.assertTrue((recovered <= 1).all())
+
+            # Test _set_lower_triangle_corr with batch shape
+            with self.subTest("set_lower_triangle_corr_batch", dtype=dtype):
+                batch_shape = torch.Size([2])
+                kernel = PositiveIndexKernel(
+                    num_tasks=3, rank=3, batch_shape=batch_shape
+                ).to(dtype=dtype)
+                target_corr = torch.rand(*batch_shape, 3, dtype=dtype)
+                kernel._set_lower_triangle_corr(target_corr)
+                covar = kernel.covar_matrix
+                eigvals = torch.linalg.eigvalsh(covar)
+                self.assertTrue((eigvals > 0).all())
+                self.assertEqual(covar.shape, torch.Size([2, 3, 3]))
+
+            # Test sample_all_priors with batch shape
+            with self.subTest("sample_all_priors_batch", dtype=dtype):
+                batch_shape = torch.Size([2])
+                task_prior = UniformPrior(0.0, 1.0)
+                kernel = PositiveIndexKernel(
+                    num_tasks=3,
+                    rank=3,
+                    task_prior=task_prior,
+                    batch_shape=batch_shape,
+                ).to(dtype=dtype)
+                sample_all_priors(kernel)
+                covar = kernel.covar_matrix
+                eigvals = torch.linalg.eigvalsh(covar)
+                self.assertTrue((eigvals > 0).all())
+                self.assertEqual(covar.shape, torch.Size([2, 3, 3]))
+
+            # Test _set_lower_triangle_corr with scalar input (under-batched)
+            with self.subTest("set_lower_triangle_corr_scalar", dtype=dtype):
+                batch_shape = torch.Size([2])
+                kernel = PositiveIndexKernel(
+                    num_tasks=3, rank=3, batch_shape=batch_shape
+                ).to(dtype=dtype)
+                # Scalar value — exercises dim()==0 branch
+                kernel._set_lower_triangle_corr(torch.tensor(0.5, dtype=dtype))
+                covar = kernel.covar_matrix
+                eigvals = torch.linalg.eigvalsh(covar)
+                self.assertTrue((eigvals > 0).all())
+                self.assertEqual(covar.shape, torch.Size([2, 3, 3]))
+
+            # Test _set_lower_triangle_corr with unbatched input on batched kernel
+            with self.subTest(
+                "set_lower_triangle_corr_unbatched_on_batch", dtype=dtype
+            ):
+                batch_shape = torch.Size([2])
+                kernel = PositiveIndexKernel(
+                    num_tasks=3, rank=3, batch_shape=batch_shape
+                ).to(dtype=dtype)
+                # 1D input with correct n_lower but no batch — exercises expand branch
+                target_corr = torch.rand(3, dtype=dtype)
+                kernel._set_lower_triangle_corr(target_corr)
+                covar = kernel.covar_matrix
+                eigvals = torch.linalg.eigvalsh(covar)
+                self.assertTrue((eigvals > 0).all())
+                self.assertEqual(covar.shape, torch.Size([2, 3, 3]))
+
+            # Test _set_lower_triangle_corr with boundary values
+            with self.subTest("set_lower_triangle_corr_boundary", dtype=dtype):
+                kernel = PositiveIndexKernel(num_tasks=2, rank=2).to(dtype=dtype)
+                kernel._set_lower_triangle_corr(torch.tensor([0.0], dtype=dtype))
+                self.assertTrue(kernel._lower_triangle_corr.isfinite().all())
+                kernel._set_lower_triangle_corr(torch.tensor([0.999], dtype=dtype))
+                self.assertTrue(kernel._lower_triangle_corr.isfinite().all())
+
+            # Test _set_lower_triangle_corr with non-PD input
+            with self.subTest("set_lower_triangle_corr_non_pd", dtype=dtype):
+                kernel = PositiveIndexKernel(num_tasks=3, rank=3).to(dtype=dtype)
+                # [0.99, 0.01, 0.99] does not form a PD correlation matrix
+                non_pd_corr = torch.tensor([0.99, 0.01, 0.99], dtype=dtype)
+                kernel._set_lower_triangle_corr(non_pd_corr)
+                covar = kernel.covar_matrix
+                eigvals = torch.linalg.eigvalsh(covar)
+                self.assertTrue((eigvals > 0).all())
+
+            # Test roundtrip accuracy for full-rank
+            with self.subTest("set_lower_triangle_corr_roundtrip", dtype=dtype):
+                kernel = PositiveIndexKernel(
+                    num_tasks=3, rank=3, unit_scale_for_target=False
+                ).to(dtype=dtype)
+                # Set var to small known value to isolate correlation effect
+                kernel.initialize(raw_var=torch.full((3,), -5.0, dtype=dtype))
+                target_corr = torch.tensor([0.8, 0.5, 0.6], dtype=dtype)
+                kernel._set_lower_triangle_corr(target_corr)
+                recovered = kernel._lower_triangle_corr
+                self.assertAllClose(recovered, target_corr, atol=0.05)
+
+            # Test sample_all_priors with task_prior
+            with self.subTest("sample_all_priors_unbatched", dtype=dtype):
+                task_prior = UniformPrior(0.0, 1.0)
+                kernel = PositiveIndexKernel(
+                    num_tasks=3,
+                    rank=3,
+                    task_prior=task_prior,
+                ).to(dtype=dtype)
+
+                corr_before = kernel._lower_triangle_corr.clone()
+                sample_all_priors(kernel)
+
+                corr_after = kernel._lower_triangle_corr
+                self.assertFalse(torch.allclose(corr_before, corr_after))
+
+                covar = kernel.covar_matrix
+                eigvals = torch.linalg.eigvalsh(covar)
+                self.assertTrue((eigvals > 0).all())
+
+            # Test with BetaPrior
+            with self.subTest("beta_prior", dtype=dtype):
+                task_prior = BetaPrior(1.2, 0.9)
+                kernel = PositiveIndexKernel(
+                    num_tasks=4,
+                    rank=4,
+                    task_prior=task_prior,
+                ).to(dtype=dtype)
+                sample_all_priors(kernel)
+                covar = kernel.covar_matrix
+                eigvals = torch.linalg.eigvalsh(covar)
+                self.assertTrue((eigvals > 0).all())
+
+            # Test sample_all_priors
+            with self.subTest("sample_all_priors", dtype=dtype):
+                task_prior = UniformPrior(0.0, 1.0)
+                kernel = PositiveIndexKernel(
+                    num_tasks=3,
+                    rank=3,
+                    task_prior=task_prior,
+                ).to(dtype=dtype)
+                sample_all_priors(kernel)
+                covar = kernel.covar_matrix
+                eigvals = torch.linalg.eigvalsh(covar)
+                self.assertTrue((eigvals > 0).all())