[Bug] get_polytope_samples fails for inequalities over q-batches and the actual dimension

[AVHopp]

🐛 Bug

When using linear inequalities across multiple points resp. different q-batches, a broadcasting error within sparse_to_dense_constraints occurs when defining a constraint.

This error only happens in a "mixed" setting when defining the constraint over both the q-batch dimension and the point dimension (see "Additional context").

Thanks already for any help :)

To reproduce

import torch
from botorch.utils.sampling import get_polytope_samples

bounds = torch.tensor([[0.0, 0.0], [1.0, 1.0]])
q = 2

# Modeling x[0,0] + x[0,1] + x[1,1] >= 1 where 1st index is q-batch and 2nd is dimension
bad_indices = torch.tensor([[0.0, 0.0], [0.0, 1.0], [1.0, 1.0]])
bad_inequality_constraints = [(bad_indices, torch.ones(3), 1)]

# This fails due to a broadcasting error in sparse_to_dense_constraints
get_polytope_samples(
    n=q, bounds=bounds, inequality_constraints=bad_inequality_constraints
)

** Stack trace/error message **

Traceback (most recent call last):
  File "/Users/M290244/Programming/BayBE/test.py", line 29, in <module>
    get_polytope_samples(
  File "/Users/M290244/miniforge3/envs/baybe310/lib/python3.10/site-packages/botorch/utils/sampling.py", line 918, in get_polytope_samples
    A, b = sparse_to_dense_constraints(
  File "/Users/M290244/miniforge3/envs/baybe310/lib/python3.10/site-packages/botorch/utils/sampling.py", line 970, in sparse_to_dense_constraints
    A[i, indices.long()] = coefficients
RuntimeError: shape mismatch: value tensor of shape [3] cannot be broadcast to indexing result of shape [3, 2]

Expected Behavior

I would have expected this code to calculate a sample of points x of shape (2,2) fulfilling x[0,0] + x[0,1] + x[1,0] >= 1, as imposing such a constraint is supported across both individual axes.

System information

Please complete the following information:

• BoTorch version 0.11.1
• Mac OS

Additional context

Defining the inequalities either only along the q-batch or the "point" dimension works as expected, that is, the following code behaves as expected:

import torch
from botorch.utils.sampling import get_polytope_samples

bounds = torch.tensor([[0.0, 0.0], [1.0, 1.0]])
q = 2

# Modeling x[0,0] + x[0,1]>= 0.5
good_indices_1 = torch.tensor([[0.0, 0.0], [0.0, 1.0]])
good_inequality_constraints_1 = [(good_indices_1, torch.ones(2), 1)]

# Modeling x[0,0] + x[1,0]>= 1
good_indices_2 = torch.tensor([[0.0, 0.0], [1.0, 0.0]])
good_inequality_constraints_2 = [(good_indices_2, torch.ones(2), 1)]

# These test do not fail
sample_1 = get_polytope_samples(
    n=q, bounds=bounds, inequality_constraints=good_inequality_constraints_1
)
assert sample_1[0, 0] + sample_1[0, 1] >= 1
sample_2 = get_polytope_samples(
    n=q, bounds=bounds, inequality_constraints=good_inequality_constraints_2
)
assert sample_2[0, 0] + sample_2[1, 0] >= 1

[AVHopp]

An additional note: These kind of constraints are supported by optimize_acqf, which is why I would expect them to also work for get_polytope_sampling.

[Balandat]

There is a misunderstanding here - the n argument of get_polytope_samples means that you want to draw n i.i.d (ideally, if the chain has mixed properly) samples from the domain. Therefore, "inter-point constraints" over q points as they are supported in optimize_acqf cannot be achieved by passing n=q to get_polytope_samples - these joint constraints cannot be expressed in this way.

Instead, what you will need to do is augment your domain - instead of drawing samples in [0, 1]^2, you'll need to "stack" the q=2 points and sample in [0,1]^4 while applying the constraints appropriately (a similar construction happens in optimize_acqf where the optimization also is performed over this joint domain):

bounds = torch.tensor([[0.0, 0.0], [1.0, 1.0]])
q = 2
n = 3

# Modeling x[0,0] + x[0,1] + x[1,1] >= 1 where 1st index is q-batch and 2nd is dimension

# We need to express this problem as a sampling problem on the joint domain [0,1]^4 and apply the constraints on that.
bounds_joint = torch.cat([bounds for _ in range(q)], dim=-1)
indices_joint = torch.tensor([[0, 1, 3]])
inequality_constraints_joint = [(indices_joint, torch.ones(3), 1)]

sample_joint = get_polytope_samples(
    n=n, bounds=bounds_joint, inequality_constraints=inequality_constraints_joint
)
for s in sample_joint: 
    assert s[0] + s[1] + s[3] >= 1

The constraint x[0,0] + x[1,0]>= 1 only seemingly works - you were (un?)lucky - if you run a bunch of replications eventually that assertion will fail.

sparse_to_dense_constraints was never meant to operate on multi-dimensional indices. The underlying HitAndRunPolytopeSampler that is called in get_polytope_samples also does not support this kind of constraint. So I think what I'll do is put up a PR to make that more clear in the docstring and raise an error if the inputs do not conform to that assumption.

Aside: If you want to draw multiple samples while using some of them before drawing new ones, yo'll want to consider using HitAndRunPolytopeSampler (see the docstring here: https://github.com/pytorch/botorch/blob/main/botorch/utils/sampling.py#L887-L895)

An additional note: These kind of constraints are supported by optimize_acqf, which is why I would expect them to also work for get_polytope_sampling.

This is a reasonable point from an API perspective - in principle one could build an interface that based on the dimension (and non-zero locations) of the constraint indices passed in does the domain-augmentation automatically under the hood. However, that would still leave some ambiguity in terms of the return type of the function - if d=3 and the constraints indicate that q=2 and you call this with n=4, what is the shape of the output? n x (qd)? n x q x d? (nq x d)? For that reason I don't think it's high enough RoI to enable that functionality out-of-the-box (but if you wanted to take a stab at coming up with a design & PR you're welcome to).

[AVHopp]

Thanks you very much for your detailed answer 😄 I did some more testing and agree, it seems like I have just been (un-)lucky with drawing samples in my first round of testing. Also, thanks for the hint towards the HitAndRunPolytopeSampler, I will have a look at that.
I also see the issue with the ambiguity regarding a potential implementation of this, but I do not think that I will be able to come up with a design/PR for this. In any case, form my side this issue can be closed as all my questions have been answered, thanks again :)