Fix problem with polytope sampler (#1341)

Summary:
## Motivation

As discussed in issue #1225, the polytope sampler runs into problems if the variables live in different dimensions. This PR fixes the issue using the approach mentioned in the issue above.

### Have you read the [Contributing Guidelines on pull requests](https://github.com/pytorch/botorch/blob/main/CONTRIBUTING.md#pull-requests)?

Yes

Pull Request resolved: #1341

Test Plan: Unit tests

Reviewed By: esantorella

Differential Revision: D38485035

Pulled By: Balandat

fbshipit-source-id: 8e6ed1c1b81f636d108becb1ab5cf2ca13294420

Author: jduerholt

botorch/utils/sampling.py

@@ -830,6 +830,30 @@ def draw(self, n: int = 1, seed: Optional[int] = None) -> Tensor:
         return samples
 
 
+def normalize_linear_constraints(
+    bounds: Tensor, constraints: List[Tuple[Tensor, Tensor, float]]
+) -> List[Tuple[Tensor, Tensor, float]]:
+    r"""Normalize linear constraints to the unit cube.
+
+    Args:
+        bounds (Tensor): A `2 x d`-dim tensor containing the box bounds.
+        constraints (List[Tuple[Tensor, Tensor, float]]): A list of
+            tuples (indices, coefficients, rhs), with each tuple encoding
+            an inequality constraint of the form
+            `\sum_i (X[indices[i]] * coefficients[i]) >= rhs` or
+            `\sum_i (X[indices[i]] * coefficients[i]) = rhs`.
+    """
+
+    new_constraints = []
+    for index, coefficient, rhs in constraints:
+        lower, upper = bounds[:, index]
+        s = upper - lower
+        new_constraints.append(
+            (index, s * coefficient, (rhs - torch.dot(coefficient, lower)).item())
+        )
+    return new_constraints
+
+
 def get_polytope_samples(
     n: int,
     bounds: Tensor,
@@ -865,9 +889,11 @@ def get_polytope_samples(
     # create tensors representing linear inequality constraints
     # of the form Ax >= b.
     if inequality_constraints:
+        # normalize_linear_constraints is called to solve this issue:
+        # https://github.com/pytorch/botorch/issues/1225
         A, b = sparse_to_dense_constraints(
             d=bounds.shape[-1],
-            constraints=inequality_constraints,
+            constraints=normalize_linear_constraints(bounds, inequality_constraints),
         )
         # Note the inequality constraints are of the form Ax >= b,
         # but PolytopeSampler expects inequality constraints of the
@@ -876,19 +902,24 @@ def get_polytope_samples(
     else:
         dense_inequality_constraints = None
     if equality_constraints:
+        # normalize_linear_constraints is called to solve this issue:
+        # https://github.com/pytorch/botorch/issues/1225
         dense_equality_constraints = sparse_to_dense_constraints(
             d=bounds.shape[-1],
-            constraints=equality_constraints,
+            constraints=normalize_linear_constraints(bounds, equality_constraints),
         )
     else:
         dense_equality_constraints = None
+    normalized_bounds = torch.zeros_like(bounds)
+    normalized_bounds[1, :] = 1.0
     polytope_sampler = HitAndRunPolytopeSampler(
+        bounds=normalized_bounds,
         inequality_constraints=dense_inequality_constraints,
-        bounds=bounds,
         equality_constraints=dense_equality_constraints,
         n_burnin=n_burnin,
     )
-    return polytope_sampler.draw(n=n * thinning, seed=seed)[::thinning]
+    samples = polytope_sampler.draw(n=n * thinning, seed=seed)[::thinning]
+    return bounds[0] + samples * (bounds[1] - bounds[0])
 
 
 def sparse_to_dense_constraints(

test/optim/test_initializers.py

@@ -275,14 +275,14 @@ def test_gen_batch_initial_conditions_constraints(self):
             bounds = torch.tensor([[0, 0], [1, 1]], device=self.device, dtype=dtype)
             inequality_constraints = [
                 (
-                    torch.tensor([1], device=self.device, dtype=dtype),
+                    torch.tensor([1], device=self.device, dtype=torch.int64),
                     torch.tensor([-4], device=self.device, dtype=dtype),
                     torch.tensor(-3, device=self.device, dtype=dtype),
                 )
             ]
             equality_constraints = [
                 (
-                    torch.tensor([0], device=self.device, dtype=dtype),
+                    torch.tensor([0], device=self.device, dtype=torch.int64),
                     torch.tensor([1], device=self.device, dtype=dtype),
                     torch.tensor(0.5, device=self.device, dtype=dtype),
                 )

test/utils/test_sampling.py

@@ -28,6 +28,7 @@
     get_polytope_samples,
     HitAndRunPolytopeSampler,
     manual_seed,
+    normalize_linear_constraints,
     PolytopeSampler,
     sample_hypersphere,
     sample_simplex,
@@ -289,6 +290,28 @@ def test_sparse_to_dense_constraints(self):
             expected_b = torch.tensor([[3.0]], **tkwargs)
             self.assertTrue(torch.equal(b, expected_b))
 
+    def test_normalize_linear_constraints(self):
+        tkwargs = {"device": self.device}
+        for dtype in (torch.float, torch.double):
+            tkwargs["dtype"] = dtype
+            constraints = [
+                (
+                    torch.tensor([1, 2, 0], dtype=torch.int64, device=self.device),
+                    torch.tensor([1.0, 1.0, 1.0], **tkwargs),
+                    1.0,
+                )
+            ]
+            bounds = torch.tensor(
+                [[0.1, 0.3, 0.1, 30.0], [0.6, 0.7, 0.7, 700.0]], **tkwargs
+            )
+            new_constraints = normalize_linear_constraints(bounds, constraints)
+            expected_coefficients = torch.tensor([0.4000, 0.6000, 0.5000], **tkwargs)
+            self.assertTrue(
+                torch.allclose(new_constraints[0][1], expected_coefficients)
+            )
+            expected_rhs = 0.5
+            self.assertAlmostEqual(new_constraints[0][-1], expected_rhs)
+
     def test_find_interior_point(self):
         # basic problem: 1 <= x_1 <= 2, 2 <= x_2 <= 3
         A = np.concatenate([np.eye(2), -np.eye(2)], axis=0)
@@ -318,14 +341,14 @@ def test_get_polytope_samples(self):
             bounds[1] = 1
             inequality_constraints = [
                 (
-                    torch.tensor([3], **tkwargs),
+                    torch.tensor([3], dtype=torch.int64, device=self.device),
                     torch.tensor([-4], **tkwargs),
                     -3,
                 )
             ]
             equality_constraints = [
                 (
-                    torch.tensor([0], **tkwargs),
+                    torch.tensor([0], dtype=torch.int64, device=self.device),
                     torch.tensor([1], **tkwargs),
                     0.5,
                 )