def prepare_experiment_for_plotting(
experiment: Experiment,
ignored_names: Optional[List[str]] = None,
) -> pd.DataFrame:
"""Strip variables not desired in the final plot and truncate names for readability
Args:
experiment: Experiment containing trials to plot
ignored_names: Metrics present in the experiment data we wish to exclude from
the final plot. By default we ignore ["generation_method", "trial_status",
"arm_name"]
Returns:
df.DataFrame: data frame ready for ingestion by plotly
"""
ignored_names = (["generation_method", "trial_status", "arm_name"]
if ignored_names is None else ignored_names)
df = exp_to_df(experiment)
dropped = df.drop(ignored_names, axis=1)
renamed = dropped.rename(
# pyre-fixme[6] Expected `typing.Union[
# typing.Callable[[Optional[typing.Hashable]], Optional[typing.Hashable]],
# None, typing.Mapping[Optional[typing.Hashable], typing.Any]]` for 1st
# parameter `columns` to call `pd.core.frame.DataFrame.rename` but got
# `typing.Dict[str, str]`.
columns=_get_shortest_unique_suffix_dict(
[str(c) for c in dropped.columns]))
return renamed
def test_get_shortest_unique_suffix_dict(self):
expected_output = {
"abc.123": "abc.123",
"asdf.abc.123": "asdf.abc.123",
"def.123": "def.123",
"abc.456": "456",
"": "",
"no_delimiter": "no_delimiter",
}
actual_output = _get_shortest_unique_suffix_dict([
"abc.123", "abc.456", "def.123", "asdf.abc.123", "", "no_delimiter"
])
self.assertDictEqual(expected_output, actual_output)