diff --git a/botorch/optim/initializers.py b/botorch/optim/initializers.py index 8b392019c9..54103745a8 100644 --- a/botorch/optim/initializers.py +++ b/botorch/optim/initializers.py @@ -63,6 +63,177 @@ ] +def transform_constraints( + constraints: Union[List[Tuple[Tensor, Tensor, float]], None], q: int, d: int +) -> List[Tuple[Tensor, Tensor, float]]: + r"""Transform constraints to sample from a d*q-dimensional space instead of a + d-dimensional state. + + This function assumes that constraints are the same for each input batch, + and broadcasts the constraints accordingly to the input batch shape. + + Args: + constraints: A list of tuples (indices, coefficients, rhs), with each tuple + encoding an (in-)equality constraint of the form + `\sum_i (X[indices[i]] * coefficients[i]) (>)= rhs`. + If `indices` is a 2-d Tensor, this supports specifying constraints across + the points in the `q`-batch (inter-point constraints). If `None`, this + function is a nullop and simply returns `None`. + q: Size of the `q`-batch. + d: Dimensionality of the problem. + + Returns: + List[Tuple[Tensor, Tensor, float]]: List of transformed constraints. + """ + if constraints is None: + return None + transformed = [] + for constraint in constraints: + if len(constraint[0].shape) == 1: + transformed += transform_intra_point_constraint(constraint, d, q) + else: + transformed.append(transform_inter_point_constraint(constraint, d)) + return transformed + + +def transform_intra_point_constraint( + constraint: Tuple[Tensor, Tensor, float], d: int, q: int +) -> List[Tuple[Tensor, Tensor, float]]: + r"""Transforms an intra-point/pointwise constraint from + d-dimensional space to a d*q-dimesional space. + + Args: + constraints: A list of tuples (indices, coefficients, rhs), with each tuple + encoding an (in-)equality constraint of the form + `\sum_i (X[indices[i]] * coefficients[i]) (>)= rhs`. Here `indices` must + be one-dimensional, and the constraint is applied to all points within the + `q`-batch. + d: Dimensionality of the problem. + + Raises: + ValueError: If indices in the constraints are larger than the + dimensionality d of the problem. + + Returns: + List[Tuple[Tensor, Tensor, float]]: List of transformed constraints. + """ + indices, coefficients, rhs = constraint + if indices.max() >= d: + raise ValueError( + f"Constraint indices cannot exceed the problem dimension {d=}." + ) + return [ + ( + torch.tensor( + [i * d + j for j in indices], dtype=torch.int64, device=indices.device + ), + coefficients, + rhs, + ) + for i in range(q) + ] + + +def transform_inter_point_constraint( + constraint: Tuple[Tensor, Tensor, float], d: int +) -> Tuple[Tensor, Tensor, float]: + r"""Transforms an inter-point constraint from + d-dimensional space to a d*q dimesional space. + + Args: + constraints: A list of tuples (indices, coefficients, rhs), with each tuple + encoding an (in-)equality constraint of the form + `\sum_i (X[indices[i]] * coefficients[i]) (>)= rhs`. `indices` must be a + 2-d Tensor, where in each row `indices[i] = (k_i, l_i)` the first index + `k_i` corresponds to the `k_i`-th element of the `q`-batch and the second + index `l_i` corresponds to the `l_i`-th feature of that element. + + Raises: + ValueError: If indices in the constraints are larger than the + dimensionality d of the problem. + + Returns: + List[Tuple[Tensor, Tensor, float]]: Transformed constraint. + """ + indices, coefficients, rhs = constraint + if indices[:, 1].max() >= d: + raise ValueError( + f"Constraint indices cannot exceed the problem dimension {d=}." + ) + return ( + torch.tensor( + [r[0] * d + r[1] for r in indices], dtype=torch.int64, device=indices.device + ), + coefficients, + rhs, + ) + + +def sample_q_batches_from_polytope( + n: int, + q: int, + bounds: Tensor, + n_burnin: int, + thinning: int, + seed: int, + inequality_constraints: Optional[List[Tuple[Tensor, Tensor, float]]] = None, + equality_constraints: Optional[List[Tuple[Tensor, Tensor, float]]] = None, +) -> Tensor: + r"""Samples `n` q-baches from a polytope of dimension `d`. + + Args: + n: Number of q-batches to sample. + q: Number of samples per q-batch + bounds: A `2 x d` tensor of lower and upper bounds for each column of `X`. + n_burnin: The number of burn-in samples for the Markov chain sampler. + thinning: The amount of thinning (number of steps to take between + returning samples). + seed: The random seed. + inequality constraints: 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`. + equality constraints: 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`. + + Returns: + A `n x q x d`-dim tensor of samples. + """ + + # check if inter-point constraints are present + inter_point = any( + len(indices.shape) > 1 + for constraints in (inequality_constraints or [], equality_constraints or []) + for indices, _, _ in constraints + ) + + if inter_point: + samples = get_polytope_samples( + n=n, + bounds=torch.hstack([bounds for _ in range(q)]), + inequality_constraints=transform_constraints( + constraints=inequality_constraints, q=q, d=bounds.shape[1] + ), + equality_constraints=transform_constraints( + constraints=equality_constraints, q=q, d=bounds.shape[1] + ), + seed=seed, + n_burnin=n_burnin, + thinning=thinning * q, + ) + else: + samples = get_polytope_samples( + n=n * q, + bounds=bounds, + inequality_constraints=inequality_constraints, + equality_constraints=equality_constraints, + seed=seed, + n_burnin=n_burnin, + thinning=thinning, + ) + return samples.view(n, q, -1).cpu() + + def gen_batch_initial_conditions( acq_function: AcquisitionFunction, bounds: Tensor, @@ -167,18 +338,15 @@ def gen_batch_initial_conditions( ) X_rnd = bounds_cpu[0] + (bounds_cpu[1] - bounds_cpu[0]) * X_rnd_nlzd else: - X_rnd = ( - get_polytope_samples( - n=n * q, - bounds=bounds, - inequality_constraints=inequality_constraints, - equality_constraints=equality_constraints, - seed=seed, - n_burnin=options.get("n_burnin", 10000), - thinning=options.get("thinning", 32), - ) - .view(n, q, -1) - .cpu() + X_rnd = sample_q_batches_from_polytope( + n=n, + q=q, + bounds=bounds, + n_burnin=options.get("n_burnin", 10000), + thinning=options.get("thinning", 32), + seed=seed, + equality_constraints=equality_constraints, + inequality_constraints=inequality_constraints, ) # sample points around best if sample_around_best: diff --git a/botorch/optim/optimize.py b/botorch/optim/optimize.py index 30b7ed5dcf..065f1c0816 100644 --- a/botorch/optim/optimize.py +++ b/botorch/optim/optimize.py @@ -97,6 +97,22 @@ def __post_init__(self) -> None: "bounds should be a `2 x d` tensor, current shape: " f"{list(self.bounds.shape)}." ) + # validate that linear constraints across the q-dim and + # self.sequential are not present together + if self.inequality_constraints is not None and self.sequential is True: + for constraint in self.inequality_constraints: + if len(constraint[0].shape) > 1: + raise UnsupportedError( + "Linear inequality constraints across the q-dimension are not " + "supported for sequential optimization." + ) + if self.equality_constraints is not None and self.sequential is True: + for constraint in self.equality_constraints: + if len(constraint[0].shape) > 1: + raise UnsupportedError( + "Linear equality constraints across the q-dimension are not " + "supported for sequential optimization." + ) # TODO: Validate constraints if provided: # https://github.com/pytorch/botorch/pull/1231 diff --git a/test/optim/test_initializers.py b/test/optim/test_initializers.py index 9da9632a97..e95d4278a6 100644 --- a/test/optim/test_initializers.py +++ b/test/optim/test_initializers.py @@ -33,7 +33,11 @@ gen_value_function_initial_conditions, sample_perturbed_subset_dims, sample_points_around_best, + sample_q_batches_from_polytope, sample_truncated_normal_perturbations, + transform_constraints, + transform_inter_point_constraint, + transform_intra_point_constraint, ) from botorch.sampling.normal import IIDNormalSampler from botorch.utils.sampling import draw_sobol_samples @@ -270,6 +274,186 @@ def test_gen_batch_initial_conditions_warning(self): ) ) + def test_gen_batch_initial_conditions_transform_intra_point_constraint(self): + for dtype in (torch.float, torch.double): + constraint = ( + torch.tensor([0, 1], dtype=torch.int64, device=self.device), + torch.tensor([-1, -1]).to(dtype=dtype, device=self.device), + -1.0, + ) + constraints = transform_intra_point_constraint( + constraint=constraint, d=3, q=3 + ) + self.assertEqual(len(constraints), 3) + self.assertAllClose( + constraints[0][0], + torch.tensor([0, 1], dtype=torch.int64, device=self.device), + ) + self.assertAllClose( + constraints[1][0], + torch.tensor([3, 4], dtype=torch.int64, device=self.device), + ) + self.assertAllClose( + constraints[2][0], + torch.tensor([6, 7], dtype=torch.int64, device=self.device), + ) + for constraint in constraints: + self.assertAllClose( + torch.tensor([-1, -1]).to(dtype=dtype), constraint[1] + ) + self.assertEqual(constraint[2], -1.0) + # test failure on invalid d + constraint = ( + torch.tensor([[0, 3]]), + torch.tensor([-1.0, -1.0], dtype=dtype), + 0, + ) + with self.assertRaisesRegex( + ValueError, + "Constraint indices cannot exceed the problem dimension d=3.", + ): + transform_intra_point_constraint(constraint=constraint, d=3, q=2) + + def test_gen_batch_intial_conditions_transform_inter_point_constraint(self): + for dtype in (torch.float, torch.double): + constraint = ( + torch.tensor([[0, 1], [1, 1]], dtype=torch.int64, device=self.device), + torch.tensor([1.0, -1.0], dtype=dtype, device=self.device), + 0, + ) + transformed = transform_inter_point_constraint(constraint=constraint, d=3) + self.assertAllClose( + transformed[0], + torch.tensor([1, 4], dtype=torch.int64, device=self.device), + ) + self.assertAllClose( + transformed[1], + torch.tensor([1.0, -1.0]).to(dtype=dtype, device=self.device), + ) + self.assertEqual(constraint[2], 0.0) + # test failure on invalid d + constraint = ( + torch.tensor([[0, 1], [1, 3]], dtype=torch.int64, device=self.device), + torch.tensor([1.0, -1.0], dtype=dtype, device=self.device), + 0, + ) + with self.assertRaisesRegex( + ValueError, + "Constraint indices cannot exceed the problem dimension d=3.", + ): + transform_inter_point_constraint(constraint=constraint, d=3) + + def test_gen_batch_initial_conditions_transform_constraints(self): + for dtype in (torch.float, torch.double): + # test with None + self.assertIsNone(transform_constraints(constraints=None, d=3, q=3)) + constraints = [ + ( + torch.tensor([0, 1], dtype=torch.int64, device=self.device), + torch.tensor([-1.0, -1.0], dtype=dtype, device=self.device), + -1.0, + ), + ( + torch.tensor( + [[0, 1], [1, 1]], device=self.device, dtype=torch.int64 + ), + torch.tensor([1.0, -1.0], dtype=dtype, device=self.device), + 0, + ), + ] + transformed = transform_constraints(constraints=constraints, d=3, q=3) + self.assertEqual(len(transformed), 4) + self.assertAllClose( + transformed[0][0], + torch.tensor([0, 1], dtype=torch.int64, device=self.device), + ) + self.assertAllClose( + transformed[1][0], + torch.tensor([3, 4], dtype=torch.int64, device=self.device), + ) + self.assertAllClose( + transformed[2][0], + torch.tensor([6, 7], dtype=torch.int64, device=self.device), + ) + for constraint in transformed[:3]: + self.assertAllClose( + torch.tensor([-1, -1], dtype=dtype, device=self.device), + constraint[1], + ) + self.assertEqual(constraint[2], -1.0) + self.assertAllClose( + transformed[-1][0], + torch.tensor([1, 4], dtype=torch.int64, device=self.device), + ) + self.assertAllClose( + transformed[-1][1], + torch.tensor([1.0, -1.0], dtype=dtype, device=self.device), + ) + self.assertEqual(transformed[-1][2], 0.0) + + def test_gen_batch_initial_conditions_sample_q_batches_from_polytope(self): + for dtype in (torch.float, torch.double): + bounds = torch.tensor( + [[0, 0, 0], [1, 1, 1]], device=self.device, dtype=dtype + ) + inequality_constraints = [ + ( + torch.tensor([0, 1], device=self.device, dtype=torch.int64), + torch.tensor([-1, 1], device=self.device, dtype=dtype), + torch.tensor(-0.5, device=self.device, dtype=dtype), + ) + ] + inter_point_inequality_constraints = [ + ( + torch.tensor([0, 1], device=self.device, dtype=torch.int64), + torch.tensor([-1, 1], device=self.device, dtype=dtype), + torch.tensor(-0.5, device=self.device, dtype=dtype), + ), + ( + torch.tensor( + [[0, 1], [1, 1]], device=self.device, dtype=torch.int64 + ), + torch.tensor([1, 1], device=self.device, dtype=dtype), + torch.tensor(0.3, device=self.device, dtype=dtype), + ), + ] + equality_constraints = [ + ( + torch.tensor([0, 1, 2], device=self.device, dtype=torch.int64), + torch.tensor([1, 1, 1], device=self.device, dtype=dtype), + torch.tensor(1, device=self.device, dtype=dtype), + ) + ] + inter_point_equality_constraints = [ + ( + torch.tensor([0, 1, 2], device=self.device, dtype=torch.int64), + torch.tensor([1, 1, 1], device=self.device, dtype=dtype), + torch.tensor(1, device=self.device, dtype=dtype), + ), + ( + torch.tensor( + [[0, 0], [1, 0]], device=self.device, dtype=torch.int64 + ), + torch.tensor([1.0, -1.0], device=self.device, dtype=dtype), + 0, + ), + ] + for equalities, inequalities in product( + [None, equality_constraints, inter_point_equality_constraints], + [None, inequality_constraints, inter_point_inequality_constraints], + ): + samples = sample_q_batches_from_polytope( + n=5, + q=3, + bounds=bounds, + n_burnin=10000, + thinning=32, + seed=42, + inequality_constraints=inequalities, + equality_constraints=equalities, + ) + self.assertEqual(samples.shape, torch.Size((5, 3, 3))) + def test_gen_batch_initial_conditions_constraints(self): for dtype in (torch.float, torch.double): bounds = torch.tensor([[0, 0], [1, 1]], device=self.device, dtype=dtype) @@ -339,6 +523,77 @@ def test_gen_batch_initial_conditions_constraints(self): torch.all(batch_initial_conditions[..., idx] == val) ) + def test_gen_batch_initial_conditions_interpoint_constraints(self): + for dtype in (torch.float, torch.double): + bounds = torch.tensor([[0, 0], [1, 1]], device=self.device, dtype=dtype) + inequality_constraints = [ + ( + torch.tensor([0, 1], device=self.device, dtype=torch.int64), + torch.tensor([-1, -1.0], device=self.device, dtype=dtype), + torch.tensor(-1.0, device=self.device, dtype=dtype), + ) + ] + equality_constraints = [ + ( + torch.tensor( + [[0, 0], [1, 0]], device=self.device, dtype=torch.int64 + ), + torch.tensor([1.0, -1.0], device=self.device, dtype=dtype), + 0, + ), + ( + torch.tensor( + [[0, 0], [2, 0]], device=self.device, dtype=torch.int64 + ), + torch.tensor([1.0, -1.0], device=self.device, dtype=dtype), + 0, + ), + ] + + for nonnegative, seed in product([True, False], [None, 1234]): + mock_acqf = MockAcquisitionFunction() + with mock.patch.object( + MockAcquisitionFunction, + "__call__", + wraps=mock_acqf.__call__, + ): + batch_initial_conditions = gen_batch_initial_conditions( + acq_function=mock_acqf, + bounds=bounds, + q=3, + num_restarts=2, + raw_samples=10, + options={ + "nonnegative": nonnegative, + "eta": 0.01, + "alpha": 0.1, + "seed": seed, + "init_batch_limit": None, + "thinning": 2, + "n_burnin": 3, + }, + inequality_constraints=inequality_constraints, + equality_constraints=equality_constraints, + ) + expected_shape = torch.Size([2, 3, 2]) + self.assertEqual(batch_initial_conditions.shape, expected_shape) + self.assertEqual(batch_initial_conditions.device, bounds.device) + self.assertEqual(batch_initial_conditions.dtype, bounds.dtype) + + self.assertTrue((batch_initial_conditions.sum(dim=-1) <= 1).all()) + + self.assertAllClose( + batch_initial_conditions[0, 0, 0], + batch_initial_conditions[0, 1, 0], + batch_initial_conditions[0, 2, 0], + ) + + self.assertAllClose( + batch_initial_conditions[1, 0, 0], + batch_initial_conditions[1, 1, 0], + batch_initial_conditions[1, 2, 0], + ) + def test_error_equality_constraints_with_sample_around_best(self): tkwargs = {"device": self.device, "dtype": torch.double} # this will give something that does not respect the constraints diff --git a/test/optim/test_optimize.py b/test/optim/test_optimize.py index 2930b8a252..4b4cff3c17 100644 --- a/test/optim/test_optimize.py +++ b/test/optim/test_optimize.py @@ -379,6 +379,51 @@ def test_optimize_acqf_sequential_notimplemented(self): sequential=True, ) + def test_optimize_acqf_sequential_q_constraint_notimplemented(self): + # Sequential acquisition function not supported with q-constraints + with self.assertRaises(UnsupportedError): + optimize_acqf( + acq_function=MockAcquisitionFunction(), + bounds=torch.stack([torch.zeros(3), 4 * torch.ones(3)]), + equality_constraints=[ + ( + torch.tensor( + [[0, 0], [1, 0]], device=self.device, dtype=torch.int64 + ), + torch.tensor( + [1.0, -1.0], device=self.device, dtype=torch.float64 + ), + 0, + ), + ], + q=3, + num_restarts=2, + raw_samples=10, + return_best_only=True, + sequential=True, + ) + with self.assertRaises(UnsupportedError): + optimize_acqf( + acq_function=MockAcquisitionFunction(), + bounds=torch.stack([torch.zeros(3), 4 * torch.ones(3)]), + inequality_constraints=[ + ( + torch.tensor( + [[0, 0], [1, 0]], device=self.device, dtype=torch.int64 + ), + torch.tensor( + [1.0, -1.0], device=self.device, dtype=torch.float64 + ), + 0, + ), + ], + q=3, + num_restarts=2, + raw_samples=10, + return_best_only=True, + sequential=True, + ) + def test_optimize_acqf_batch_limit(self) -> None: num_restarts = 3 raw_samples = 5