Non-equidistant discrete dimensions for `optimize_acqf_mixed_alternating` #2923

jduerholt · 2025-07-11T08:58:45Z

Motivation

As dicussed in #2904, this is the first PR in planned series of PRs regarding adding additional functionality to optimize_acqf_mixed_alternating. In this one, non-equidistant discrete dimensions are introduced.

CC: @TobyBoyne

Have you read the Contributing Guidelines on pull requests?

Yes.

Test Plan

Unit tests.

codecov · 2025-07-11T09:31:05Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 100.00%. Comparing base (0eea0b7) to head (a1a4235).
Report is 10 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff            @@
##              main     #2923   +/-   ##
=========================================
  Coverage   100.00%   100.00%           
=========================================
  Files          212       212           
  Lines        19819     19852   +33     
=========================================
+ Hits         19819     19852   +33

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

Balandat

I didn't go through in full detail; at a high level this makes sense, the main thing I'm worried is that if there are many dimensions and many values this could end up becoming quite slow. Do you have any timing / profiling data on this?

botorch/optim/optimize_mixed.py

jduerholt · 2025-07-11T14:10:49Z

the main thing I'm worried is that if there are many dimensions

In principle this should be of the same speed as the original implementation, main difference is the rounding operation, are you interested in profiling of this operation? Or do you see issues with speed also in other parts of the code?

jduerholt · 2025-07-11T14:13:59Z

Ah, now I see, you see the main problem in get_nearest_neighbors, or? Are interested in profiling data on this? I can generate it, but have not looked into it, as I wanted to first get things started and working again. I will have a look into this.

jduerholt · 2025-07-11T14:40:43Z

Here is some code for profiling for the mentioned operation:

from botorch.optim.optimize_mixed import get_nearest_neighbors
import time
import torch

d = 100
nd = 90
n_discrete_values = 300

current_x = torch.rand(d)

discrete_dims = {i: list(range(n_discrete_values+1)) for i in range(nd)}

for i in range(nd):
    current_x[i] = torch.randint(0, n_discrete_values, (1,)).item()

bounds = torch.tensor([[0.0] * d, [n_discrete_values] * d])

start = time.time()
neighbors = get_nearest_neighbors(current_x=current_x, discrete_dims=discrete_dims, bounds=bounds)
end = time.time()
print(f"Time taken to find nearest neighbors: {end - start:.4f} seconds")

For this rather extreme case, it takes around 0.006 seconds on my macbook with a M1 processor. The original one takes for this case around 0.001 seconds, the new one is defintely slower, but I think that it will be not the rate limiting step (speeking as a chemist :D), or?

Balandat · 2025-07-11T15:04:03Z

Yeah, that seems like it should be fine.

jduerholt · 2025-07-11T21:10:11Z

Hi @Balandat,

I implemented your comments, so it should be ready for rereview.

Best,

Johannes

Balandat

Generally this lgtm - @saitcakmak could you take a look as well since this changes some things deep down in some code that I believe you wrote?

botorch/optim/optimize_mixed.py

saitcakmak

This seems reasonable overall. My main concern is with clamping operations, which may be introducing a bug. Curious to hear your take on extending this to categoricals as well

saitcakmak · 2025-07-17T15:27:38Z

botorch/optim/optimize_mixed.py

+        t_values = torch.tensor(
+            values, device=unique_neighbors.device, dtype=unique_neighbors.dtype
+        )
+        idx = unique_neighbors[:, dim].long().clamp(0, len(values) - 1)


Why do we need to clamp here? Both plus_neighbors and minus_neighbors have been clamped, right? If we actually need to clamp here, I think that may produce duplicate neighbors.

I will investigate this, after fixing the clamping from above.

saitcakmak · 2025-07-17T15:29:20Z

botorch/optim/optimize_mixed.py

-    plus_neighbors = current_x.repeat(num_discrete, 1)
-    plus_neighbors[:, discrete_dims] += diag_ones
+    plus_neighbors = current_x_int.repeat(num_discrete, 1)
+    plus_neighbors[:, t_discrete_dims] += diag_ones
    plus_neighbors.clamp_(max=bounds[1])


Does it make sense to clamp to bounds anymore, since the neighbors now have index values that are between 0 & len(values) - 1?

Good catch, we need to clamp, but not with the bounds ;) I overlooked this. I will change it.

@jduerholt this still needs to be changed it seems

saitcakmak · 2025-07-17T15:30:37Z

botorch/optim/optimize_mixed.py

-    minus_neighbors = current_x.repeat(num_discrete, 1)
-    minus_neighbors[:, discrete_dims] -= diag_ones
+    minus_neighbors = current_x_int.repeat(num_discrete, 1)
+    minus_neighbors[:, t_discrete_dims] -= diag_ones
    minus_neighbors.clamp_(min=bounds[0])


Same question about clamping here. I think this one is particularly problematic since the bounds[0] may be some large number even though the index value has a fixed range. I guess the same argument would apply to bounds[1] as well since the bounds can be negative as well.

Totally correct, we need to clamp with zero here.

botorch/optim/optimize_mixed.py

saitcakmak · 2025-07-17T15:42:55Z

botorch/optim/optimize_mixed.py

+        discrete_dims: A dictionary mapping indices of discrete and binary
+            dimensions to a list of allowed values for that dimension.
        cat_dims: A list of indices corresponding to categorical parameters.


So, the categoricals are still required to be [0, 1, 2, ...] but discrete dimensions are relaxed. I think it'd be good to treat them equally here if it's not too much extra work.

Yeah, this PR is only dealing with the discrete ones. My plan is to file a consecutive PR in which I introduce the same syntax for categoricals. My plan was to get first the feedback on the discrete ones to have then less work with the categoricals ;) Do you want to have the categoricals also in this PR?

I think it's fine to work incrementally and deal with the categoricals in a separate PR, so long as this PR clearly states hat limitation in the docstring and errors our if the input doesn't comply with that requirement.

botorch/optim/optimize_mixed.py

Balandat

Looks like a couple of items are still open, o/w this is getting there!

Balandat · 2025-07-18T13:46:37Z

botorch/optim/optimize_mixed.py

-    plus_neighbors = current_x.repeat(num_discrete, 1)
-    plus_neighbors[:, discrete_dims] += diag_ones
+    plus_neighbors = current_x_int.repeat(num_discrete, 1)
+    plus_neighbors[:, t_discrete_dims] += diag_ones
    plus_neighbors.clamp_(max=bounds[1])


@jduerholt this still needs to be changed it seems

Balandat · 2025-07-18T13:46:49Z

botorch/optim/optimize_mixed.py

-    minus_neighbors = current_x.repeat(num_discrete, 1)
-    minus_neighbors[:, discrete_dims] -= diag_ones
+    minus_neighbors = current_x_int.repeat(num_discrete, 1)
+    minus_neighbors[:, t_discrete_dims] -= diag_ones
    minus_neighbors.clamp_(min=bounds[0])


jduerholt added 2 commits July 11, 2025 10:46

tests are passing

be62f01

remove unnecesary sorted

1a08b24

facebook-github-bot added the CLA Signed Do not delete this pull request or issue due to inactivity. label Jul 11, 2025

jduerholt added 2 commits July 11, 2025 11:11

fix failing test

ed8f878

update docstrings

307cc31

Balandat reviewed Jul 11, 2025

View reviewed changes

botorch/optim/optimize_mixed.py Outdated Show resolved Hide resolved

botorch/optim/optimize_mixed.py Show resolved Hide resolved

botorch/optim/optimize_mixed.py Outdated Show resolved Hide resolved

jduerholt added 2 commits July 11, 2025 22:30

improve continuous relaxation

1230c4b

implement max suggestion

9799141

Balandat reviewed Jul 16, 2025

View reviewed changes

botorch/optim/optimize_mixed.py Outdated Show resolved Hide resolved

botorch/optim/optimize_mixed.py Outdated Show resolved Hide resolved

botorch/optim/optimize_mixed.py Outdated Show resolved Hide resolved

implement suggestions by Max

a62caad

saitcakmak requested changes Jul 17, 2025

View reviewed changes

remove sorting operations

a1a4235

Balandat reviewed Jul 18, 2025

View reviewed changes

Non-equidistant discrete dimensions for optimize_acqf_mixed_alternating #2923

Are you sure you want to change the base?

Non-equidistant discrete dimensions for optimize_acqf_mixed_alternating #2923

Conversation

jduerholt commented Jul 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Have you read the Contributing Guidelines on pull requests?

Test Plan

Uh oh!

codecov bot commented Jul 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Balandat left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jduerholt commented Jul 11, 2025

Uh oh!

jduerholt commented Jul 11, 2025

Uh oh!

jduerholt commented Jul 11, 2025

Uh oh!

Balandat commented Jul 11, 2025

Uh oh!

jduerholt commented Jul 11, 2025

Uh oh!

Balandat left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

saitcakmak left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Balandat Jul 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Balandat left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Non-equidistant discrete dimensions for `optimize_acqf_mixed_alternating` #2923

Non-equidistant discrete dimensions for `optimize_acqf_mixed_alternating` #2923

jduerholt commented Jul 11, 2025 •

edited

Loading

codecov bot commented Jul 11, 2025 •

edited

Loading

Balandat Jul 18, 2025 •

edited

Loading