Add better Error for initializers on intra-point constraints (#2775)

Summary:
Pull Request resolved: #2775

I have added the wanted exception.

Reviewed By: esantorella

Differential Revision: D71204155

fbshipit-source-id: dfe5ff0820bffbf69b4ac5c3e453daf1a20f1fdd

Author: SamuelGabriel

botorch/optim/initializers.py

@@ -541,6 +541,21 @@ def gen_one_shot_kg_initial_conditions(
         )
     q_aug = acq_function.get_augmented_q_batch_size(q=q)
 
+    # Check for inter-point (in)equality constraints
+    for constraints, type_of_constraints in [
+        (inequality_constraints, "inequality"),
+        (equality_constraints, "equality"),
+    ]:
+        if constraints is not None and len(constraints) > 0:
+            for indices, _, _ in constraints:
+                if len(indices.shape) > 1:
+                    raise NotImplementedError(
+                        "Indices must be one-dimensional "
+                        "in gen_one_shot_kg_initial_conditions. "
+                        f"Received indices {indices} "
+                        f"for {type_of_constraints} constrained."
+                    )
+
     # TODO: Avoid unnecessary computation by not generating all candidates
     ics = gen_batch_initial_conditions(
         acq_function=acq_function,
@@ -636,10 +651,12 @@ def gen_one_shot_hvkg_initial_conditions(
             restarts and raw samples for solving the posterior objective
             maximization problem, respectively) and `eta` (temperature parameter
             for sampling heuristic from posterior objective maximizers).
-        inequality constraints: A list of tuples (indices, coefficients, rhs),
+        inequality constraints: Optionally, list of tuples (indices, coefficients, rhs),
             with each tuple encoding an inequality constraint of the form
-            `\sum_i (X[indices[i]] * coefficients[i]) >= rhs`.
-        equality constraints: A list of tuples (indices, coefficients, rhs),
+            `\sum_i (X[indices[i]] * coefficients[i]) >= rhs`. Each
+            tensor of indices must be one-dimensional, since inter-point
+            constraints are not supported here.
+        equality constraints: Optionally, 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`.
 

test/optim/test_initializers.py

@@ -1046,40 +1046,89 @@ def test_gen_one_shot_kg_initial_conditions(self):
                     raw_samples=raw_samples,
                     options={"frac_random": 2.0},
                 )
+
+            forbidden_indices = torch.tensor([[0, 1], [0, 1]])
+            for constraint_type in [
+                "inequality_constraints",
+                "equality_constraints",
+            ]:
+                constraint_kwarg = {
+                    constraint_type: [
+                        (
+                            forbidden_indices,
+                            torch.ones(2).to(mean),
+                            0.5,
+                        )
+                    ]
+                }
+                with self.subTest(
+                    f"intra-point {constraint_type} not supported"
+                ), self.assertRaisesRegex(
+                    NotImplementedError,
+                    "Indices must be one-dimensional "
+                    "in gen_one_shot_kg_initial_conditions. "
+                    "Received indices",
+                ):
+                    gen_one_shot_kg_initial_conditions(
+                        acq_function=mock_kg,
+                        bounds=bounds,
+                        q=1,
+                        num_restarts=num_restarts,
+                        raw_samples=raw_samples,
+                        **constraint_kwarg,
+                    )
+
             # test generation logic
             q = 2
             mock_random_ics = torch.rand(num_restarts, q + num_fantasies, 2)
             mock_fantasy_cands = torch.ones(20, 1, 2)
             mock_fantasy_vals = torch.randn(20)
-            with ExitStack() as es:
-                mock_gbics = es.enter_context(
-                    mock.patch(
-                        "botorch.optim.initializers.gen_batch_initial_conditions",
-                        return_value=mock_random_ics,
-                    )
-                )
-                mock_optacqf = es.enter_context(
-                    mock.patch(
-                        "botorch.optim.optimize.optimize_acqf",
-                        return_value=(mock_fantasy_cands, mock_fantasy_vals),
+            for constraint_kwargs in [
+                {},
+                {
+                    "inequality_constraints": [
+                        (torch.tensor([0, 1]), torch.ones(2).to(mean), 0.5)
+                    ]
+                },  # test that no error is raised
+                {
+                    "equality_constraints": [
+                        (torch.tensor([0, 1]), torch.ones(2).to(mean), 0.5)
+                    ]
+                    # test that no error is raised
+                },
+            ]:
+                with ExitStack() as es:
+                    mock_gbics = es.enter_context(
+                        mock.patch(
+                            "botorch.optim.initializers.gen_batch_initial_conditions",
+                            return_value=mock_random_ics,
+                        )
                     )
-                )
-                ics = gen_one_shot_kg_initial_conditions(
-                    acq_function=mock_kg,
-                    bounds=bounds,
-                    q=q,
-                    num_restarts=num_restarts,
-                    raw_samples=raw_samples,
-                )
-                mock_gbics.assert_called_once()
-                mock_optacqf.assert_called_once()
-                n_value = int((1 - 0.1) * num_fantasies)
-                self.assertTrue(
-                    torch.equal(
-                        ics[..., :-n_value, :], mock_random_ics[..., :-n_value, :]
+                    mock_optacqf = es.enter_context(
+                        mock.patch(
+                            "botorch.optim.optimize.optimize_acqf",
+                            return_value=(mock_fantasy_cands, mock_fantasy_vals),
+                        )
                     )
-                )
-                self.assertTrue(torch.all(ics[..., -n_value:, :] == 1))
+                    with self.subTest(f"Main test with {constraint_kwargs}"):
+                        ics = gen_one_shot_kg_initial_conditions(
+                            acq_function=mock_kg,
+                            bounds=bounds,
+                            q=q,
+                            num_restarts=num_restarts,
+                            raw_samples=raw_samples,
+                            **constraint_kwargs,
+                        )
+                        mock_gbics.assert_called_once()
+                        mock_optacqf.assert_called_once()
+                        n_value = int((1 - 0.1) * num_fantasies)
+                        self.assertTrue(
+                            torch.equal(
+                                ics[..., :-n_value, :],
+                                mock_random_ics[..., :-n_value, :],
+                            )
+                        )
+                        self.assertTrue(torch.all(ics[..., -n_value:, :] == 1))
 
 
 class TestGenOneShotHVKGInitialConditions(BotorchTestCase):