From 29226d8f9451e6dec2246672035284cd90f29f96 Mon Sep 17 00:00:00 2001
From: Sait Cakmak <saitcakmak@meta.com>
Date: Mon, 21 Oct 2024 20:56:15 -0700
Subject: [PATCH] Add an option for parameter constraints to
 get_branin_experiment (#2933)

Summary:
Pull Request resolved: https://github.com/facebook/Ax/pull/2933

Adds an option to add parameter constraints in `get_branin_experiment`, which will be used in a follow up diff for testing.

Reviewed By: sdaulton

Differential Revision: D64697748

fbshipit-source-id: cf0f04d18d04e3f9366ad9a8bf88f7ed2b4908dc
---
 ax/utils/testing/core_stubs.py | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/ax/utils/testing/core_stubs.py b/ax/utils/testing/core_stubs.py
index dbb7ede8e5c..2305dd74acc 100644
--- a/ax/utils/testing/core_stubs.py
+++ b/ax/utils/testing/core_stubs.py
@@ -246,6 +246,7 @@ def get_branin_experiment(
     with_fidelity_parameter: bool = False,
     with_choice_parameter: bool = False,
     with_str_choice_param: bool = False,
+    with_parameter_constraint: bool = False,
     search_space: SearchSpace | None = None,
     minimize: bool = False,
     named: bool = True,
@@ -257,6 +258,7 @@ def get_branin_experiment(
         with_fidelity_parameter=with_fidelity_parameter,
         with_choice_parameter=with_choice_parameter,
         with_str_choice_param=with_str_choice_param,
+        with_parameter_constraint=with_parameter_constraint,
     )
     exp = Experiment(
         name="branin_test_experiment" if named else None,
@@ -919,6 +921,7 @@ def get_branin_search_space(
     with_fidelity_parameter: bool = False,
     with_choice_parameter: bool = False,
     with_str_choice_param: bool = False,
+    with_parameter_constraint: bool = False,
 ) -> SearchSpace:
     parameters = [
         RangeParameter(
@@ -955,8 +958,16 @@ def get_branin_search_space(
                 target_value=1.0,
             )
         )
+    if with_parameter_constraint:
+        constraints = [
+            ParameterConstraint(constraint_dict={"x1": 1, "x2": 1}, bound=15.0)
+        ]
+    else:
+        constraints = None
 
-    return SearchSpace(parameters=cast(list[Parameter], parameters))
+    return SearchSpace(
+        parameters=cast(list[Parameter], parameters), parameter_constraints=constraints
+    )
 
 
 def get_factorial_search_space() -> SearchSpace: