Skip to content

Preserve treatment SEMs for zero relative changes #5302

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
saitcakmak wants to merge 1 commit into from
Closed

Preserve treatment SEMs for zero relative changes #5302

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 13 additions & 3 deletions ax/adapter/transforms/relativize.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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`
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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,
Expand Down
175 changes: 175 additions & 0 deletions ax/adapter/transforms/tests/test_relativize_transform.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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):
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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))
},
Expand Down
79 changes: 78 additions & 1 deletion ax/adapter/transforms/tests/test_transform_to_new_sq.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down
1 change: 1 addition & 0 deletions ax/adapter/transforms/transform_to_new_sq.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down
23 changes: 13 additions & 10 deletions ax/utils/stats/math_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Loading
Loading