Botorch equality_constraints argument example in the optimize_acqf function

Question

I'm doing a Bayesian Optimization code with BoTorch and I want to apply a constraint like x2+x4+x6=1 to the candidates, so I use the optimize_acqf function with this signature:

 new_x, value = optimize_acqf(
        acq_function=acq_func,
        bounds=to.tensor([[0.0] * 20, [1.0] * 20]),
        q=BATCH_SIZE,
        num_restarts=10,
        raw_samples=100,  # used for intialization heuristic
        options={"batch_limit": 5, "maxiter": 200, "nonnegative": True},
        equality_constraints=[(to.tensor([2,4,6]),to.tensor([1,1,1]),1)],
        sequential=True,
    )

But when I try to run the code, I get a RunTimeError (dot : expected both vectors to have same dtype, but found Long and Float) in the equality_constraints argument (if I comment that part out, the error disappears). I've tried setting the dtype of the tensors to float (as I reckon that the float tensor is the one from the data, and then it gives ValueError: No feasible point found. Constraint polytope appears empty. Check your constraints.)

I've searched the docs, and here it says that the signature of the argument I want is

constraints (equality) – 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

In another page referring to the same function (afaik) It says that the list of tuples should contain tensors and floats

equality_constraints (List[Tuple[Tensor, Tensor, float]] | None) – A list of tuples (indices, coefficients, rhs), with each tuple encoding an equality constraint of the form sum_i (X[indices[i]] * coefficients[i]) = rhs

From what I understand, the argument performs something like this in the case of the example given, where equality_constraints=[(to.tensor([2,4,6]),to.tensor([1,1,1]),1)]

for i in range(num_rows):
    X[i,2]+X[i,4]+X[i,6] = 1

But I haven't been able to figure out how I should define those tensors to run correctly.

Can someone give me a simple example of how to use this argument?

Answer 1

After a lot of searching, and a bunch of hours lost tracking code, I've found out that the optimize_acqf makes a call, under the hood, to gen_candidates_scipy which, in turn, when equality_constraints is given, calls the function make_scipy_linear_constraints.

In the docstring of that function (or the docs, it doesn't really matter) we can see a good example of how to use the equality_constraints argument:

Example:

The following will enforce that x[1] + 0.5 x[3] >= -0.1 for each x in both elements of the q-batch, and each of the 3 t-batches:

    >>> constraints = make_scipy_linear_constraints(
    >>>     torch.Size([3, 2, 4]),
    >>>     [(torch.tensor([1, 3]), torch.tensor([1.0, 0.5]), -0.1)],
    >>> )

The following will enforce that x[0, 1] + 0.5 x[1, 3] >= -0.1 where x[0, :] is the first element of the q-batch and x[1, :] is the second element of the q-batch, for each of the 3 t-batches:

    >>> constraints = make_scipy_linear_constraints(
    >>>     torch.size([3, 2, 4])
    >>>     [(torch.tensor([[0, 1], [1, 3]), torch.tensor([1.0, 0.5]), -0.1)],
    >>> )

So, going back to the example in the question:

From what I understand, the argument performs something like this in the case of the example given, where equality_constraints=[(to.tensor([2,4,6]),to.tensor([1,1,1]),1)]

for i in range(num_rows):
   X[i,2]+X[i,4]+X[i,6] = 1

That understanding is correct, and the function is eliminating the candidates that do not fit that constraint.