Remove parameter constraints when expanding model space in Adapter (#3408)

Summary:
Pull Request resolved: #3408

By default we expand range parameters in the model space so that the
training data include all modelable data, not just those that fall within the
box bounds.

A similar argument suggests that we should also include training data that
violate parameter constraints in the model. They are modelable and can be used
to improve predictions at points that satisfy parameter constraints.

We do this by simply dropping parameter constraints from the model space, if
the model space is specified to be expanded.

This will also fix #1568 . The issue there
is that generated points violated parameter constraints due to numerical
issues, and so then were left out of the training data and repeatedly
generated. With this change, those points would be included in the training
data and so likely not re-generated. Note, however, that the purpose of this
change is not just to fix this issue; it is a change that is right because it
better fits the concept of the modeling space.

Reviewed By: sdaulton, saitcakmak

Differential Revision: D69880066

fbshipit-source-id: c5bc5d042b567eceefe9ef956240b4577c0fb5cb

Author: bletham

ax/modelbridge/base.py

@@ -140,7 +140,8 @@ def __init__(
             expand_model_space: If True, expand range parameter bounds in model
                 space to cover given training data. This will make the modeling
                 space larger than the search space if training data fall outside
-                the search space.
+                the search space. Will also include training points that violate
+                parameter constraints in the modeling.
             fit_out_of_design: If specified, all training data are used.
                 Otherwise, only in design points are used.
             fit_abandoned: Whether data for abandoned arms or trials should be
@@ -442,6 +443,8 @@ def _set_model_space(self, observations: list[Observation]) -> None:
             if isinstance(p, RangeParameter):
                 p.lower = min(p.lower, min(param_vals[p.name]))
                 p.upper = max(p.upper, max(param_vals[p.name]))
+        # Remove parameter constraints from the model space.
+        self._model_space.set_parameter_constraints([])
 
     def _set_status_quo(
         self,

ax/modelbridge/tests/test_base_modelbridge.py

@@ -22,6 +22,7 @@
 from ax.core.observation import ObservationData, ObservationFeatures
 from ax.core.optimization_config import OptimizationConfig
 from ax.core.parameter import FixedParameter, ParameterType, RangeParameter
+from ax.core.parameter_constraint import SumConstraint
 from ax.core.search_space import SearchSpace
 from ax.exceptions.core import UnsupportedError, UserInputError
 from ax.modelbridge.base import (
@@ -1000,10 +1001,21 @@ def test_SetModelSpace(self) -> None:
         experiment.attach_data(get_branin_data_batch(batch=trial, fill_vals=sq_vals))
         trial.mark_completed()
         data = experiment.lookup_data()
+        # Make search space with a parameter constraint
+        ss = experiment.search_space.clone()
+        ss.set_parameter_constraints(
+            [
+                SumConstraint(
+                    parameters=list(ss.parameters.values()),
+                    is_upper_bound=True,
+                    bound=30.0,
+                )
+            ]
+        )
 
         # Check that SQ and custom are OOD
         m = Adapter(
-            search_space=experiment.search_space,
+            search_space=ss,
             model=None,
             experiment=experiment,
             data=data,
@@ -1014,10 +1026,11 @@ def test_SetModelSpace(self) -> None:
         self.assertEqual(set(ood_arms), {"status_quo", "custom"})
         self.assertEqual(m.model_space.parameters["x1"].lower, -5.0)  # pyre-ignore[16]
         self.assertEqual(m.model_space.parameters["x2"].upper, 15.0)  # pyre-ignore[16]
+        self.assertEqual(len(m.model_space.parameter_constraints), 1)
 
         # With expand model space, custom is not OOD, and model space is expanded
         m = Adapter(
-            search_space=experiment.search_space,
+            search_space=ss,
             model=None,
             experiment=experiment,
             data=data,
@@ -1027,10 +1040,11 @@ def test_SetModelSpace(self) -> None:
         self.assertEqual(set(ood_arms), {"status_quo"})
         self.assertEqual(m.model_space.parameters["x1"].lower, -20.0)
         self.assertEqual(m.model_space.parameters["x2"].upper, 18.0)
+        self.assertEqual(m.model_space.parameter_constraints, [])
 
         # With fill values, SQ is also in design, and x2 is further expanded
         m = Adapter(
-            search_space=experiment.search_space,
+            search_space=ss,
             model=None,
             experiment=experiment,
             data=data,
@@ -1039,6 +1053,7 @@ def test_SetModelSpace(self) -> None:
         )
         self.assertEqual(sum(m.training_in_design), 7)
         self.assertEqual(m.model_space.parameters["x2"].upper, 20)
+        self.assertEqual(m.model_space.parameter_constraints, [])
 
     @mock.patch(
         "ax.modelbridge.base.observations_from_data",