Enabling candidate generation with both non_linear_constraints and fixed_features

[jduerholt]

Motivation

This PR adds the possibility to use fixed_features in combination with nonlinear_inequality_constraints as discussed in issue #1904.

@Balandat: concerning the solution that you suggested: I think it is much easier as the f_np_wrapper is also used for the nonlinear constraints:

botorch/botorch/generation/gen.py

Line 224 in 38e2027

f_np_wrapper=f_np_wrapper,

and fix_features is already applied there:

botorch/botorch/generation/gen.py

Line 189 in 38e2027

X_fix = fix_features(X, fixed_features=fixed_features)

So no need for another wrapper, or?

A wrapper function would be only needed if one wants to have the possibility to use the nonlinear constraints also with the reduced domain, then one would need a method, which takes the reduced X as input and augments it with the fixed features.

Here is some example code, that shows that it works:

import torch
from botorch.models import SingleTaskGP
from botorch.fit import fit_gpytorch_mll
from botorch.utils import standardize
from gpytorch.mlls import ExactMarginalLogLikelihood
from botorch.utils.sampling import get_polytope_samples

def constraint(X):
    return torch.sum(X[...,:3], dim=-1) -1.80

train_X = torch.rand(10, 3)
Y = 1 - torch.norm(train_X[:,:2] - 0.5, dim=-1, keepdim=True)
Y = Y + 0.1 * torch.randn_like(Y)  # add some noise
train_Y = standardize(Y)

gp = SingleTaskGP(train_X, train_Y)
mll = ExactMarginalLogLikelihood(gp.likelihood, gp)
fit_gpytorch_mll(mll)

from botorch.acquisition import UpperConfidenceBound

UCB = UpperConfidenceBound(gp, beta=0.1)

from botorch.optim import optimize_acqf


bounds = torch.stack([torch.zeros(3), torch.ones(3)])

batch_initial_conditions = get_polytope_samples(
    n=5,
    bounds=bounds,
    inequality_constraints=[(torch.tensor([0,1,2]),torch.tensor([1.,1.,1.]),1.8)],
    n_burnin=1000
).unsqueeze(-2)

candidate, acq_value = optimize_acqf(
    UCB, 
    bounds=bounds, 
    q=1, 
    num_restarts=5, 
    raw_samples=20,
    fixed_features={2:0.5}, 
    nonlinear_inequality_constraints = [constraint],
    batch_initial_conditions = batch_initial_conditions
)
print(candidate)  # tensor([0.4887, 0.5063])
print(candidate.sum())

Have you read the Contributing Guidelines on pull requests?

Yes.

Test Plan

Unit tests.

[codecov]

Codecov Report

Merging #1912 (dfa9b64) into main (38e2027) will not change coverage.
The diff coverage is 100.00%.

❗ Current head dfa9b64 differs from pull request most recent head 0e71cdb. Consider uploading reports for the commit 0e71cdb to get more accurate results

@@            Coverage Diff            @@
##              main     #1912   +/-   ##
=========================================
  Coverage   100.00%   100.00%           
=========================================
  Files          173       173           
  Lines        15232     15290   +58     
=========================================
+ Hits         15232     15290   +58     

Impacted Files	Coverage Δ
botorch/models/contextual.py	100.00% <ø> (ø)
botorch/models/pairwise_gp.py	100.00% <ø> (ø)
botorch/models/utils/assorted.py	100.00% <ø> (ø)
botorch/generation/gen.py	100.00% <100.00%> (ø)
botorch/generation/utils.py	100.00% <100.00%> (ø)
botorch/models/kernels/contextual_lcea.py	100.00% <100.00%> (ø)
botorch/models/likelihoods/pairwise.py	100.00% <100.00%> (ø)
botorch/optim/parameter_constraints.py	100.00% <100.00%> (ø)
botorch/utils/probability/utils.py	100.00% <100.00%> (ø)

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

[Balandat]

Yeah I think this makes sense. If there are multiple fixed features it's going to be slower than optimizing over the reduced domain since the optimizer needs to (i) identify the fixed features and (ii) SLSQP scales pretty poorly with the problem dimension anyway.

But that's probably fine for now (maybe add this comment in code for future reference?

tensor([0.4887, 0.5063])

This doesn't seem to be the right output? Should be 3dim with the third feature fixed to 0.5?

[Balandat]

Hmm so with this setup this test may pass even if fixing the features doesn't work (given that it only narrows this to a specific solution of the previous solution set). Can you use a different value here? I might also make sense to use a different mock acquisition function so that the optimizer is unique...

[jduerholt]

I will have a deeper look to the test. All this mocking is sometimes a bit hard for me to understand ;)

[Balandat]

I don't think there is actually any mocking going on here - it's just that mock_acq_function is a SquaredAcquisitionFunction rather than a true acquisition function.

If you don't want to change that, one thing you could do would be to just replace one of the constraints to result in a different optimum / optimizer.

[Balandat]

Test failure seems unrelated I'll take a look at that

[jduerholt]

Yeah I think this makes sense. If there are multiple fixed features it's going to be slower than optimizing over the reduced domain since the optimizer needs to (i) identify the fixed features and (ii) SLSQP scales pretty poorly with the problem dimension anyway.

I am already working on the proper solution with the reduced domain. I will give it a few more minutes and if I think it is feasible without taking too much time, I will add it to the PR.

This doesn't seem to be the right output? Should be 3dim with the third feature fixed to 0.5?

You care correct, this is copied from the botorch website minimal example

and I forgot to remove it :D

[jduerholt]

Hi @Balandat,

I implemented now also the solution using the reduced domain. Which one do you prefer?

Best,

Johannes

[Balandat]

Looks great, just a couple of minor style nits. Thanks for putting this up!

[jduerholt]

@Balandat, I incorporated your changes and added another test which was still missing (from my perspective). From my side, this is now ready for merge.

[Balandat]

Great, thanks you. I'll pull it in and merge it.

[facebook-github-bot]

@Balandat has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[SebastianAment]

Test failure seems unrelated I'll take a look at that

This was solved by #1919, should pass upon re-running the tests.

[SebastianAment]

Thanks for adding this @jduerholt! Mainly left cosmetic suggestions that will aid readability, as well as completed type signatures.

Last, could you rename the PR something a little more descriptive like "Enabling candidate generation with both non_linear_constraints and fixed_features"? This will make it easier to understand what's going on when browsing the commit history on a high level.

[SebastianAment]

Two subjective suggestions to make this quicker to parse:

1. How about _gen_nonlin_constraints_of_variable_features? Could apply same convention to the linear case below. (cc @Balandat)
2. Could you add a doc-string here, something like "Given a dictionary of fixed_features, returns a list of non-linear constraints on the variable (non-fixed) features."

[jduerholt]

I added a docsting. Concerning renaming: should I do it, or should it be done in a seperate PR for both methods (nonlinear and linear) to keep it consistent?

[SebastianAment]

Thank you! If you're up for it, you can do both in this commit. The only mentions of _generate_unfixed_lin_constraints are in:

• generation/utils.py,
• optim/parameter_constraints.py, and
• test/optim/parameter_constraints.py.

A quick search and replace in these files should do it.

[jduerholt]

Hmm, but do you really prefer the new name? What do you do not like about _generate_unfixed_lin_constraints?

[jduerholt]

But ofc, I can do it ;)

[SebastianAment]

_generate_unfixed_(non)lin_constraints makes it sound like the constraints are "unfixed", whereas it is the features that are. But let's keep it then, especially since the docstring makes it clear now. Thanks for adding it!

[facebook-github-bot]

@SebastianAment has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[facebook-github-bot]

@Balandat has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[facebook-github-bot]

@Balandat merged this pull request in 6f98187.