Constrain optimization over unfixed features if fixed_features is passed. (#839)

Summary:
Pull Request resolved: #839

This diff implements a change previously discussed in https://www.internalfb.com/diff/D28869203?dst_version_fbid=180392114012471&transaction_fbid=280143027144841. This is a sensible thing to do and will potentially save a lot of compute.

Reviewed By: Balandat

Differential Revision: D29265778

fbshipit-source-id: e8018a45e17394c865f414df44f9ceafaf61f0b8

Author: vishwakftw

botorch/generation/gen.py

@@ -10,10 +10,12 @@
 
 from __future__ import annotations
 
-from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union
+from typing import Any, Dict, List, Optional, Tuple, Type, Union
 
 import numpy as np
 import torch
+from botorch.acquisition import AcquisitionFunction
+from botorch.generation.utils import _remove_fixed_features_from_optimization
 from botorch.optim.parameter_constraints import (
     _arrayify,
     make_scipy_bounds,
@@ -23,13 +25,12 @@
 from botorch.optim.utils import _filter_kwargs, columnwise_clamp, fix_features
 from scipy.optimize import minimize
 from torch import Tensor
-from torch.nn import Module
 from torch.optim import Optimizer
 
 
 def gen_candidates_scipy(
     initial_conditions: Tensor,
-    acquisition_function: Module,
+    acquisition_function: AcquisitionFunction,
     lower_bounds: Optional[Union[float, Tensor]] = None,
     upper_bounds: Optional[Union[float, Tensor]] = None,
     inequality_constraints: Optional[List[Tuple[Tensor, Tensor, float]]] = None,
@@ -83,9 +84,45 @@ def gen_candidates_scipy(
             )
     """
     options = options or {}
+
+    # REDUCED is used indicate if we are optimizing over a reduced domain dimension
+    # after considering fixed_features.
+    # REDUCED mode if fixed_features is not None except for when fixed_features.values()
+    # contains None and linear constraints are passed.
+    REDUCED = fixed_features is not None
+    if inequality_constraints or equality_constraints:
+        REDUCED = REDUCED and (None not in fixed_features.values())
+
+    if REDUCED:
+        _no_fixed_features = _remove_fixed_features_from_optimization(
+            fixed_features=fixed_features,
+            acquisition_function=acquisition_function,
+            initial_conditions=initial_conditions,
+            lower_bounds=lower_bounds,
+            upper_bounds=upper_bounds,
+            inequality_constraints=inequality_constraints,
+            equality_constraints=equality_constraints,
+        )
+
+        # call the routine with no fixed_features
+        clamped_candidates, batch_acquisition = gen_candidates_scipy(
+            initial_conditions=_no_fixed_features.initial_conditions,
+            acquisition_function=_no_fixed_features.acquisition_function,
+            lower_bounds=_no_fixed_features.lower_bounds,
+            upper_bounds=_no_fixed_features.upper_bounds,
+            inequality_constraints=_no_fixed_features.inequality_constraints,
+            equality_constraints=_no_fixed_features.equality_constraints,
+            options=options,
+            fixed_features=None,
+        )
+        clamped_candidates = _no_fixed_features.acquisition_function._construct_X_full(
+            clamped_candidates
+        )
+        return clamped_candidates, batch_acquisition
+
     clamped_candidates = columnwise_clamp(
         X=initial_conditions, lower=lower_bounds, upper=upper_bounds
-    ).requires_grad_(True)
+    )
 
     shapeX = clamped_candidates.shape
     x0 = _arrayify(clamped_candidates.view(-1))
@@ -111,7 +148,7 @@ def f(x):
             .contiguous()
             .requires_grad_(True)
         )
-        X_fix = fix_features(X=X, fixed_features=fixed_features)
+        X_fix = fix_features(X, fixed_features=fixed_features)
         loss = -acquisition_function(X_fix).sum()
         # compute gradient w.r.t. the inputs (does not accumulate in leaves)
         gradf = _arrayify(torch.autograd.grad(loss, X)[0].contiguous().view(-1))
@@ -137,20 +174,22 @@ def f(x):
         options={k: v for k, v in options.items() if k not in ["method", "callback"]},
     )
     candidates = fix_features(
-        X=torch.from_numpy(res.x).to(initial_conditions).view(shapeX).contiguous(),
+        X=torch.from_numpy(res.x).to(initial_conditions).reshape(shapeX),
         fixed_features=fixed_features,
     )
+
     clamped_candidates = columnwise_clamp(
         X=candidates, lower=lower_bounds, upper=upper_bounds, raise_on_violation=True
     )
     with torch.no_grad():
         batch_acquisition = acquisition_function(clamped_candidates)
+
     return clamped_candidates, batch_acquisition
 
 
 def gen_candidates_torch(
     initial_conditions: Tensor,
-    acquisition_function: Callable,
+    acquisition_function: AcquisitionFunction,
     lower_bounds: Optional[Union[float, Tensor]] = None,
     upper_bounds: Optional[Union[float, Tensor]] = None,
     optimizer: Type[Optimizer] = torch.optim.Adam,
@@ -199,10 +238,41 @@ def gen_candidates_torch(
             )
     """
     options = options or {}
+
+    # REDUCED is used indicate if we are optimizing over a reduced domain dimension
+    # after considering fixed_features.
+    REDUCED = fixed_features is not None
+
+    if REDUCED:
+        _no_fixed_features = _remove_fixed_features_from_optimization(
+            fixed_features=fixed_features,
+            acquisition_function=acquisition_function,
+            initial_conditions=initial_conditions,
+            lower_bounds=lower_bounds,
+            upper_bounds=upper_bounds,
+            inequality_constraints=None,
+            equality_constraints=None,
+        )
+
+        # call the routine with no fixed_features
+        clamped_candidates, batch_acquisition = gen_candidates_torch(
+            initial_conditions=_no_fixed_features.initial_conditions,
+            acquisition_function=_no_fixed_features.acquisition_function,
+            lower_bounds=_no_fixed_features.lower_bounds,
+            upper_bounds=_no_fixed_features.upper_bounds,
+            optimizer=optimizer,
+            options=options,
+            verbose=verbose,
+            fixed_features=None,
+        )
+        clamped_candidates = _no_fixed_features.acquisition_function._construct_X_full(
+            clamped_candidates
+        )
+        return clamped_candidates, batch_acquisition
+
     clamped_candidates = columnwise_clamp(
         X=initial_conditions, lower=lower_bounds, upper=upper_bounds
     ).requires_grad_(True)
-    candidates = fix_features(clamped_candidates, fixed_features)
     bayes_optimizer = optimizer(
         params=[clamped_candidates], lr=options.get("lr", 0.025)
     )
@@ -215,29 +285,33 @@ def gen_candidates_torch(
     )
     while not stop:
         i += 1
-        loss = -acquisition_function(candidates).sum()
+        loss = -acquisition_function(clamped_candidates).sum()
         if verbose:
             print("Iter: {} - Value: {:.3f}".format(i, -(loss.item())))
         loss_trajectory.append(loss.item())
-        param_trajectory["candidates"].append(candidates.clone())
+        param_trajectory["candidates"].append(clamped_candidates.clone())
 
         def closure():
             bayes_optimizer.zero_grad()
             loss.backward()
             return loss
 
         bayes_optimizer.step(closure)
-        clamped_candidates.data = columnwise_clamp(
-            clamped_candidates, lower_bounds, upper_bounds
-        )
-        candidates = fix_features(clamped_candidates, fixed_features)
+        with torch.no_grad():
+            clamped_candidates = columnwise_clamp(
+                X=clamped_candidates, lower=lower_bounds, upper=upper_bounds
+            )
         stop = stopping_criterion.evaluate(fvals=loss.detach())
     clamped_candidates = columnwise_clamp(
-        X=candidates, lower=lower_bounds, upper=upper_bounds, raise_on_violation=True
+        X=clamped_candidates,
+        lower=lower_bounds,
+        upper=upper_bounds,
+        raise_on_violation=True,
     )
     with torch.no_grad():
-        batch_acquisition = acquisition_function(candidates)
-    return candidates, batch_acquisition
+        batch_acquisition = acquisition_function(clamped_candidates)
+
+    return clamped_candidates, batch_acquisition
 
 
 def get_best_candidates(batch_candidates: Tensor, batch_values: Tensor) -> Tensor:

botorch/generation/utils.py

@@ -6,9 +6,12 @@
 
 from __future__ import annotations
 
-import typing  # noqa F401
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Tuple, Union
 
 import torch
+from botorch.acquisition import AcquisitionFunction, FixedFeatureAcquisitionFunction
+from botorch.optim.parameter_constraints import _generate_unfixed_lin_constraints
 from torch import Tensor
 
 
@@ -44,3 +47,104 @@ def _flip_sub_unique(x: Tensor, k: int) -> Tensor:
         if len(out) >= k:
             break
     return x[idcs[: len(out)]]
+
+
+@dataclass(frozen=True, repr=False, eq=False)
+class _NoFixedFeatures:
+    """
+    Dataclass to store the objects after removing fixed features.
+    Objects here refer to the acquisition function, initial conditions,
+    bounds and parameter constraints.
+    """
+
+    acquisition_function: FixedFeatureAcquisitionFunction
+    initial_conditions: Tensor
+    lower_bounds: Optional[Union[float, Tensor]]
+    upper_bounds: Optional[Union[float, Tensor]]
+    inequality_constraints: Optional[List[Tuple[Tensor, Tensor, float]]]
+    equality_constraints: Optional[List[Tuple[Tensor, Tensor, float]]]
+
+
+def _remove_fixed_features_from_optimization(
+    fixed_features: Dict[int, Optional[float]],
+    acquisition_function: AcquisitionFunction,
+    initial_conditions: Tensor,
+    lower_bounds: Optional[Union[float, Tensor]],
+    upper_bounds: Optional[Union[float, Tensor]],
+    inequality_constraints: Optional[List[Tuple[Tensor, Tensor, float]]],
+    equality_constraints: Optional[List[Tuple[Tensor, Tensor, float]]],
+) -> _NoFixedFeatures:
+    """
+    Given a set of non-empty fixed features, this function effectively reduces the
+    dimensionality of the domain that the acquisition function is being optimized
+    over by removing the set of fixed features. Consequently, this function returns a
+    new `FixedFeatureAcquisitionFunction`, new constraints, and bounds defined over
+    unfixed features.
+
+    Args:
+        fixed_features: This is a dictionary of feature indices to values, where
+            all generated candidates will have features fixed to these values.
+            If the dictionary value is None, then that feature will just be
+            fixed to the clamped value and not optimized. Assumes values to be
+            compatible with lower_bounds and upper_bounds!
+        acquisition_function: Acquisition function over the original domain being
+            maximized.
+        initial_conditions: Starting points for optimization w.r.t. the complete domain.
+        lower_bounds: Minimum values for each column of initial_conditions.
+        upper_bounds: Minimum values for each column of initial_conditions.
+        inequality constraints: A list of tuples (indices, coefficients, rhs),
+            with each tuple encoding an inequality constraint of the form
+            `sum_i (X[indices[i]] * coefficients[i]) >= rhs`.
+        equality constraints: A list of tuples (indices, coefficients, rhs),
+            with each tuple encoding an inequality constraint of the form
+            `sum_i (X[indices[i]] * coefficients[i]) = rhs`.
+
+    Returns:
+        _NoFixedFeatures dataclass object.
+    """
+    # sort the keys for consistency
+    sorted_keys = sorted(fixed_features)
+    sorted_values = []
+    for key in sorted_keys:
+        if fixed_features[key] is None:
+            val = initial_conditions[..., [key]]
+        else:
+            val = fixed_features[key]
+        sorted_values.append(val)
+
+    d = initial_conditions.shape[-1]
+    acquisition_function = FixedFeatureAcquisitionFunction(
+        acq_function=acquisition_function,
+        d=d,
+        columns=sorted_keys,
+        values=sorted_values,
+    )
+
+    # extract initial_conditions, bounds at unfixed indices
+    unfixed_indices = sorted(set(range(d)) - set(sorted_keys))
+    initial_conditions = initial_conditions[..., unfixed_indices]
+    if isinstance(lower_bounds, Tensor):
+        lower_bounds = lower_bounds[..., unfixed_indices]
+    if isinstance(upper_bounds, Tensor):
+        upper_bounds = upper_bounds[..., unfixed_indices]
+
+    inequality_constraints = _generate_unfixed_lin_constraints(
+        constraints=inequality_constraints,
+        fixed_features=fixed_features,
+        dimension=d,
+        eq=False,
+    )
+    equality_constraints = _generate_unfixed_lin_constraints(
+        constraints=equality_constraints,
+        fixed_features=fixed_features,
+        dimension=d,
+        eq=True,
+    )
+    return _NoFixedFeatures(
+        acquisition_function=acquisition_function,
+        initial_conditions=initial_conditions,
+        lower_bounds=lower_bounds,
+        upper_bounds=upper_bounds,
+        inequality_constraints=inequality_constraints,
+        equality_constraints=equality_constraints,
+    )

botorch/optim/parameter_constraints.py

@@ -15,7 +15,7 @@
 
 import numpy as np
 import torch
-from botorch.exceptions.errors import UnsupportedError
+from botorch.exceptions.errors import UnsupportedError, CandidateGenerationError
 from scipy.optimize import Bounds
 from torch import Tensor
 
@@ -267,3 +267,66 @@ def _make_linear_constraints(
     else:
         raise ValueError("`indices` must be at least one-dimensional")
     return constraints
+
+
+def _generate_unfixed_lin_constraints(
+    constraints: Optional[List[Tuple[Tensor, Tensor, float]]],
+    fixed_features: Dict[int, float],
+    dimension: int,
+    eq: bool,
+) -> Optional[List[Tuple[Tensor, Tensor, float]]]:
+
+    # If constraints is None or an empty list, then return itself
+    if not constraints:
+        return constraints
+
+    # replace_index generates the new indices for the unfixed dimensions
+    # after eliminating the fixed dimensions.
+    # Example: dimension = 5, ff.keys() = [1, 3], replace_index = {0: 0, 2: 1, 4: 2}
+    unfixed_keys = sorted(set(range(dimension)) - set(fixed_features))
+    unfixed_keys = torch.tensor(unfixed_keys).to(constraints[0][0])
+    replace_index = torch.arange(dimension - len(fixed_features)).to(constraints[0][0])
+
+    new_constraints = []
+    # parse constraints one-by-one
+    for constraint_id, (indices, coefficients, rhs) in enumerate(constraints):
+        new_rhs = rhs
+        new_indices = []
+        new_coefficients = []
+        # the following unsqueeze is done to facilitate a simpler for-loop.
+        indices_2dim = indices if indices.ndim == 2 else indices.unsqueeze(-1)
+        for coefficient, index in zip(coefficients, indices_2dim):
+            ffval_or_None = fixed_features.get(index[-1].item())
+            # if ffval_or_None is None, then the index is not fixed
+            if ffval_or_None is None:
+                new_indices.append(index)
+                new_coefficients.append(coefficient)
+            # otherwise, we "remove" the constraints corresponding to that index
+            else:
+                new_rhs -= coefficient.item() * ffval_or_None
+
+        # all indices were fixed, so the constraint is gone.
+        if len(new_indices) == 0:
+            if (eq and new_rhs != 0) or (not eq and new_rhs > 0):
+                prefix = "Eq" if eq else "Ineq"
+                raise CandidateGenerationError(
+                    f"{prefix}ality constraint {constraint_id} not met "
+                    "with fixed_features."
+                )
+        else:
+            # However, one key transformation has to be noted.
+            # new_indices is with respect to the older (fuller) domain, and so it will
+            # have to be converted using replace_index.
+            new_indices = torch.stack(new_indices, dim=0)
+            # generate new index location after the removal of fixed_features indices
+            new_indices_dim_d = new_indices[:, -1].unsqueeze(-1)
+            new_indices_dim_d = replace_index[
+                torch.nonzero(new_indices_dim_d == unfixed_keys, as_tuple=True)[1]
+            ]
+            new_indices[:, -1] = new_indices_dim_d
+            # squeeze(-1) is a no-op if dim -1 is not singleton
+            new_indices.squeeze_(-1)
+            # convert new_coefficients to Tensor
+            new_coefficients = torch.stack(new_coefficients)
+            new_constraints.append((new_indices, new_coefficients, new_rhs))
+    return new_constraints