From 5f26309c5f9c82209cc1f427f343179295036d37 Mon Sep 17 00:00:00 2001 From: Sait Cakmak Date: Mon, 19 Oct 2026 11:57:19 -0700 Subject: [PATCH] Preserve treatment SEMs for zero relative changes Summary: ## What Changed - Keep constant-control unrelativize SEMs on the rescaled treatment value, including zero relative means. - Identify the real status quo by arm name so a treatment with identical statistics is not mistaken for the control in constant-control mode. - Restore true status-quo (0, 0) observations inside BaseRelativize._get_rel_mean_sem, preserving subclass dispatch and TransformToNewSQ absolute-scale interpretation. - Restrict the raw-statistics identity shortcut to forward relativization, avoiding accidental equality between relative predictions and absolute controls during inverse transforms. - Clarify that the utility zero-relative-mean override applies only outside constant-control mode. Addresses review comments 1, 2, and 4. Nonconstant-control forward variance parity (comment 3) is isolated in a separate stacked change tracked by T293694908. Based on https://github.com/facebook/Ax/issues/5301. Differential Revision: D126108875 --- ax/adapter/transforms/relativize.py | 16 +- .../tests/test_relativize_transform.py | 175 ++++++++++++++++++ .../tests/test_transform_to_new_sq.py | 79 +++++++- ax/adapter/transforms/transform_to_new_sq.py | 1 + ax/utils/stats/math_utils.py | 23 ++- ax/utils/stats/tests/test_math_utils.py | 53 ++++++ 6 files changed, 333 insertions(+), 14 deletions(-) diff --git a/ax/adapter/transforms/relativize.py b/ax/adapter/transforms/relativize.py index d3fb5dd2376..58e4d5d5447 100644 --- a/ax/adapter/transforms/relativize.py +++ b/ax/adapter/transforms/relativize.py @@ -290,6 +290,7 @@ def _get_relative_data_from_obs( data=obs.data, status_quo_data=self.status_quo_data_by_trial[idx], rel_op=rel_op, + is_status_quo=obs.arm_name == self.status_quo_name, ) def _rel_op_on_observations( @@ -311,6 +312,7 @@ def _get_relative_data( data: ObservationData, status_quo_data: ObservationData, rel_op: Callable[..., tuple[npt.NDArray, npt.NDArray]], + is_status_quo: bool, ) -> ObservationData: r""" Relativize or unrelativize `data` based on `status_quo_data` based on `rel_op` @@ -319,7 +321,7 @@ def _get_relative_data( data: ObservationData object to relativize status_quo_data: The status quo data (un)relativization is based upon rel_op: relativize or unrelativize operator. - control_as_constant: if treating the control metric as constant + is_status_quo: Whether the observation belongs to the status quo arm. Returns: (un)relativized ObservationData @@ -345,6 +347,7 @@ def _get_relative_data( sem_c=sem_c, metric=metric, rel_op=rel_op, + is_status_quo=is_status_quo, ) result.means[i] = means_rel result.covariance[i][i] = sems_rel**2 @@ -358,10 +361,17 @@ def _get_rel_mean_sem( sem_c: float, metric: str, rel_op: Callable[..., tuple[npt.NDArray, npt.NDArray]], + is_status_quo: bool, ) -> tuple[float | npt.NDArray, float | npt.NDArray]: """Compute (un)relativized mean and sem for a single metric.""" - # if the is the status quo - if means_t == mean_c and sems_t == sem_c: + if is_status_quo and rel_op == unrelativize and means_t == sems_t == 0: + return mean_c, sem_c + if ( + rel_op == relativize + and (is_status_quo or not self.control_as_constant) + and means_t == mean_c + and sems_t == sem_c + ): return 0, 0 return rel_op( means_t=means_t, diff --git a/ax/adapter/transforms/tests/test_relativize_transform.py b/ax/adapter/transforms/tests/test_relativize_transform.py index 8e196203bb7..7a3a42d3bb4 100644 --- a/ax/adapter/transforms/tests/test_relativize_transform.py +++ b/ax/adapter/transforms/tests/test_relativize_transform.py @@ -267,6 +267,7 @@ def _check_transform_observations( features=obs_features[2], arm_name=arm_names[2], ), + status_quo_name=arm_names[2], status_quo_data_by_trial={0: obs_data[0], 1: obs_data[2]}, ) tf = relativize_cls( @@ -282,6 +283,178 @@ def _check_transform_observations( expected_mean_and_covar=expected_mean_and_covar[2:4], ) + def test_constant_control_equal_means_round_trip(self) -> None: + """Restore the true SQ without overwriting equal-mean treatment variances.""" + for mean, sq_variance in product((10.0, -10.0), (0.0, 0.01)): + with self.subTest(mean=mean, sq_variance=sq_variance): + observations = [ + Observation( + features=ObservationFeatures( + parameters={"x": i}, trial_index=0 + ), + data=ObservationData( + metric_signatures=["m"], + means=np.array([mean]), + covariance=np.array([[variance]]), + ), + arm_name="status_quo" if i == 0 else str(i), + ) + for i, variance in enumerate((sq_variance, 4.0, 0.0, sq_variance)) + ] + transform = RelativizeWithConstantControl( + adapter=Mock( + status_quo_name="status_quo", + status_quo_data_by_trial={0: observations[0].data}, + _experiment=None, + ) + ) + relative = transform.transform_observations(observations) + self.assertAllClose(relative[0].data.means, np.array([0.0])) + self.assertAllClose(relative[0].data.covariance, np.array([[0.0]])) + self.assertAllClose( + relative[3].data.covariance, np.array([[100.0 * sq_variance]]) + ) + restored = transform.untransform_observations(relative) + for original, result in zip(observations, restored): + self.assertAllClose(result.data.means, original.data.means) + self.assertAllClose( + result.data.covariance, original.data.covariance + ) + + def test_nonconstant_control_identical_data_round_trip(self) -> None: + """Keep the legacy identical-data shortcut outside constant-control mode.""" + data = ObservationData( + metric_signatures=["m"], + means=np.array([10.0]), + covariance=np.array([[0.01]]), + ) + transform = Relativize( + adapter=Mock( + status_quo_name="status_quo", + status_quo_data_by_trial={0: data}, + _experiment=None, + ) + ) + observation = Observation( + features=ObservationFeatures(parameters={}, trial_index=0), + data=data, + arm_name="treatment", + ) + relative = transform.transform_observations([observation]) + self.assertAllClose(relative[0].data.means, np.array([0.0])) + self.assertAllClose(relative[0].data.covariance, np.array([[0.0]])) + restored = transform.untransform_observations(relative)[0] + self.assertAllClose(restored.data.means, data.means) + self.assertAllClose(restored.data.covariance, data.covariance) + + def test_untransform_status_quo_by_trial_and_metric(self) -> None: + """Restore each trial's SQ in observation metric order, including fallback.""" + sq_data = { + 0: ObservationData( + metric_signatures=["a", "b"], + means=np.array([10.0, 20.0]), + covariance=np.diag([0.01, 0.09]), + ), + 1: ObservationData( + metric_signatures=["a", "b"], + means=np.array([-10.0, 30.0]), + covariance=np.diag([0.04, 0.16]), + ), + } + transform = RelativizeWithConstantControl( + adapter=Mock( + status_quo_name="status_quo", + status_quo_data_by_trial=sq_data, + _experiment=None, + ) + ) + for trial_index in (0, 1, None): + with self.subTest(trial_index=trial_index): + observation = Observation( + features=ObservationFeatures( + parameters={}, trial_index=trial_index + ), + data=ObservationData( + metric_signatures=["b", "a"], + means=np.zeros(2), + covariance=np.zeros((2, 2)), + ), + arm_name="status_quo", + ) + restored = transform.untransform_observations([observation])[0] + expected = sq_data[1 if trial_index is None else trial_index] + self.assertAllClose(restored.data.means, expected.means[::-1].copy()) + self.assertAllClose( + restored.data.covariance, expected.covariance[::-1, ::-1].copy() + ) + + def test_status_quo_and_unnamed_predictions_are_untransformed(self) -> None: + """Unnamed predictions retain uncertainty; only the true SQ at (0, 0) resets.""" + transform = RelativizeWithConstantControl( + adapter=Mock( + status_quo_name="status_quo", + status_quo_data_by_trial={ + 0: ObservationData( + metric_signatures=["m"], + means=np.array([10.0]), + covariance=np.array([[0.01]]), + ) + }, + _experiment=None, + ) + ) + for arm_name, mean, variance in product( + (None, "status_quo"), (0.0, 5.0, 10.0), (0.0, 0.01, 400.0) + ): + with self.subTest(arm_name=arm_name, mean=mean, variance=variance): + observation = Observation( + features=ObservationFeatures(parameters={}, trial_index=0), + data=ObservationData( + metric_signatures=["m"], + means=np.array([mean]), + covariance=np.array([[variance]]), + ), + arm_name=arm_name, + ) + restored = transform.untransform_observations([observation])[0] + expected_variance = ( + 0.01 + if arm_name == "status_quo" and mean == variance == 0 + else variance / 100.0 + ) + self.assertAllClose(restored.data.means, np.array([10.0 + mean / 10.0])) + self.assertAllClose( + restored.data.covariance, np.array([[expected_variance]]) + ) + + def test_nonconstant_inverse_does_not_use_forward_shortcut(self) -> None: + """Matching relative and absolute numbers must not zero an inverse result.""" + transform = Relativize( + adapter=Mock( + status_quo_name="status_quo", + status_quo_data_by_trial={ + 0: ObservationData( + metric_signatures=["m"], + means=np.array([10.0]), + covariance=np.array([[0.01]]), + ) + }, + _experiment=None, + ) + ) + observation = Observation( + features=ObservationFeatures(parameters={}, trial_index=0), + data=ObservationData( + metric_signatures=["m"], + means=np.array([10.0]), + covariance=np.array([[0.01]]), + ), + arm_name="treatment", + ) + restored = transform.untransform_observations([observation])[0] + self.assertAllClose(restored.data.means, np.array([11.0 / (1.0 - 0.0001)])) + self.assertAllClose(restored.data.covariance, np.array([[0.0]])) + def test_bad_relativize(self) -> None: # Check instantiation and subclassing of BaseRelativize class BadRelativize(BaseRelativize): @@ -322,6 +495,7 @@ def test_transform_status_quos_always_zero(self) -> None: status_quo=Mock( data=obs_data[0], features=obs_features[0], arm_name=arm_names[0] ), + status_quo_name=arm_names[0], status_quo_data_by_trial={0: obs_data[0]}, ) observations = recombine_observations(obs_features, obs_data, arm_names) @@ -369,6 +543,7 @@ def test_multitask_data(self) -> None: ), arm_name="status_quo", ), + status_quo_name="status_quo", status_quo_data_by_trial={ i: sq_obs_data[i] for i in range(len(sq_obs_data)) }, diff --git a/ax/adapter/transforms/tests/test_transform_to_new_sq.py b/ax/adapter/transforms/tests/test_transform_to_new_sq.py index a25af09cf6c..8228f353799 100644 --- a/ax/adapter/transforms/tests/test_transform_to_new_sq.py +++ b/ax/adapter/transforms/tests/test_transform_to_new_sq.py @@ -20,11 +20,12 @@ from ax.adapter.transforms.tests.test_relativize_transform import RelativizeDataTest from ax.adapter.transforms.transform_to_new_sq import TransformToNewSQ from ax.core.batch_trial import BatchTrial +from ax.core.observation import Observation, ObservationData, ObservationFeatures from ax.core.observation_utils import observations_from_data from ax.exceptions.core import DataRequiredError from ax.generators.base import Generator from ax.utils.common.testutils import TestCase -from ax.utils.stats.math_utils import relativize +from ax.utils.stats.math_utils import relativize, unrelativize from ax.utils.testing.core_stubs import ( get_branin_data_batch, get_branin_experiment, @@ -96,6 +97,82 @@ def _refresh_adapter(self) -> None: data=self.exp.lookup_data(), ) + def test_equal_mean_treatment_round_trip(self) -> None: + """Moving to a new SQ and back must preserve equal-mean treatment SEMs.""" + for mean_c in (10.0, -10.0): + with self.subTest(mean_c=mean_c): + sq_data = { + idx: ObservationData( + metric_signatures=["m"], + means=np.array([mean]), + covariance=np.array([[0.01]]), + ) + for idx, mean in enumerate((mean_c, 30.0)) + } + transform = TransformToNewSQ( + adapter=mock.Mock( + status_quo_name="status_quo", + status_quo_data_by_trial=sq_data, + _experiment=None, + ), + config={"target_trial_index": 1}, + ) + observations = [ + Observation( + features=ObservationFeatures( + parameters={"x": i}, trial_index=0 + ), + data=ObservationData( + metric_signatures=["m"], + means=np.array([mean_c]), + covariance=np.array([[variance]]), + ), + arm_name=str(i), + ) + for i, variance in enumerate((4.0, 0.0)) + ] + transformed = transform.transform_observations(observations) + self.assertAllClose(transformed[0].data.means, np.array([30.0])) + self.assertAllClose(transformed[0].data.covariance, np.array([[36.0]])) + restored = transform.untransform_observations(transformed) + for original, result in zip(observations, restored): + self.assertAllClose(result.data.means, original.data.means) + self.assertAllClose( + result.data.covariance, original.data.covariance + ) + + def test_status_quo_inverse_uses_target_absolute_scale(self) -> None: + """The subclass hook, not the base data loop, interprets inverse values.""" + transform = TransformToNewSQ( + adapter=mock.Mock( + status_quo_name="status_quo", + status_quo_data_by_trial={ + idx: ObservationData( + metric_signatures=["m"], + means=np.array([mean]), + covariance=np.array([[0.01]]), + ) + for idx, mean in enumerate((10.0, 30.0)) + }, + _experiment=None, + ), + config={"target_trial_index": 1}, + ) + observation = Observation( + features=ObservationFeatures(parameters={}, trial_index=0), + data=ObservationData( + metric_signatures=["m"], + means=np.zeros(1), + covariance=np.zeros((1, 1)), + ), + arm_name="status_quo", + ) + restored = transform._get_relative_data_from_obs( + obs=observation, rel_op=unrelativize + ) + self.assertAllClose(restored.means, np.zeros(1)) + self.assertAllClose(restored.covariance, np.zeros((1, 1))) + def test_adapter_without_status_quo_name(self) -> None: self.adapter._status_quo = None self.adapter._status_quo_name = None diff --git a/ax/adapter/transforms/transform_to_new_sq.py b/ax/adapter/transforms/transform_to_new_sq.py index 906d899687f..a3d23f15a62 100644 --- a/ax/adapter/transforms/transform_to_new_sq.py +++ b/ax/adapter/transforms/transform_to_new_sq.py @@ -250,6 +250,7 @@ def _get_rel_mean_sem( sem_c: float, metric: str, rel_op: Callable[..., tuple[npt.NDArray, npt.NDArray]], + is_status_quo: bool, ) -> tuple[float, float]: """Compute (un)transformed mean and sem for a single metric.""" target_status_quo_data = self.status_quo_data_by_trial[self.default_trial_idx] diff --git a/ax/utils/stats/math_utils.py b/ax/utils/stats/math_utils.py index 546aff4a918..c372fb4c613 100644 --- a/ax/utils/stats/math_utils.py +++ b/ax/utils/stats/math_utils.py @@ -143,6 +143,9 @@ def unrelativize( (i.e., 1 means 1%). control_as_constant: If true, control is treated as a constant. bias_correction, sem_c, and cov_means are ignored when this is true. + SEMs are rescaled by abs(mean_c), including for zero relative changes. + The zero-relative-mean override that restores the control mean and SEM + applies only when control_as_constant=False. Returns: m_t: Inferred sample (test) means in the unrelativized scale @@ -174,15 +177,15 @@ def unrelativize( # If above condition cannot be guaranteed, use control_as_constant = True s_t = np.sqrt(s_t2.clip(min=0.0)) - # if means_t is 0.0 exactly, return control mean and sem directly - if np.isscalar(means_t): - if means_t == 0.0: - m_t = mean_c - s_t = sem_c - else: - m_t = np.array(m_t) - s_t = np.array(s_t) - m_t[means_t == 0.0] = mean_c - s_t[means_t == 0.0] = sem_c + # if means_t is 0.0 exactly, return control mean and sem directly + if np.isscalar(means_t): + if means_t == 0.0: + m_t = mean_c + s_t = sem_c + else: + m_t = np.array(m_t) + s_t = np.array(s_t) + m_t[means_t == 0.0] = mean_c + s_t[means_t == 0.0] = sem_c return np.asarray(m_t), np.asarray(s_t) diff --git a/ax/utils/stats/tests/test_math_utils.py b/ax/utils/stats/tests/test_math_utils.py index 6b063829950..398be825992 100644 --- a/ax/utils/stats/tests/test_math_utils.py +++ b/ax/utils/stats/tests/test_math_utils.py @@ -341,3 +341,56 @@ def test_unrelativize(self) -> None: ) self.assertTrue(np.allclose(means_t, unrel_mean_t)) self.assertTrue(np.allclose(sems_t, unrel_sems_t)) + + def test_constant_control_zero_change_round_trip(self) -> None: + """Equal means must not replace treatment SEMs with control uncertainty.""" + for mean_c, sem_c, as_percent in product( + (10.0, -10.0), (0.0, 0.1, np.nan), (False, True) + ): + with self.subTest(mean_c=mean_c, sem_c=sem_c, as_percent=as_percent): + means = np.array([mean_c, mean_c, mean_c + 1.0]) + sems = np.array([2.0, 0.0, 3.0]) + relative_means, relative_sems = relativize( + means_t=means, + sems_t=sems, + mean_c=mean_c, + sem_c=sem_c, + as_percent=as_percent, + control_as_constant=True, + ) + restored_means, restored_sems = unrelativize( + means_t=relative_means, + sems_t=relative_sems, + mean_c=mean_c, + sem_c=sem_c, + as_percent=as_percent, + control_as_constant=True, + ) + self.assertAllClose(restored_means, means) + self.assertAllClose(restored_sems, sems) + scalar_mean, scalar_sem = unrelativize( + means_t=float(relative_means[0]), + sems_t=float(relative_sems[0]), + mean_c=mean_c, + sem_c=sem_c, + as_percent=as_percent, + control_as_constant=True, + ) + self.assertAlmostEqual(float(scalar_mean), mean_c) + self.assertAlmostEqual(float(scalar_sem), 2.0) + + def test_nonconstant_control_zero_change(self) -> None: + """Preserve the legacy status-quo override for the delta-method inverse.""" + for bias_correction, as_percent in product((False, True), (False, True)): + with self.subTest(bias_correction=bias_correction, as_percent=as_percent): + means, sems = unrelativize( + means_t=np.array([0.0]), + sems_t=np.array([0.2]), + mean_c=10.0, + sem_c=0.1, + bias_correction=bias_correction, + as_percent=as_percent, + control_as_constant=False, + ) + self.assertAllClose(means, np.array([10.0])) + self.assertAllClose(sems, np.array([0.1]))