From 6d125fd1236513a10c9529d746652a37f41f5fc8 Mon Sep 17 00:00:00 2001 From: Xavier Bouthillier Date: Sat, 16 Sep 2023 15:45:38 -0400 Subject: [PATCH] Add new analysis method for symbolic explanation from Segel et al --- src/orion/analysis/__init__.py | 1 + .../analysis/partial_dependency_utils.py | 7 +- src/orion/analysis/symbolic_explanation.py | 256 +++++++++++++++++ src/orion/plotting/backend_plotly.py | 263 +++++++++++++++++- src/orion/plotting/base.py | 96 +++++++ 5 files changed, 619 insertions(+), 4 deletions(-) create mode 100644 src/orion/analysis/symbolic_explanation.py diff --git a/src/orion/analysis/__init__.py b/src/orion/analysis/__init__.py index 9df822b23..f99864856 100644 --- a/src/orion/analysis/__init__.py +++ b/src/orion/analysis/__init__.py @@ -7,5 +7,6 @@ from orion.analysis.lpi_utils import lpi from orion.analysis.partial_dependency_utils import partial_dependency from orion.analysis.regret_utils import regret +from orion.analysis.symbolic_explanation import symbolic_explanation __all__ = ["average", "ranking", "lpi", "partial_dependency", "regret"] diff --git a/src/orion/analysis/partial_dependency_utils.py b/src/orion/analysis/partial_dependency_utils.py index 9989b95db..8c7c0ea10 100644 --- a/src/orion/analysis/partial_dependency_utils.py +++ b/src/orion/analysis/partial_dependency_utils.py @@ -76,9 +76,10 @@ def partial_dependency( if trials.empty or trials.shape[0] == 0: return {} - data = to_numpy(trials, space) - data = flatten_numpy(data, flattened_space) - model = train_regressor(model, data, **kwargs) + if isinstance(model, str): + data = to_numpy(trials, space) + data = flatten_numpy(data, flattened_space) + model = train_regressor(model, data, **kwargs) data = [ format_trials.trial_to_tuple(trial, flattened_space) diff --git a/src/orion/analysis/symbolic_explanation.py b/src/orion/analysis/symbolic_explanation.py new file mode 100644 index 000000000..b6054e51b --- /dev/null +++ b/src/orion/analysis/symbolic_explanation.py @@ -0,0 +1,256 @@ +""" +Module for symbolic explanation. + +Large chunk of codes are copied from this repository: https://github.com/automl/symbolic-explanations +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass, field +from typing import List + +import numpy as np +import pandas as pd +import sympy +from gplearn import functions +from gplearn.functions import make_function +from gplearn.genetic import SymbolicRegressor + +from orion.algo.space import Categorical, Space +from orion.analysis.base import flatten_numpy, flatten_params, to_numpy, train_regressor +from orion.core.utils import format_trials +from orion.core.worker.transformer import build_required_space + +# TODO: Add dependencies: +# gplearn +# sympy + +logger = logging.getLogger(__name__) + + +# TODO: May need to set a_max as a parameter so user can change it. +def safe_exp(x): + return np.clip(x, a_min=None, a_max=100000) + + +@dataclass +class SymbolicRegressorParams: + population_size: int = 5000 + generations: int = 20 + metric: str = "rmse" + parsimony_coefficient: float = 0.0001 + random_state: int | None = None + verbose: int = 1 + function_set: list = field( + default_factory=lambda: [ + "add", + "sub", + "mul", + "div", + "sqrt", + "log", + "sin", + "cos", + "abs", + make_function(function=safe_exp, arity=1, name="exp"), + ] + ) + + +SymbolicRegressorParams(metric="test") + + +def convert_symb(symb, space: Space, n_decimals: int = None) -> sympy.core.expr: + """ + Convert a fitted symbolic regression to a simplified and potentially rounded mathematical expression. + Warning: eval is used in this function, thus it should not be used on unsanitized input (see + https://docs.sympy.org/latest/modules/core.html?highlight=eval#module-sympy.core.sympify). + + Parameters + ---------- + symb: Fitted symbolic regressor to find a simplified expression for. + n_dim: Number of input dimensions. If input has only a single dimension, X0 in expression is exchanged by x. + n_decimals: If set, round floats in the expression to this number of decimals. + + Returns + ------- + symb_conv: Converted mathematical expression. + """ + + # sqrt is protected function in gplearn, always returning sqrt(abs(x)) + sqrt_pos = [] + prev_sqrt_inserts = 0 + for i, f in enumerate(symb._program.program): + if isinstance(f, functions._Function) and f.name == "sqrt": + sqrt_pos.append(i) + for i in sqrt_pos: + symb._program.program.insert(i + prev_sqrt_inserts + 1, functions.abs1) + prev_sqrt_inserts += 1 + + # log is protected function in gplearn, always returning sqrt(abs(x)) + log_pos = [] + prev_log_inserts = 0 + for i, f in enumerate(symb._program.program): + if isinstance(f, functions._Function) and f.name == "log": + log_pos.append(i) + for i in log_pos: + symb._program.program.insert(i + prev_log_inserts + 1, functions.abs1) + prev_log_inserts += 1 + + symb_str = str(symb._program) + + converter = { + "sub": lambda x, y: x - y, + "div": lambda x, y: x / y, + "mul": lambda x, y: x * y, + "add": lambda x, y: x + y, + "neg": lambda x: -x, + "pow": lambda x, y: x**y, + } + + converter.update( + { + f"X{i}": sympy.symbols(dim.name, real=True) + for i, dim in enumerate(space.values()) + } + ) + + if symb._program.length_ > 300: + print( + f"Expression of length {symb._program._length} too long to convert, return raw string." + ) + return symb_str + + symb_conv = sympy.sympify( + symb_str.replace("[", "").replace("]", ""), locals=converter + ) + if n_decimals: + # Make sure also floats deeper in the expression tree are rounded + for a in sympy.preorder_traversal(symb_conv): + if isinstance(a, sympy.core.numbers.Float): + symb_conv = symb_conv.subs(a, round(a, n_decimals)) + + return symb_conv + + +def simplify_formula( + symb_model: SymbolicRegressor, space: Space, n_decimals: int = 3 +) -> str: + conv_expr = convert_symb(symb_model, space=space, n_decimals=n_decimals) + + return conv_expr + + +def symbolic_explanation( + trials: pd.DataFrame, + space: Space, + params: List[str] | None = None, + model="RandomForestRegressor", + n_samples: int = 50, + sampling_seed: int | None = None, + timeout: float = 300, + symbolic_regressor_params: SymbolicRegressorParams = SymbolicRegressorParams(), + **kwargs, +) -> SymbolicRegressor: + """ + Calculates a symbolic explanation of the effect of parameters + on the objective based on a collection of + :class:`orion.core.worker.trial.Trial`. + + For more information on the method, + see original paper at https://openreview.net/forum?id=JQwAc91sg_x. + + Segel, Sarah, et al. "Symbolic explanations for hyperparameter + optimization." AutoML Conference 2023. + + Parameters + ---------- + trials: DataFrame or dict + A dataframe of trials containing, at least, the columns 'objective' and 'id'. Or a dict + equivalent. + + space: Space object + A space object from an experiment. + + params: list of str, optional + The parameters to include in the computation. All parameters are included by default. + + model: str + Name of the regression model to use. Can be one of + - AdaBoostRegressor + - BaggingRegressor + - ExtraTreesRegressor + - GradientBoostingRegressor + - RandomForestRegressor (Default) + + n_samples: int + Number of samples to randomly generate for fitting the surrogate model. + Default is 50. + + sampling_seed: int + Seed used to sample the points for fitting the surrogate model. + + timeout: float + Number of seconds before the evolutionary algorithm is stopped. + + symbolic_regressor_params: SymbolicRegressorParams + Dataclass for the parameters for the ``SymbolicRegressor``. + + **kwargs + Arguments for the regressor model. + + Returns + ------- + SymbolicRegressor + A SymbolicRegressor fitted on the trials. + """ + + # TODO: Need to handle multi-fidelity. Maybe only pick the highest fidelity... + + # TODO: Validate that the history (`trials`) is large enough to make sense compared to the number of sampled + # points from the surrogate. + + if any(isinstance(dim, Categorical) for dim in space): + raise ValueError( + "Symbolic explanation does not support categorical dimensions yet." + ) + + params = flatten_params(space, params) + + flattened_space = build_required_space( + space, + dist_requirement="linear", + type_requirement="numerical", + shape_requirement="flattened", + ) + + if trials.empty or trials.shape[0] == 0: + return {} + + data = to_numpy(trials, space) + data = flatten_numpy(data, flattened_space) + model = train_regressor(model, data, **kwargs) + + # Sample random points for fitting the symbolic regressor. + data = [ + format_trials.trial_to_tuple(trial, flattened_space) + for trial in flattened_space.sample(n_samples, seed=sampling_seed) + ] + X_train = pd.DataFrame(data, columns=flattened_space.keys()).to_numpy() + + Y_train = model.predict(X_train) + + symbolic_regressor = SymbolicRegressor( + population_size=symbolic_regressor_params.population_size, + generations=symbolic_regressor_params.generations, + function_set=symbolic_regressor_params.function_set, + metric=symbolic_regressor_params.metric, + parsimony_coefficient=symbolic_regressor_params.parsimony_coefficient, + verbose=symbolic_regressor_params.verbose, + random_state=symbolic_regressor_params.random_state, + ) + + symbolic_regressor.fit(X_train, Y_train) + + return symbolic_regressor diff --git a/src/orion/plotting/backend_plotly.py b/src/orion/plotting/backend_plotly.py index 5e1fa5079..7dd63344e 100644 --- a/src/orion/plotting/backend_plotly.py +++ b/src/orion/plotting/backend_plotly.py @@ -10,11 +10,16 @@ import pandas as pd import plotly.express as px import plotly.graph_objects as go +import sympy from plotly.subplots import make_subplots import orion.analysis import orion.analysis.base from orion.algo.space import Categorical, Fidelity +from orion.analysis.symbolic_explanation import ( + SymbolicRegressorParams, + simplify_formula, +) from orion.core.worker.transformer import build_required_space @@ -194,7 +199,6 @@ def reformat_competitions(experiments): return competitions def build_groups(competitions): - if not isinstance(competitions, dict): rankings = [] for competition in competitions: @@ -488,6 +492,263 @@ def _plot_scatter(x, y, df): return fig +def symbolic_explanation( + experiment, + with_evc_tree=True, + params=None, + smoothing=0.85, + n_grid_points=10, + n_samples=50, + colorscale="Blues", + model="RandomForestRegressor", + model_kwargs=None, + timeout=300, + n_decimals=3, + sampling_seed=None, + symbolic_regressor_params=SymbolicRegressorParams(), + verbose_hover=True, +): + """Plotly implementation of `orion.plotting.symbolic_explanation`""" + + def build_data(): + """Builds the dataframe for the plot""" + df = experiment.to_pandas(with_evc_tree=with_evc_tree) + + names = list(experiment.space.keys()) + df["params"] = df[names].apply(_format_hyperparameters, args=(names,), axis=1) + + df = df.loc[df["status"] == "completed"] + symbolic_regressor = orion.analysis.symbolic_explanation( + trials=df, + space=experiment.space, + params=params, + model=model, + n_samples=n_samples, + timeout=timeout, + n_decimals=n_decimals, + sampling_seed=sampling_seed, + symbolic_regressor_params=symbolic_regressor_params, + **model_kwargs, + ) + + data = orion.analysis.partial_dependency( + experiment.to_pandas(), + experiment.space, + model=symbolic_regressor, + n_grid_points=n_grid_points, + ) + + latex_formula = sympy.latex( + simplify_formula( + symbolic_regressor, experiment.space, n_decimals=n_decimals + ) + ) + + df = _flatten_dims(df, experiment.space) + return (df, data, latex_formula) + + def _set_scale(figure, dims, x, y): + for axis, dim in zip("xy", dims): + if "reciprocal" in dim.prior_name or dim.type == "fidelity": + getattr(figure, f"update_{axis}axes")(type="log", row=y, col=x) + + def _plot_marginalized_avg(data, x_name): + return go.Scatter( + x=data[0][x_name], + y=data[1], + mode="lines", + name=None, + showlegend=False, + line=dict( + color=px.colors.qualitative.D3[0], + ), + ) + + def _plot_marginalized_std(data, x_name): + return go.Scatter( + x=list(data[0][x_name]) + list(data[0][x_name])[::-1], + y=list(data[1] - data[2]) + list(data[1] + data[2])[::-1], + mode="lines", + name=None, + fill="toself", + showlegend=False, + line=dict( + color=px.colors.qualitative.D3[0], + width=0, + ), + ) + + def _plot_contour(data, x_name, y_name): + return go.Contour( + x=data[0][x_name], + y=data[0][y_name], + z=data[1], + connectgaps=True, + # Share the same color range across contour plots + coloraxis="coloraxis", + line_smoothing=smoothing, + # To show labels + contours=dict( + coloring="heatmap", + showlabels=True, # show labels on contours + labelfont=dict( # label font properties + size=12, + color="white", + ), + ), + ) + + def _plot_scatter(x, y, df): + return go.Scatter( + x=x, + y=y, + marker={ + "line": {"width": 0.5, "color": "Grey"}, + "color": "black", + "size": 5, + }, + mode="markers", + opacity=0.5, + showlegend=False, + customdata=list(zip(df["id"], df["suggested"], df["params"])), + hovertemplate=_template_trials(verbose_hover), + ) + + if model_kwargs is None: + model_kwargs = {} + + df, data, latex_formula = build_data() + + if not data: + return go.Figure() + + params = [ + param_names for param_names in data.keys() if isinstance(param_names, str) + ] + + flattened_space = build_required_space( + experiment.space, + shape_requirement="flattened", + ) + + fig = make_subplots( + rows=len(params), + cols=len(params), + shared_xaxes=True, + shared_yaxes=False, + ) + + fig.update_layout(paper_bgcolor="rgba(0,0,0,0)", plot_bgcolor="rgba(0,0,0,0)") + cmin = float("inf") + cmax = -float("inf") + + for x_i in range(len(params)): + x_name = params[x_i] + fig.add_trace( + _plot_marginalized_avg(data[x_name], x_name), + row=x_i + 1, + col=x_i + 1, + ) + fig.add_trace( + _plot_marginalized_std(data[x_name], x_name), + row=x_i + 1, + col=x_i + 1, + ) + fig.add_trace( + _plot_scatter(df[x_name], df["objective"], df), + row=x_i + 1, + col=x_i + 1, + ) + + _set_scale(fig, [flattened_space[x_name]], x_i + 1, x_i + 1) + + fig.update_xaxes(title_text=x_name, row=len(params), col=x_i + 1) + if x_i > 0: + fig.update_yaxes(title_text=x_name, row=x_i + 1, col=1) + else: + fig.update_yaxes(title_text="Objective", row=x_i + 1, col=x_i + 1) + + for y_i in range(x_i + 1, len(params)): + y_name = params[y_i] + fig.add_trace( + _plot_contour( + data[(x_name, y_name)], + x_name, + y_name, + ), + row=y_i + 1, + col=x_i + 1, + ) + fig.add_trace( + _plot_scatter(df[x_name], df[y_name], df), + row=y_i + 1, + col=x_i + 1, + ) + + cmin = min(cmin, data[(x_name, y_name)][1].min()) + cmax = max(cmax, data[(x_name, y_name)][1].max()) + + _set_scale( + fig, + [flattened_space[name] for name in [x_name, y_name]], + x_i + 1, + y_i + 1, + ) + + for x_i in range(len(params)): + plot_id = len(params) * x_i + x_i + 1 + if plot_id > 1: + key = f"yaxis{plot_id}_range" + else: + key = "yaxis_range" + fig.update_layout(**{key: [cmin, cmax]}) + + fig.update_layout( + title=f"Partial dependencies on symbolic explanation\nfor experiment '{experiment.name}'", + ) + fig.layout.coloraxis.colorbar.title = "Objective" + + fig.update_layout(coloraxis=dict(colorscale=colorscale), showlegend=False) + + fmin = min([data[params[0]][1].min(), df["objective"].min()]) + fmax = max([data[params[0]][1].max(), df["objective"].max()]) + + fig.update_layout( + **{ + f"yaxis{len(experiment.space)}_visible": False, + f"yaxis{len(experiment.space)}_range": [fmin, fmax], + } + ) + + last_param = list(experiment.space.values())[-1] + xlims = last_param = last_param.interval() + fig.add_annotation( + x=sum(xlims) / 2, + y=fmax, + font=dict(size=24), + height=30, + text="Best formula found", + showarrow=False, + align="right", + valign="bottom", + col=0, + row=1, + ) + fig.add_annotation( + x=sum(xlims) / 2, + y=(fmax - fmin) / (len(params) * 0.75 + 0.5) + fmin, + font=dict(size=24), + text=r"${}$".format(latex_formula), + showarrow=False, + align="right", + valign="top", + col=0, + row=1, + ) + + return fig + + def regret( experiment, with_evc_tree=True, order_by="suggested", verbose_hover=True, **kwargs ): diff --git a/src/orion/plotting/base.py b/src/orion/plotting/base.py index 8694be0c0..4bb4381a7 100644 --- a/src/orion/plotting/base.py +++ b/src/orion/plotting/base.py @@ -3,6 +3,7 @@ ============================= """ import orion.plotting.backend_plotly as backend +from orion.analysis.symbolic_explanation import SymbolicRegressorParams def lpi( @@ -188,6 +189,94 @@ def partial_dependencies( ) +def symbolic_explanation( + experiment, + with_evc_tree=True, + params=None, + smoothing=0.85, + verbose_hover=True, + n_grid_points=10, + n_samples=50, + colorscale="Blues", + model="RandomForestRegressor", + model_kwargs=None, + timeout=300, + n_decimals=3, + sampling_seed=None, + symbolic_regressor_params=SymbolicRegressorParams(), +): + """ + Make contour plots to visualize the search space of each combination of params. + + Parameters + ---------- + experiment: ExperimentClient or Experiment + The orion object containing the experiment data + + with_evc_tree: bool, optional + Fetch all trials from the EVC tree. + Default: True + + params: list of str, optional + Indicates the parameters to include in the plots. All parameters are included by default. + + smoothing: float, optional + Smoothing applied to the countor plot. 0 corresponds to no smoothing. Default is 0.85. + + verbose_hover: bool + Indicates whether to display the hyperparameter in hover tooltips. True by default. + + colorscale: str, optional + The colorscale used for the contour plots. Supported values depends on the backend. + Default is 'Blues'. + + n_grid_points: int, optional + Number of points in the grid to compute partial dependency. Default is 10. + + n_samples: int, optional + Number of samples to randomly generate the grid used to compute the partial dependency. + Default is 50. + + model: str + Name of the regression model to use. Can be one of + - AdaBoostRegressor + - BaggingRegressor + - ExtraTreesRegressor + - GradientBoostingRegressor + - RandomForestRegressor (Default) + + model_kwargs: dict, optional + Arguments for the regressor model. + + Returns + ------- + plotly.graph_objects.Figure + + Raises + ------ + ValueError + If no experiment is provided. + + """ + + return backend.symbolic_explanation( + experiment, + with_evc_tree=with_evc_tree, + params=params, + smoothing=smoothing, + verbose_hover=verbose_hover, + n_grid_points=n_grid_points, + n_samples=n_samples, + colorscale=colorscale, + model=model, + model_kwargs=model_kwargs, + timeout=timeout, + n_decimals=n_decimals, + sampling_seed=sampling_seed, + symbolic_regressor_params=symbolic_regressor_params, + ) + + def rankings(experiments, with_evc_tree=True, order_by="suggested", **kwargs): """ Make a plot to visually compare the ranking of different hyper-optimization processes. @@ -434,6 +523,7 @@ def regrets(experiments, with_evc_tree=True, order_by="suggested", **kwargs): "lpi": lpi, "parallel_coordinates": parallel_coordinates, "partial_dependencies": partial_dependencies, + "symbolic_explanation": symbolic_explanation, "regret": regret, "regrets": regrets, "rankings": rankings, @@ -446,6 +536,7 @@ def regrets(experiments, with_evc_tree=True, order_by="suggested", **kwargs): "lpi": lpi, "parallel_coordinates": parallel_coordinates, "partial_dependencies": partial_dependencies, + "symbolic_explanation": symbolic_explanation, "regret": regret, } @@ -507,6 +598,11 @@ def partial_dependencies(self, **kwargs): __doc__ = partial_dependencies.__doc__ return self(kind="partial_dependencies", **kwargs) + def symbolic_explanation(self, **kwargs): + """Compute symbolic explanation and make partial dependency plots based on the formula found.""" + __doc__ = partial_dependencies.__doc__ + return self(kind="symbolic_explanation", **kwargs) + def regret(self, **kwargs): """Make a plot to visualize the performance of the hyper-optimization process.""" __doc__ = regret.__doc__