Using nonlinear inequality constraints on discrete variables

[creekmar]

I have a nonlinear constraint on a continuous variable based on a discrete variable. Where we are constraining the temperature based on the boiling point chosen. The boiling points are [61.2, 110.6, 132, 214.4, 180.1], and the bounds of the temperature we choose is continuous from [20, 140]. And we want the chosen temperature less than the boiling point.
I am using optimize_acqf_mixed and ran into the error

  File "C:\Users\mcree\anaconda3\envs\test\lib\site-packages\botorch\optim\parameter_constraints.py", line 590, in make_scipy_nonlinear_inequality_constraints
    raise ValueError(
ValueError: `batch_initial_conditions` must satisfy the non-linear inequality constraints.

After looking into the code, I realized that the fixed features list is forcibly changing batched_initial_conditions to the specified fixed feature. Therefore one of the candidates which originally looked like tensor([ 80.0000, 214.4000]) would be forcibly changed to tensor([ 80.0000, 61.2000]), which would not meet the constraints. Is there any way to deal with this?

[Balandat]

@creekmar could you provide a bit more details on how the fixed features come into play here? Maybe provide some more extensive code example to repro this so we can better understand the exact issue?

[creekmar]

Here is the code I use to run

# Other imports
import time
import warnings

# Botorch imports
from botorch.models import MixedSingleTaskGP
from gpytorch.mlls.exact_marginal_log_likelihood import ExactMarginalLogLikelihood
from botorch.optim import optimize_acqf_mixed
from botorch import fit_gpytorch_mll
from botorch.acquisition.monte_carlo import (
     qNoisyExpectedImprovement,
)
from botorch.exceptions import BadInitialCandidatesWarning
from botorch.sampling.normal import SobolQMCNormalSampler
from botorch.utils.sampling import draw_sobol_samples
from botorch.models.transforms import Normalize, Standardize
import torch

# ML constants
NOISE_SE = 0.5
DEVICE = torch.device("cuda:3" if torch.cuda.is_available() else "cpu")

BP = [61.2, 110.6, 132, 214.4, 180.1]
SPACE_BOUNDS = torch.tensor([(20, 61.2), (140, 180.1)], device=DEVICE)

def temp_constraint(params: torch.Tensor) -> torch.Tensor:
    return params[1]-params[0]

def dummy_measure(params):
    heater, bp = params
    return pow(bp-heater, 2)

def take_dummy_measures(data):
    result = torch.zeros(len(data))
    for i in range(len(data)):
        result[i] = dummy_measure(data[i])
    return result

def generate_init_data(n=10, maximize = True):
    # need to add comma so it recognizes it as a tuple and not a single listing
    b = torch.tensor([(0.0,), (1.0,)], device=DEVICE)
    raw_samples = draw_sobol_samples(bounds = b, n=1, q=n)
    train_X = raw_samples.flatten(0,1)

    # add the bp
    count = 0
    blank = torch.zeros(n, 1)
    while count < n:
        bp = BP[count%5]
        blank[count, 0] = bp
        if bp > 140:
            high = 140
        else: 
            high = bp

        # if we want random temps in the bounds
        train_X[count, 0] = train_X[count, 0].mul(high-20).add(20)

        count += 1

    train_X = torch.cat((train_X, blank), 1)
    # TODO replace with actual measure from experiment
    exact_obj = take_dummy_measures(train_X).unsqueeze(-1)  # add output dimension
    train_Y = exact_obj + NOISE_SE * torch.randn_like(exact_obj)

    if maximize:
        best_observed_value = train_Y.max().item()
    else:
        best_observed_value = train_Y.min().item()
    return train_X, train_Y, best_observed_value

def gen_batch():
    # Note: the shape of batch should be len(shape) = 3 and shape[0] = 1
    # having 5 batch initial points that are in the midpoint of the temp range
    n = 5
    # should have comma to differentiate tuple from single item
    b = torch.tensor([(0.0,), (1.0,)], device=DEVICE)
    raw_samples = draw_sobol_samples(bounds = b, n=1, q=n)
    train_X = raw_samples

    # add the bp
    count = 0
    blank = torch.zeros(1,n, 1)
    while count < n:
        bp = BP[count%5] 
        blank[0, count, 0] = bp
        # temperatures are based on the midpoint of temp bounds
        if bp > 140:
            high = 140
        else: 
            high = bp
        # use the midpoint of the temp range
        train_X[0,count, 0] = (high+20)/2
        count += 1
    train_X = torch.cat((train_X, blank), 2)
    return train_X
    
def initialize_model(train_x, train_obj, state_dict=None):
    # combine into a multi-output GP model
    model = MixedSingleTaskGP(train_x, train_obj, [1], outcome_transform=Standardize(m=1), input_transform=Normalize(d=2))
    mll = ExactMarginalLogLikelihood(model.likelihood, model)
    # load state dict if it is passed
    if state_dict is not None:
        model.load_state_dict(state_dict)
    return mll, model

def optimize_acqf_and_get_observation(acq_func, batch_size=1, num_restarts=10, raw_samples=64):
    candidates, _ = optimize_acqf_mixed(
        acq_function=acq_func,
        bounds=SPACE_BOUNDS,
        q=batch_size,
        num_restarts=num_restarts,
        raw_samples=raw_samples,  # used for intialization heuristic
        fixed_features_list=[{1: 61.2}, {1: 110.6}, {1: 132}, {1: 214.4}, {1: 180.1}],
        # TODO figure this out
        # True means intrapoint which applies constraint to each individual point in a batch
        nonlinear_inequality_constraints=[(temp_constraint,True)],
        batch_initial_conditions=gen_batch(),
        # batch limit must be 1
        options={"batch_limit": 1, "maxiter": 200}
        )

    # observe new values
    new_x = candidates.detach()
    print("\tNew Candidate:", new_x)
    
    # TODO Replace with actual measurement
    exact_obj = take_dummy_measures(new_x).unsqueeze(-1)  # add output dimension
    new_obj = exact_obj + NOISE_SE * torch.randn_like(exact_obj)

    return new_x, new_obj

def run_optimization(n_start_runs=10, n_loop_runs=20, mc_samples=256, maximize = True, verbose=True,
                     state_dict_file='model_state_dict2.pth'):
    warnings.filterwarnings("ignore", category=BadInitialCandidatesWarning)
    warnings.filterwarnings("ignore", category=RuntimeWarning)

    best_observed_nei  = []
    time_list = torch.zeros(n_loop_runs+n_start_runs, 1)
    # call helper functions to generate initial training data and initialize model
    (
        train_x_nei,
        train_obj_nei,
        best_observed_value_nei,
    ) = generate_init_data(n_start_runs, maximize)
    mll_nei, model_nei = initialize_model(train_x_nei, train_obj_nei)

    best_observed_nei.append(best_observed_value_nei)

    # run N_RUNS runs of BayesOpt after the initial random data
    for iteration in range(1, 21):

        t0 = time.monotonic()

        # fit the models
        fit_gpytorch_mll(mll_nei)

        # define the qEI and qNEI acquisition modules using a QMC sampler
        qmc_sampler = SobolQMCNormalSampler(sample_shape=torch.Size([mc_samples]))

        # for best_f, we use the best observed noisy values as an approximation
        qNEI = qNoisyExpectedImprovement(
            model=model_nei,
            X_baseline=train_x_nei,
            sampler=qmc_sampler,
        )

        # optimize and get new observation
        new_x_nei, new_obj_nei = optimize_acqf_and_get_observation(qNEI)

        # update training points
        train_x_nei = torch.cat([train_x_nei, new_x_nei])
        train_obj_nei = torch.cat([train_obj_nei, new_obj_nei])

        # update progress
        if maximize:
            best_value_nei = train_obj_nei.max().item()
            best_x_nei = train_x_nei[torch.argmax(train_obj_nei).item()]
        else:
            best_value_nei = train_obj_nei.min().item()
            best_x_nei = train_x_nei[torch.argmin(train_obj_nei).item()]
        best_observed_nei.append(best_value_nei)


        # reinitialize the models so they are ready for fitting on next iteration
        # use the current state dict to speed up fitting
        mll_nei, model_nei = initialize_model(
            train_x_nei,
            train_obj_nei,
            model_nei.state_dict(),
        )

        t1 = time.monotonic()
        time_list[iteration-1+n_start_runs] = t1-t0

        if verbose:
            print(
                f"\nRUN {iteration:>2}: \n\t"
                f"qNEI Best Value: {best_value_nei:>4.2f} at {best_x_nei.tolist()}\n\t"
                f"Time = {t1-t0:>4.2f}.\n",
                end="",
            )
        else:
            print(".", end="")

    # save data and plot optimization trace
    torch.save(model_nei.state_dict(), state_dict_file)

if __name__ == "__main__":
    run_optimization(maximize=False)

[Balandat]

I believe the issue is that in gen_batch() you're generating initial conditions that aren't necessarily compatible with all boiling points. Since optimize_acqf_mixed for q=1 is really just a simple loop around optimize_acqf, the easiest approach here woudl be to just call optimize_acqf in a loop where the batch_initial_conditions are feasible w.r.t. to the currently chosen boiling point:

botorch/botorch/optim/optimize.py

Lines 927 to 952 in c895a8d

	if q == 1:
	ff_candidate_list, ff_acq_value_list = [], []
	for fixed_features in fixed_features_list:
	candidate, acq_value = optimize_acqf(
	acq_function=acq_function,
	bounds=bounds,
	q=q,
	num_restarts=num_restarts,
	raw_samples=raw_samples,
	options=options or {},
	inequality_constraints=inequality_constraints,
	equality_constraints=equality_constraints,
	nonlinear_inequality_constraints=nonlinear_inequality_constraints,
	fixed_features=fixed_features,
	post_processing_func=post_processing_func,
	batch_initial_conditions=batch_initial_conditions,
	ic_generator=ic_generator,
	return_best_only=True,
	**ic_gen_kwargs,
	)
	ff_candidate_list.append(candidate)
	ff_acq_value_list.append(acq_value)
	ff_acq_values = torch.stack(ff_acq_value_list)
	best = torch.argmax(ff_acq_values)
	return ff_candidate_list[best], ff_acq_values[best]