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)