Source code for ax.plot.trace

#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

from typing import Any, Dict, List, Optional, Tuple, Union

import numpy as np
import plotly.graph_objs as go
from ax.plot.base import AxPlotConfig, AxPlotTypes
from ax.plot.color import COLORS, DISCRETE_COLOR_SCALE, rgba


# type aliases
Traces = List[Dict[str, Any]]


[docs]def mean_trace_scatter( y: np.ndarray, trace_color: Tuple[int] = COLORS.STEELBLUE.value, legend_label: str = "mean", hover_labels: Optional[List[str]] = None, ) -> go.Scatter: """Creates a graph object for trace of the mean of the given series across runs. Args: y: (r x t) array with results from r runs and t trials. trace_color: tuple of 3 int values representing an RGB color. Defaults to blue. legend_label: label for this trace. hover_labels: optional, text to show on hover; list where the i-th value corresponds to the i-th value in the value of the `y` argument. Returns: go.Scatter: plotly graph object """ return go.Scatter( name=legend_label, legendgroup=legend_label, x=np.arange(1, y.shape[1] + 1), y=np.mean(y, axis=0), mode="lines", line={"color": rgba(trace_color)}, fillcolor=rgba(trace_color, 0.3), fill="tonexty", text=hover_labels, )
[docs]def sem_range_scatter( y: np.ndarray, trace_color: Tuple[int] = COLORS.STEELBLUE.value, legend_label: str = "", ) -> Tuple[go.Scatter, go.Scatter]: """Creates a graph object for trace of mean +/- 2 SEMs for y, across runs. Args: y: (r x t) array with results from r runs and t trials. trace_color: tuple of 3 int values representing an RGB color. Defaults to blue. legend_label: Label for the legend group. Returns: Tuple[go.Scatter]: plotly graph objects for lower and upper bounds """ mean = np.mean(y, axis=0) sem = np.std(y, axis=0) / np.sqrt(y.shape[0]) return ( go.Scatter( x=np.arange(1, y.shape[1] + 1), y=mean - 2 * sem, legendgroup=legend_label, mode="lines", line={"width": 0}, showlegend=False, hoverinfo="none", ), go.Scatter( x=np.arange(1, y.shape[1] + 1), y=mean + 2 * sem, legendgroup=legend_label, mode="lines", line={"width": 0}, fillcolor=rgba(trace_color, 0.3), fill="tonexty", showlegend=False, hoverinfo="none", ), )
[docs]def optimum_objective_scatter( optimum: float, num_iterations: int, optimum_color: Tuple[int] = COLORS.ORANGE.value ) -> go.Scatter: """Creates a graph object for the line representing optimal objective. Args: optimum: value of the optimal objective num_iterations: how many trials were in the optimization (used to determine the width of the plot) trace_color: tuple of 3 int values representing an RGB color. Defaults to orange. Returns: go.Scatter: plotly graph objects for the optimal objective line """ return go.Scatter( x=[1, num_iterations], y=[optimum] * 2, mode="lines", line={"dash": "dash", "color": rgba(optimum_color)}, name="Optimum", )
[docs]def model_transitions_scatter( model_transitions: List[int], y_range: List[float], generator_change_color: Tuple[int] = COLORS.TEAL.value, ) -> List[go.Scatter]: """Creates a graph object for the line(s) representing generator changes. Args: model_transitions: iterations, before which generators changed y_range: upper and lower values of the y-range of the plot generator_change_color: tuple of 3 int values representing an RGB color. Defaults to orange. Returns: go.Scatter: plotly graph objects for the lines representing generator changes """ if len(y_range) != 2: # pragma: no cover raise ValueError("y_range should have two values, lower and upper.") data: List[go.Scatter] = [] for change in model_transitions: data.append( go.Scatter( x=[change] * 2, y=y_range, mode="lines", line={"dash": "dash", "color": rgba(generator_change_color)}, name="Generator change", ) ) return data
[docs]def optimization_trace_single_method( y: np.ndarray, optimum: Optional[float] = None, model_transitions: Optional[List[int]] = None, title: str = "", ylabel: str = "", hover_labels: Optional[List[str]] = None, trace_color: Tuple[int] = COLORS.STEELBLUE.value, optimum_color: Tuple[int] = COLORS.ORANGE.value, generator_change_color: Tuple[int] = COLORS.TEAL.value, ) -> AxPlotConfig: """Plots an optimization trace with mean and 2 SEMs Args: y: (r x t) array; result to plot, with r runs and t trials optimum: value of the optimal objective model_transitions: iterations, before which generators changed title: title for this plot. ylabel: label for the Y-axis. hover_labels: optional, text to show on hover; list where the i-th value corresponds to the i-th value in the value of the `y` argument. trace_color: tuple of 3 int values representing an RGB color. Defaults to orange. optimum_color: tuple of 3 int values representing an RGB color. Defaults to orange. generator_change_color: tuple of 3 int values representing an RGB color. Defaults to orange. Returns: AxPlotConfig: plot of the optimization trace with IQR """ trace = mean_trace_scatter(y=y, trace_color=trace_color, hover_labels=hover_labels) lower, upper = sem_range_scatter(y=y, trace_color=trace_color) layout = go.Layout( title=title, showlegend=True, yaxis={"title": ylabel}, xaxis={"title": "Iteration"}, ) data = [lower, trace, upper] if optimum is not None: data.append( optimum_objective_scatter( optimum=optimum, num_iterations=y.shape[1], optimum_color=optimum_color ) ) if model_transitions is not None: # pragma: no cover y_lower = np.min(np.percentile(y, 25, axis=0)) y_upper = np.max(np.percentile(y, 75, axis=0)) if optimum is not None and optimum < y_lower: y_lower = optimum if optimum is not None and optimum > y_upper: y_upper = optimum data.extend( model_transitions_scatter( model_transitions=model_transitions, y_range=[y_lower, y_upper], generator_change_color=generator_change_color, ) ) return AxPlotConfig( data=go.Figure(layout=layout, data=data), plot_type=AxPlotTypes.GENERIC )
[docs]def optimization_trace_all_methods( y_dict: Dict[str, np.ndarray], optimum: Optional[float] = None, title: str = "", ylabel: str = "", hover_labels: Optional[List[str]] = None, trace_colors: List[Tuple[int]] = DISCRETE_COLOR_SCALE, optimum_color: Tuple[int] = COLORS.ORANGE.value, ) -> AxPlotConfig: """Plots a comparison of optimization traces with 2-SEM bands for multiple methods on the same problem. Args: y: a mapping of method names to (r x t) arrays, where r is the number of runs in the test, and t is the number of trials. optimum: value of the optimal objective. title: title for this plot. ylabel: label for the Y-axis. hover_labels: optional, text to show on hover; list where the i-th value corresponds to the i-th value in the value of the `y` argument. trace_colors: tuples of 3 int values representing RGB colors to use for different methods shown in the combination plot. Defaults to Ax discrete color scale. optimum_color: tuple of 3 int values representing an RGB color. Defaults to orange. Returns: AxPlotConfig: plot of the comparison of optimization traces with IQR """ data: List[go.Scatter] = [] for i, (method, y) in enumerate(y_dict.items()): # If there are more traces than colors, start reusing colors. color = trace_colors[i % len(trace_colors)] trace = mean_trace_scatter(y=y, trace_color=color, legend_label=method) lower, upper = sem_range_scatter(y=y, trace_color=color, legend_label=method) data.extend([lower, trace, upper]) if optimum is not None: num_iterations = max(y.shape[1] for y in y_dict.values()) data.append( optimum_objective_scatter( optimum=optimum, num_iterations=num_iterations, optimum_color=optimum_color, ) ) layout = go.Layout( title=title, showlegend=True, yaxis={"title": ylabel}, xaxis={"title": "Iteration"}, ) return AxPlotConfig( data=go.Figure(layout=layout, data=data), plot_type=AxPlotTypes.GENERIC )
[docs]def optimization_times( fit_times: Dict[str, List[float]], gen_times: Dict[str, List[float]], title: str = "", ) -> AxPlotConfig: """Plots wall times for each method as a bar chart. Args: fit_times: A map from method name to a list of the model fitting times. gen_times: A map from method name to a list of the gen times. title: Title for this plot. Returns: AxPlotConfig with the plot """ # Compute means and SEs methods = list(fit_times.keys()) fit_res: Dict[str, Union[str, List[float]]] = {"name": "Fitting"} fit_res["mean"] = [np.mean(fit_times[m]) for m in methods] fit_res["2sems"] = [ 2 * np.std(fit_times[m]) / np.sqrt(len(fit_times[m])) for m in methods ] gen_res: Dict[str, Union[str, List[float]]] = {"name": "Generation"} gen_res["mean"] = [np.mean(gen_times[m]) for m in methods] gen_res["2sems"] = [ 2 * np.std(gen_times[m]) / np.sqrt(len(gen_times[m])) for m in methods ] total_mean: List[float] = [] total_2sems: List[float] = [] for m in methods: totals = np.array(fit_times[m]) + np.array(gen_times[m]) total_mean.append(np.mean(totals)) total_2sems.append(2 * np.std(totals) / np.sqrt(len(totals))) total_res: Dict[str, Union[str, List[float]]] = { "name": "Total", "mean": total_mean, "2sems": total_2sems, } # Construct plot data: List[go.Bar] = [] for i, res in enumerate([fit_res, gen_res, total_res]): data.append( go.Bar( x=methods, y=res["mean"], text=res["name"], textposition="auto", error_y={"type": "data", "array": res["2sems"], "visible": True}, marker={ "color": rgba(DISCRETE_COLOR_SCALE[i]), "line": {"color": "rgb(0,0,0)", "width": 1.0}, }, opacity=0.6, name=res["name"], ) ) layout = go.Layout( title=title, showlegend=False, yaxis={"title": "Time"}, xaxis={"title": "Method"}, ) return AxPlotConfig( data=go.Figure(layout=layout, data=data), plot_type=AxPlotTypes.GENERIC )