Utilizing optimize_acqf in an Ax Service API call

[Abrikosoff]

Dear Ax Team,

First of all, thanks a lot for your efforts put into Ax, and for maintaining such a lively and timely Q&A space! My questions is, generally, in regards to the implementation of a NChooseK constraint in the Ax Service API. From discussions in 769 and , it seems that setting a nonlinear inequality constraint (L0) is the way to go, and something like the following (BoTorch) implementation would work:

import torch
from botorch.acquisition import ExpectedImprovement, qExpectedImprovement
from botorch.fit import fit_gpytorch_model
from botorch.models import SingleTaskGP
from botorch.models.transforms import Standardize
from botorch.optim import optimize_acqf
from botorch.test_functions import Hartmann
from gpytorch.mlls import ExactMarginalLogLikelihood
from torch.quasirandom import SobolEngine

K =  2
def narrow_gaussian(x, ell):
    return torch.exp(-0.5 * (x / ell) ** 2)


def ineq_constraint(x, ell=1e-3 ):
    """
    Each
            callable is expected to take a `(num_restarts) x q x d`-dim tensor as an
            input and return a `(num_restarts) x q`-dim tensor with the constraint
            values.
    """
    # Approximation of || x ||_0 <= 3. The constraint is >= 0 to conform with SLSQP
    return narrow_gaussian(x, ell).sum(dim=-1) - K


def get_feasible_sobol_points(n,k):
    """Sobol sequence where we randomly set three of the parameters to zero to satisfy the constraint"""
    X = SobolEngine(dimension=6, scramble=True).draw(n).to(torch.double)
    inds = torch.argsort(torch.rand(n, 6), dim=-1)[:, :k]
    X[torch.arange(X.shape[0]).unsqueeze(-1), inds] = 0
    return X


def get_batch_initial_conditions(num_restarts, raw_samples, q, acqf):
    X = get_feasible_sobol_points(n=raw_samples*q, k=k).unsqueeze(1)
    X = X.reshape((torch.Size((raw_samples,q,6))))
    acq_vals = acqf(X)
    return X[acq_vals.topk(num_restarts).indices]


hartmann = Hartmann(dim=6)
k = 2
q = 1

X = get_feasible_sobol_points(n=10,k=K)
Y = hartmann(X).unsqueeze(-1)
print(f"Best initial point: {Y.min().item():.3f}")

gp = SingleTaskGP(X, Y, outcome_transform=Standardize(m=1))
mll = ExactMarginalLogLikelihood(gp.likelihood, gp)
fit_gpytorch_model(mll)

EI = qExpectedImprovement(model=gp, best_f=Y.min())

batch_initial_conditions = get_batch_initial_conditions(num_restarts=1, raw_samples=512, acqf=EI, q=q)

print(batch_initial_conditions)

candidate, acq_value = optimize_acqf(
    EI,
    bounds=torch.cat((torch.zeros(1, 6), torch.ones(1, 6))),
    q=q,
    nonlinear_inequality_constraints=[ineq_constraint],
    batch_initial_conditions=batch_initial_conditions,
    num_restarts=20,
    options={"batch_limit": 1, "maxiter": 200},
)

print(candidate)
print(ineq_constraint(candidate)) 

However, as far as I can understand, this is not compatible with the Service API. Can you give me some pointers on how to modify this code to fit the Service, or is there any other way of implementing the NChooseK constraint which is more Service API-friendly?

Thank you very much in advance!

Edit: Perhaps related to this issue.