def setUp(self):
self.experiment = get_multi_type_experiment(add_trials=True)
self.data = self.experiment.fetch_data()
self.observations = observations_from_data(self.experiment, self.data)
self.observation_data = [o.data for o in self.observations]
self.observation_features = [o.features for o in self.observations]
self.tconfig = tconfig_from_mt_experiment(self.experiment)
def testObservationsFromDataWithFidelities(self):
truth = {
0.5: {
"arm_name": "0_0",
"parameters": {"x": 0, "y": "a", "z": 1},
"mean": 2.0,
"sem": 2.0,
"trial_index": 1,
"metric_name": "a",
"fidelities": json.dumps({"z": 0.5}),
"updated_parameters": {"x": 0, "y": "a", "z": 0.5},
"mean_t": np.array([2.0]),
"covariance_t": np.array([[4.0]]),
},
0.25: {
"arm_name": "0_1",
"parameters": {"x": 1, "y": "b", "z": 0.5},
"mean": 3.0,
"sem": 3.0,
"trial_index": 2,
"metric_name": "a",
"fidelities": json.dumps({"z": 0.25}),
"updated_parameters": {"x": 1, "y": "b", "z": 0.25},
"mean_t": np.array([3.0]),
"covariance_t": np.array([[9.0]]),
},
1: {
"arm_name": "0_0",
"parameters": {"x": 0, "y": "a", "z": 1},
"mean": 4.0,
"sem": 4.0,
"trial_index": 1,
"metric_name": "b",
"fidelities": json.dumps({"z": 1}),
"updated_parameters": {"x": 0, "y": "a", "z": 1},
"mean_t": np.array([4.0]),
"covariance_t": np.array([[16.0]]),
},
}
arms = {
obs["arm_name"]: Arm(parameters=obs["parameters"]) for obs in truth.values()
}
experiment = Mock()
type(experiment).arms_by_name = PropertyMock(return_value=arms)
df = pd.DataFrame(list(truth.values()))[
["arm_name", "trial_index", "mean", "sem", "metric_name", "fidelities"]
]
data = Data(df=df)
observations = observations_from_data(experiment, data)
self.assertEqual(len(observations), 3)
for obs in observations:
t = truth[obs.features.parameters["z"]]
self.assertEqual(obs.features.parameters, t["updated_parameters"])
self.assertEqual(obs.features.trial_index, t["trial_index"])
self.assertEqual(obs.data.metric_names, [t["metric_name"]])
self.assertTrue(np.array_equal(obs.data.means, t["mean_t"]))
self.assertTrue(np.array_equal(obs.data.covariance, t["covariance_t"]))
self.assertEqual(obs.arm_name, t["arm_name"])
def update(self, new_data: Data, experiment: Experiment) -> None:
"""Update the model bridge and the underlying model with new data. This
method should be used instead of `fit`, in cases where the underlying
model does not need to be re-fit from scratch, but rather updated.
Note: `update` expects only new data (obtained since the model initialization
or last update) to be passed in, not all data in the experiment.
Args:
new_data: Data from the experiment obtained since the last call to
`update`.
experiment: Experiment, in which this data was obtained.
"""
t_update_start = time.time()
observations = (observations_from_data(
experiment=experiment, data=new_data) if experiment is not None
and new_data is not None else [])
obs_feats_raw, obs_data_raw = self._extend_training_data(
observations=observations)
obs_feats, obs_data, search_space = self._transform_data(
obs_feats=obs_feats_raw,
obs_data=obs_data_raw,
search_space=self._model_space,
transforms=self._raw_transforms,
transform_configs=self._transform_configs,
)
self._update(
search_space=search_space,
observation_features=obs_feats,
observation_data=obs_data,
)
self.fit_time += time.time() - t_update_start
self.fit_time_since_gen += time.time() - t_update_start
def testObservationsFromData(self):
truth = [
{
"arm_name": "0_0",
"parameters": {"x": 0, "y": "a"},
"mean": 2.0,
"sem": 2.0,
"trial_index": 1,
"metric_name": "a",
},
{
"arm_name": "0_1",
"parameters": {"x": 1, "y": "b"},
"mean": 3.0,
"sem": 3.0,
"trial_index": 2,
"metric_name": "a",
},
{
"arm_name": "0_0",
"parameters": {"x": 0, "y": "a"},
"mean": 4.0,
"sem": 4.0,
"trial_index": 1,
"metric_name": "b",
},
]
arms = {obs["arm_name"]: Arm(parameters=obs["parameters"]) for obs in truth}
experiment = Mock()
type(experiment).arms_by_name = PropertyMock(return_value=arms)
df = pd.DataFrame(truth)[
["arm_name", "trial_index", "mean", "sem", "metric_name"]
]
data = Data(df=df)
observations = observations_from_data(experiment, data)
self.assertEqual(len(observations), 2)
# Get them in the order we want for tests below
if observations[0].features.parameters["x"] == 1:
observations.reverse()
obsd_truth = {
"metric_names": [["a", "b"], ["a"]],
"means": [np.array([2.0, 4.0]), np.array([3])],
"covariance": [np.diag([4.0, 16.0]), np.array([[9.0]])],
}
cname_truth = ["0_0", "0_1"]
for i, obs in enumerate(observations):
self.assertEqual(obs.features.parameters, truth[i]["parameters"])
self.assertEqual(obs.features.trial_index, truth[i]["trial_index"])
self.assertEqual(obs.data.metric_names, obsd_truth["metric_names"][i])
self.assertTrue(np.array_equal(obs.data.means, obsd_truth["means"][i]))
self.assertTrue(
np.array_equal(obs.data.covariance, obsd_truth["covariance"][i])
)
self.assertEqual(obs.arm_name, cname_truth[i])
def test_update(self, _mock_update, _mock_gen):
exp = get_experiment_for_value()
exp.optimization_config = get_optimization_config_no_constraints()
ss = get_search_space_for_range_values()
exp.search_space = ss
modelbridge = ModelBridge(
search_space=ss, model=Model(), transforms=[Log], experiment=exp
)
exp.new_trial(generator_run=modelbridge.gen(1))
modelbridge._set_training_data(
observations_from_data(
data=Data(
pd.DataFrame(
[
{
"arm_name": "0_0",
"metric_name": "m1",
"mean": 3.0,
"sem": 1.0,
}
]
)
),
experiment=exp,
),
ss,
)
exp.new_trial(generator_run=modelbridge.gen(1))
modelbridge.update(
new_data=Data(
pd.DataFrame(
[{"arm_name": "1_0", "metric_name": "m1", "mean": 5.0, "sem": 0.0}]
)
),
experiment=exp,
)
exp.new_trial(generator_run=modelbridge.gen(1))
# Trying to update with unrecognised metric should error.
with self.assertRaisesRegex(ValueError, "Unrecognised metric"):
modelbridge.update(
new_data=Data(
pd.DataFrame(
[
{
"arm_name": "1_0",
"metric_name": "m2",
"mean": 5.0,
"sem": 0.0,
}
]
)
),
experiment=exp,
)
def testObservationsWithCandidateMetadata(self):
SOME_METADATA_KEY = "metadatum"
truth = [
{
"arm_name": "0_0",
"parameters": {"x": 0, "y": "a"},
"mean": 2.0,
"sem": 2.0,
"trial_index": 0,
"metric_name": "a",
},
{
"arm_name": "1_0",
"parameters": {"x": 1, "y": "b"},
"mean": 3.0,
"sem": 3.0,
"trial_index": 1,
"metric_name": "a",
},
]
arms = {
obs["arm_name"]: Arm(name=obs["arm_name"], parameters=obs["parameters"])
for obs in truth
}
experiment = Mock()
experiment._trial_indices_by_status = {status: set() for status in TrialStatus}
trials = {
obs["trial_index"]: Trial(
experiment,
GeneratorRun(
arms=[arms[obs["arm_name"]]],
candidate_metadata_by_arm_signature={
arms[obs["arm_name"]].signature: {
SOME_METADATA_KEY: f"value_{obs['trial_index']}"
}
},
),
)
for obs in truth
}
type(experiment).arms_by_name = PropertyMock(return_value=arms)
type(experiment).trials = PropertyMock(return_value=trials)
df = pd.DataFrame(truth)[
["arm_name", "trial_index", "mean", "sem", "metric_name"]
]
data = Data(df=df)
observations = observations_from_data(experiment, data)
for observation in observations:
self.assertEqual(
observation.features.metadata.get(SOME_METADATA_KEY),
f"value_{observation.features.trial_index}",
)
def get_current_regrets(experiment: Experiment, observer: ExObserver):
observations = observations_from_data(experiment, experiment.eval())
observation_features, observation_data = separate_observations(
observations)
obs_x = observation_features_to_tensor(observation_features)
if obs_x.dim() == 1:
obs_x.unsqueeze_(-1)
obs_f = observation_data_to_tensor(observation_data)
opt_x = observer.current['optimum']['x']
opt_f = observer.current['optimum']['function']
loc_regret = torch.norm(obs_x - opt_x, dim=1).min()
fun_regret = torch.abs(obs_f - opt_f).min()
return loc_regret, fun_regret
def update(self, data: Data, experiment: Experiment) -> None:
"""Update the model bridge and the underlying model with new data. This
method should be used instead of `fit`, in cases where the underlying
model does not need to be re-fit from scratch, but rather updated.
Note: `update` expects only new data (obtained since the model initialization
or last update) to be passed in, not all data in the experiment.
Args:
data: data from the experiment obtained since the last update
experiment: experiment, in which this data was obtained
"""
t_update_start = time.time()
observations = (observations_from_data(experiment, data)
if experiment is not None and data is not None else [])
obs_feats, obs_data = self._extend_training_data(
observations=observations)
for t in self.transforms.values():
obs_feats = t.transform_observation_features(obs_feats)
obs_data = t.transform_observation_data(obs_data, obs_feats)
self._update(observation_features=obs_feats, observation_data=obs_data)
self.fit_time += time.time() - t_update_start
self.fit_time_since_gen += time.time() - t_update_start
def __init__(
self,
search_space: SearchSpace,
model: Any,
transforms: Optional[List[Type[Transform]]] = None,
experiment: Optional[Experiment] = None,
data: Optional[Data] = None,
transform_configs: Optional[Dict[str, TConfig]] = None,
status_quo_name: Optional[str] = None,
status_quo_features: Optional[ObservationFeatures] = None,
optimization_config: Optional[OptimizationConfig] = None,
) -> None:
"""
Applies transforms and fits model.
Args:
experiment: Is used to get arm parameters. Is not mutated.
search_space: Search space for fitting the model. Constraints need
not be the same ones used in gen.
data: Ax Data.
model: Interface will be specified in subclass. If model requires
initialization, that should be done prior to its use here.
transforms: List of uninitialized transform classes. Forward
transforms will be applied in this order, and untransforms in
the reverse order.
transform_configs: A dictionary from transform name to the
transform config dictionary.
status_quo_name: Name of the status quo arm. Can only be used if
Data has a single set of ObservationFeatures corresponding to
that arm.
status_quo_features: ObservationFeatures to use as status quo.
Either this or status_quo_name should be specified, not both.
optimization_config: Optimization config defining how to optimize
the model.
"""
t_fit_start = time.time()
self._metric_names: Set[str] = set()
self._training_data: List[Observation] = []
self._optimization_config: Optional[
OptimizationConfig] = optimization_config
self._training_in_design: List[bool] = []
self._status_quo: Optional[Observation] = None
self._arms_by_signature: Optional[Dict[str, Arm]] = None
self.transforms: MutableMapping[str, Transform] = OrderedDict()
self._model_space = search_space.clone()
if experiment is not None:
if self._optimization_config is None:
self._optimization_config = experiment.optimization_config
self._arms_by_signature = experiment.arms_by_signature
# Get observation features and data
obs_feats: List[ObservationFeatures] = []
obs_data: List[ObservationData] = []
observations = (observations_from_data(experiment, data)
if experiment is not None and data is not None else [])
obs_feats, obs_data = self._set_training_data(observations)
# Set model status quo
if any(x is not None
for x in [experiment, status_quo_name, status_quo_features]):
self._set_status_quo(
experiment=experiment,
status_quo_name=status_quo_name,
status_quo_features=status_quo_features,
)
# Initialize transforms
if transform_configs is None:
transform_configs = {}
search_space = search_space.clone()
if transforms is not None:
for t in transforms:
t_instance = t(
search_space=search_space,
observation_features=obs_feats,
observation_data=obs_data,
config=transform_configs.get(t.__name__, None),
)
search_space = t_instance.transform_search_space(search_space)
obs_feats = t_instance.transform_observation_features(
obs_feats)
obs_data = t_instance.transform_observation_data(
obs_data, obs_feats)
self.transforms[t.__name__] = t_instance
# Apply terminal transform and fit
try:
self._fit(
model=model,
search_space=search_space,
observation_features=obs_feats,
observation_data=obs_data,
)
self.fit_time = time.time() - t_fit_start
self.fit_time_since_gen = float(self.fit_time)
except NotImplementedError:
self.fit_time = 0.0
self.fit_time_since_gen = 0.0
def __init__(
self,
search_space: SearchSpace,
model: Any,
transforms: Optional[List[Type[Transform]]] = None,
experiment: Optional[Experiment] = None,
data: Optional[Data] = None,
transform_configs: Optional[Dict[str, TConfig]] = None,
status_quo_name: Optional[str] = None,
status_quo_features: Optional[ObservationFeatures] = None,
optimization_config: Optional[OptimizationConfig] = None,
fit_out_of_design: bool = False,
) -> None:
"""
Applies transforms and fits model.
Args:
experiment: Is used to get arm parameters. Is not mutated.
search_space: Search space for fitting the model. Constraints need
not be the same ones used in gen.
data: Ax Data.
model: Interface will be specified in subclass. If model requires
initialization, that should be done prior to its use here.
transforms: List of uninitialized transform classes. Forward
transforms will be applied in this order, and untransforms in
the reverse order.
transform_configs: A dictionary from transform name to the
transform config dictionary.
status_quo_name: Name of the status quo arm. Can only be used if
Data has a single set of ObservationFeatures corresponding to
that arm.
status_quo_features: ObservationFeatures to use as status quo.
Either this or status_quo_name should be specified, not both.
optimization_config: Optimization config defining how to optimize
the model.
"""
t_fit_start = time.time()
self._metric_names: Set[str] = set()
self._training_data: List[Observation] = []
self._optimization_config: Optional[OptimizationConfig] = optimization_config
self._training_in_design: List[bool] = []
self._status_quo: Optional[Observation] = None
self._arms_by_signature: Optional[Dict[str, Arm]] = None
self.transforms: MutableMapping[str, Transform] = OrderedDict()
self._model_key: Optional[str] = None
self._model_kwargs: Optional[Dict[str, Any]] = None
self._bridge_kwargs: Optional[Dict[str, Any]] = None
self._model_space = search_space.clone()
self._raw_transforms = transforms
self._transform_configs: Optional[Dict[str, TConfig]] = transform_configs
self._fit_out_of_design = fit_out_of_design
if experiment is not None:
if self._optimization_config is None:
self._optimization_config = experiment.optimization_config
self._arms_by_signature = experiment.arms_by_signature
observations = (
# pyre-fixme[6]: Expected `Experiment` for 1st param but got `None`.
observations_from_data(experiment, data)
if experiment is not None and data is not None
else []
)
obs_feats_raw, obs_data_raw = self._set_training_data(
observations=observations, search_space=search_space
)
# Set model status quo
# NOTE: training data must be set before setting the status quo.
self._set_status_quo(
experiment=experiment,
status_quo_name=status_quo_name,
status_quo_features=status_quo_features,
)
obs_feats, obs_data, search_space = self._transform_data(
obs_feats=obs_feats_raw,
obs_data=obs_data_raw,
search_space=search_space,
transforms=transforms,
transform_configs=transform_configs,
)
# Apply terminal transform and fit
try:
self._fit(
model=model,
search_space=search_space,
observation_features=obs_feats,
observation_data=obs_data,
)
self.fit_time = time.time() - t_fit_start
self.fit_time_since_gen = float(self.fit_time)
except NotImplementedError:
self.fit_time = 0.0
self.fit_time_since_gen = 0.0
def testObservationsFromDataWithSomeMissingTimes(self):
truth = [
{
"arm_name": "0_0",
"parameters": {
"x": 0,
"y": "a"
},
"mean": 2.0,
"sem": 2.0,
"trial_index": 1,
"metric_name": "a",
"start_time": 0,
},
{
"arm_name": "0_1",
"parameters": {
"x": 1,
"y": "b"
},
"mean": 3.0,
"sem": 3.0,
"trial_index": 2,
"metric_name": "a",
"start_time": 0,
},
{
"arm_name": "0_0",
"parameters": {
"x": 0,
"y": "a"
},
"mean": 4.0,
"sem": 4.0,
"trial_index": 1,
"metric_name": "b",
"start_time": None,
},
{
"arm_name": "0_1",
"parameters": {
"x": 1,
"y": "b"
},
"mean": 5.0,
"sem": 5.0,
"trial_index": 2,
"metric_name": "b",
"start_time": None,
},
]
arms = {
obs["arm_name"]: Arm(name=obs["arm_name"],
parameters=obs["parameters"])
for obs in truth
}
experiment = Mock()
experiment._trial_indices_by_status = {
status: set()
for status in TrialStatus
}
trials = {
obs["trial_index"]:
Trial(experiment, GeneratorRun(arms=[arms[obs["arm_name"]]]))
for obs in truth
}
type(experiment).arms_by_name = PropertyMock(return_value=arms)
type(experiment).trials = PropertyMock(return_value=trials)
df = pd.DataFrame(truth)[[
"arm_name", "trial_index", "mean", "sem", "metric_name",
"start_time"
]]
data = Data(df=df)
observations = observations_from_data(experiment, data)
self.assertEqual(len(observations), 4)
# Get them in the order we want for tests below
if observations[0].features.parameters["x"] == 1:
observations.reverse()
obsd_truth = {
"metric_names": [["a"], ["a"], ["b"], ["b"]],
"means": [
np.array([2.0]),
np.array([3.0]),
np.array([4.0]),
np.array([5.0]),
],
"covariance": [
np.diag([4.0]),
np.diag([9.0]),
np.diag([16.0]),
np.diag([25.0]),
],
}
cname_truth = ["0_0", "0_1", "0_0", "0_1"]
for i, obs in enumerate(observations):
self.assertEqual(obs.features.parameters, truth[i]["parameters"])
self.assertEqual(obs.features.trial_index, truth[i]["trial_index"])
self.assertEqual(obs.data.metric_names,
obsd_truth["metric_names"][i])
self.assertTrue(
np.array_equal(obs.data.means, obsd_truth["means"][i]))
self.assertTrue(
np.array_equal(obs.data.covariance,
obsd_truth["covariance"][i]))
self.assertEqual(obs.arm_name, cname_truth[i])
def testObservationsFromDataAbandoned(self):
truth = {
0.5: {
"arm_name": "0_0",
"parameters": {"x": 0, "y": "a", "z": 1},
"mean": 2.0,
"sem": 2.0,
"trial_index": 0,
"metric_name": "a",
"updated_parameters": {"x": 0, "y": "a", "z": 0.5},
"mean_t": np.array([2.0]),
"covariance_t": np.array([[4.0]]),
"z": 0.5,
"timestamp": 50,
},
1: {
"arm_name": "1_0",
"parameters": {"x": 0, "y": "a", "z": 1},
"mean": 4.0,
"sem": 4.0,
"trial_index": 1,
"metric_name": "b",
"updated_parameters": {"x": 0, "y": "a", "z": 1},
"mean_t": np.array([4.0]),
"covariance_t": np.array([[16.0]]),
"z": 1,
"timestamp": 100,
},
0.25: {
"arm_name": "2_0",
"parameters": {"x": 1, "y": "a", "z": 0.5},
"mean": 3.0,
"sem": 3.0,
"trial_index": 2,
"metric_name": "a",
"updated_parameters": {"x": 1, "y": "b", "z": 0.25},
"mean_t": np.array([3.0]),
"covariance_t": np.array([[9.0]]),
"z": 0.25,
"timestamp": 25,
},
0.75: {
"arm_name": "2_1",
"parameters": {"x": 1, "y": "b", "z": 0.75},
"mean": 3.0,
"sem": 3.0,
"trial_index": 2,
"metric_name": "a",
"updated_parameters": {"x": 1, "y": "b", "z": 0.75},
"mean_t": np.array([3.0]),
"covariance_t": np.array([[9.0]]),
"z": 0.75,
"timestamp": 25,
},
}
arms = {
obs["arm_name"]: Arm(name=obs["arm_name"], parameters=obs["parameters"])
for _, obs in truth.items()
}
experiment = Mock()
experiment._trial_indices_by_status = {status: set() for status in TrialStatus}
trials = {
obs["trial_index"]: (
Trial(experiment, GeneratorRun(arms=[arms[obs["arm_name"]]]))
)
for _, obs in list(truth.items())[:-1]
if not obs["arm_name"].startswith("2")
}
batch = BatchTrial(experiment, GeneratorRun(arms=[arms["2_0"], arms["2_1"]]))
trials.update({2: batch})
trials.get(1).mark_abandoned()
trials.get(2).mark_arm_abandoned(arm_name="2_1")
type(experiment).arms_by_name = PropertyMock(return_value=arms)
type(experiment).trials = PropertyMock(return_value=trials)
df = pd.DataFrame(list(truth.values()))[
["arm_name", "trial_index", "mean", "sem", "metric_name"]
]
data = Data(df=df)
# 1 arm is abandoned and 1 trial is abandoned, so only 2 observations should be
# included.
obs_no_abandoned = observations_from_data(experiment, data)
self.assertEqual(len(obs_no_abandoned), 2)
# 1 arm is abandoned and 1 trial is abandoned, so only 2 observations should be
# included.
obs_with_abandoned = observations_from_data(
experiment, data, include_abandoned=True
)
self.assertEqual(len(obs_with_abandoned), 4)
def testObservationsFromMapData(self):
truth = {
0.5: {
"arm_name": "0_0",
"parameters": {
"x": 0,
"y": "a",
"z": 1
},
"mean": 2.0,
"sem": 2.0,
"trial_index": 1,
"metric_name": "a",
"updated_parameters": {
"x": 0,
"y": "a",
"z": 0.5
},
"mean_t": np.array([2.0]),
"covariance_t": np.array([[4.0]]),
"z": 0.5,
"timestamp": 50,
},
0.25: {
"arm_name": "0_1",
"parameters": {
"x": 1,
"y": "b",
"z": 0.5
},
"mean": 3.0,
"sem": 3.0,
"trial_index": 2,
"metric_name": "a",
"updated_parameters": {
"x": 1,
"y": "b",
"z": 0.25
},
"mean_t": np.array([3.0]),
"covariance_t": np.array([[9.0]]),
"z": 0.25,
"timestamp": 25,
},
1: {
"arm_name": "0_0",
"parameters": {
"x": 0,
"y": "a",
"z": 1
},
"mean": 4.0,
"sem": 4.0,
"trial_index": 1,
"metric_name": "b",
"updated_parameters": {
"x": 0,
"y": "a",
"z": 1
},
"mean_t": np.array([4.0]),
"covariance_t": np.array([[16.0]]),
"z": 1,
"timestamp": 100,
},
}
arms = {
obs["arm_name"]: Arm(name=obs["arm_name"],
parameters=obs["parameters"])
for _, obs in truth.items()
}
experiment = Mock()
experiment._trial_indices_by_status = {
status: set()
for status in TrialStatus
}
trials = {
obs["trial_index"]:
Trial(experiment, GeneratorRun(arms=[arms[obs["arm_name"]]]))
for _, obs in truth.items()
}
type(experiment).arms_by_name = PropertyMock(return_value=arms)
type(experiment).trials = PropertyMock(return_value=trials)
df = pd.DataFrame(list(truth.values()))[[
"arm_name", "trial_index", "mean", "sem", "metric_name", "z",
"timestamp"
]]
data = MapData(df=df, map_keys=["z", "timestamp"])
observations = observations_from_data(experiment, data)
self.assertEqual(len(observations), 3)
for obs in observations:
t = truth[obs.features.parameters["z"]]
self.assertEqual(obs.features.parameters, t["updated_parameters"])
self.assertEqual(obs.features.trial_index, t["trial_index"])
self.assertEqual(obs.data.metric_names, [t["metric_name"]])
self.assertTrue(np.array_equal(obs.data.means, t["mean_t"]))
self.assertTrue(
np.array_equal(obs.data.covariance, t["covariance_t"]))
self.assertEqual(obs.arm_name, t["arm_name"])
self.assertEqual(obs.features.metadata,
{"timestamp": t["timestamp"]})
def test_multitask_data(self):
experiment = get_branin_with_multi_task()
data = experiment.fetch_data()
observations = observations_from_data(
experiment=experiment,
data=data,
)
relative_observations = observations_from_data(
experiment=experiment,
data=relativize_data(
data=data,
status_quo_name="status_quo",
as_percent=True,
include_sq=True,
),
)
status_quo_row = data.df.loc[
(data.df["arm_name"] == "status_quo") & (data.df["trial_index"] == 1)
]
modelbridge = Mock(
status_quo=Observation(
data=ObservationData(
metric_names=status_quo_row["metric_name"].values,
means=status_quo_row["mean"].values,
covariance=np.array([status_quo_row["sem"].values ** 2]),
),
features=ObservationFeatures(
parameters=experiment.status_quo.parameters
),
)
)
obs_features = [obs.features for obs in observations]
obs_data = [obs.data for obs in observations]
expected_obs_data = [obs.data for obs in relative_observations]
transform = Relativize(
search_space=None,
observation_features=obs_features,
observation_data=obs_data,
modelbridge=modelbridge,
)
relative_obs_data = transform.transform_observation_data(obs_data, obs_features)
self.maxDiff = None
# this assertion just checks that order is the same, which
# is only important for the purposes of this test
self.assertEqual(
[datum.metric_names for datum in relative_obs_data],
[datum.metric_names for datum in expected_obs_data],
)
means = [
np.array([datum.means for datum in relative_obs_data]),
np.array([datum.means for datum in expected_obs_data]),
]
# `self.assertAlmostEqual(relative_obs_data, expected_obs_data)`
# fails 1% of the time, so we check with numpy.
self.assertTrue(
all(np.isclose(means[0], means[1])),
means,
)
covariances = [
np.array([datum.covariance for datum in expected_obs_data]),
np.array([datum.covariance for datum in relative_obs_data]),
]
self.assertTrue(
all(np.isclose(covariances[0], covariances[1])),
covariances,
)
