Adding parameter "baseline_arm_name" to "_format_comparison_string" (f…

…acebook#2016) Summary: Updating _format_comparison_string to take baseline_arm_name as parameter, replacing the behavior of hardcoding baseline_arm_name using BASELINE_ARM_NAME. Reviewed By: bernardbeckerman Differential Revision: D51589439
mgrange1998 · Nov 27, 2023 · 048935f · 048935f
1 parent f8560e8
commit 048935f
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/ax/service/tests/test_report_utils.py b/ax/service/tests/test_report_utils.py
@@ -582,6 +582,7 @@ def test_compare_to_baseline(self) -> None:
 
             output_text = _format_comparison_string(
                 comparison_arm_name=comparison_arm_names[0],
+                baseline_arm_name=BASELINE_ARM_NAME,
                 objective_name=OBJECTIVE_METRIC,
                 percent_change=1150.0,
                 baseline_value=0.2,
@@ -645,6 +646,7 @@ def test_compare_to_baseline_pass_in_opt(self) -> None:
 
             output_text = _format_comparison_string(
                 comparison_arm_name=comparison_arm_names[0],
+                baseline_arm_name=BASELINE_ARM_NAME,
                 objective_name=OBJECTIVE_METRIC,
                 percent_change=1150.0,
                 baseline_value=0.2,
@@ -699,6 +701,7 @@ def test_compare_to_baseline_minimize(self) -> None:
 
             output_text = _format_comparison_string(
                 comparison_arm_name=comparison_arm_names[0],
+                baseline_arm_name=BASELINE_ARM_NAME,
                 objective_name=OBJECTIVE_METRIC,
                 percent_change=50.0,
                 baseline_value=0.2,
@@ -1008,6 +1011,7 @@ def test_compare_to_baseline_moo(self) -> None:
             )
             output_text_0 = _format_comparison_string(
                 comparison_arm_name="opt_0",
+                baseline_arm_name=BASELINE_ARM_NAME,
                 objective_name="m0",
                 percent_change=150.0,
                 baseline_value=1.0,
@@ -1016,6 +1020,7 @@ def test_compare_to_baseline_moo(self) -> None:
             )
             output_text_1 = _format_comparison_string(
                 comparison_arm_name="opt_1_min",
+                baseline_arm_name=BASELINE_ARM_NAME,
                 objective_name="m1",
                 percent_change=110.0,
                 baseline_value=1.0,
@@ -1024,6 +1029,7 @@ def test_compare_to_baseline_moo(self) -> None:
             )
             output_text_3 = _format_comparison_string(
                 comparison_arm_name="opt_3",
+                baseline_arm_name=BASELINE_ARM_NAME,
                 objective_name="m3",
                 percent_change=1.0,
                 baseline_value=1.0,

diff --git a/ax/service/utils/report_utils.py b/ax/service/utils/report_utils.py
@@ -1140,13 +1140,13 @@ def _warn_and_create_warning_plot(warning_msg: str) -> go.Figure:
 
 def _format_comparison_string(
     comparison_arm_name: str,
+    baseline_arm_name: str,
     objective_name: str,
     percent_change: float,
     baseline_value: float,
     comparison_value: float,
     digits: int,
 ) -> str:
-    baseline_arm_name = BASELINE_ARM_NAME
     return (
         f"{comparison_arm_name=} "
         + "improves your objective metric "
@@ -1186,6 +1186,7 @@ def _construct_comparison_message(
 
     return _format_comparison_string(
         comparison_arm_name=comparison_arm_name,
+        baseline_arm_name=BASELINE_ARM_NAME,
         objective_name=objective_name,
         percent_change=percent_change,
         baseline_value=baseline_value,