Create helper for defining arm

Summary:
Currently, when users want to define status quo arms in Ax, they have to
enummerate every parameter key and, assuming they want defaults, set it as
None (e.g. https://fburl.com/wiki/zis37qp0). To simplify this, we introduce a helper utility that constructs a null
parameter arm by default, but allows users to override specific params if
necessary.

Reviewed By: sdsingh

Differential Revision: D20491298

fbshipit-source-id: 1aecbf6a5f15b27eb0967a78d35f02ffbb365cf0

Author: Steve Mandala

ax/core/search_space.py

@@ -17,6 +17,7 @@
     SumConstraint,
 )
 from ax.core.types import TParameterization
+from ax.utils.common.typeutils import not_none
 
 
 class SearchSpace(Base):
@@ -238,10 +239,29 @@ def out_of_design_arm(self) -> Arm:
         Returns:
             New arm w/ null parameter values.
         """
-        parameters = {}
-        for p_name in self.parameters.keys():
-            parameters[p_name] = None
-        return Arm(parameters)
+        return self.construct_arm()
+
+    def construct_arm(
+        self, parameters: Optional[TParameterization] = None, name: Optional[str] = None
+    ) -> Arm:
+        """Construct new arm using given parameters and name. Any
+        missing parameters fallback to the experiment defaults,
+        represented as None
+        """
+        final_parameters: TParameterization = {k: None for k in self.parameters.keys()}
+        if parameters is not None:
+            # Validate the param values
+            for p_name, p_value in parameters.items():
+                if p_name not in self.parameters:
+                    raise ValueError(f"`{p_name}` does not exist in search space.")
+                if p_value is not None and not self.parameters[p_name].validate(
+                    p_value
+                ):
+                    raise ValueError(
+                        f"`{p_value}` is not a valid value for parameter {p_name}."
+                    )
+            final_parameters.update(not_none(parameters))
+        return Arm(parameters=final_parameters, name=name)
 
     def clone(self) -> "SearchSpace":
         return SearchSpace(

ax/core/tests/test_search_space.py

@@ -302,3 +302,28 @@ def testOutOfDesignArm(self):
         arm1_nones = [p is None for p in arm1.parameters.values()]
         self.assertTrue(all(arm1_nones))
         self.assertTrue(arm1 == arm2)
+
+    def testConstructArm(self):
+        # Test constructing an arm of default values
+        arm = self.ss1.construct_arm(name="test")
+        self.assertEqual(arm.name, "test")
+        for p_name in self.ss1.parameters.keys():
+            self.assertTrue(p_name in arm.parameters)
+            self.assertEqual(arm.parameters[p_name], None)
+
+        # Test constructing an arm with a custom value
+        arm = self.ss1.construct_arm({"a": 1.0})
+        for p_name in self.ss1.parameters.keys():
+            self.assertTrue(p_name in arm.parameters)
+            if p_name == "a":
+                self.assertEqual(arm.parameters[p_name], 1.0)
+            else:
+                self.assertEqual(arm.parameters[p_name], None)
+
+        # Test constructing an arm with a bad param name
+        with self.assertRaises(ValueError):
+            self.ss1.construct_arm({"IDONTEXIST_a": 1.0})
+
+        # Test constructing an arm with a bad param name
+        with self.assertRaises(ValueError):
+            self.ss1.construct_arm({"a": "notafloat"})