diff --git a/doc/source/api/diagnostics.rst b/doc/source/api/diagnostics.rst index 432b4d52..8de5acc2 100644 --- a/doc/source/api/diagnostics.rst +++ b/doc/source/api/diagnostics.rst @@ -7,3 +7,4 @@ Diagnostics :toctree: _autosummary ExplorationDiagnostics + AxModelManager diff --git a/doc/source/user_guide/advanced_usage/build_gp_surrogates.ipynb b/doc/source/user_guide/advanced_usage/build_gp_surrogates.ipynb new file mode 100644 index 00000000..8e3229e5 --- /dev/null +++ b/doc/source/user_guide/advanced_usage/build_gp_surrogates.ipynb @@ -0,0 +1,294 @@ +{ + "cells": [ + { + "cell_type": "raw", + "metadata": { + "raw_mimetype": "text/restructuredtext" + }, + "source": [ + "Building GP surrogate models from optimization data\n", + "===================================================\n", + "\n", + "The :class:`~optimas.diagnostics.ExplorationDiagnostics` class\n", + "provides a simple way of fitting a Gaussian process (GP) model to any of the\n", + "objectives or analyzed parameters of an ``optimas``\n", + ":class:`~optimas.explorations.Exploration`, independently of which generator\n", + "was used. This is useful to get a better understanding of the underlying\n", + "function, make predictions, etc.\n", + "\n", + "In this example, we will illustrate how to build GP models by using\n", + "a basic optimization that runs directly on\n", + "a Jupyter notebook. This optimization uses an\n", + ":class:`~optimas.generators.RandomSamplingGenerator` and a simple\n", + ":class:`~optimas.evaluators.FunctionEvaluator` that evaluates an\n", + "analytical function.\n", + "\n", + "\n", + "Set up example optimization\n", + "~~~~~~~~~~~~~~~~~~~~~~~~~~~\n", + "\n", + "The following cell sets up and runs an optimization with two input parameters\n", + "``x1`` and ``x2``, two objectives ``f1`` and ``f2``, and one additional\n", + "analyzed parameter ``p1``.\n", + "At each evaluation, the ``eval_func_sf_moo`` function is run,\n", + "which assigns a value to each outcome parameter according to the analytical\n", + "formulas\n", + "\n", + ".. math::\n", + "\n", + " f_1(x_1, x_2) = -(x_1 + 10 \\cos(x_1)) (x_2 + 5\\cos(x_2))\n", + "\n", + ".. math::\n", + "\n", + " f_2(x_1, x_2) = 2 f_1(x_1, x_2)\n", + "\n", + ".. math::\n", + "\n", + " p_1(x_1, x_2) = \\sin(x_1) + \\cos(x_2)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "from optimas.explorations import Exploration\n", + "from optimas.core import VaryingParameter, Objective, Parameter\n", + "from optimas.generators import RandomSamplingGenerator\n", + "from optimas.evaluators import FunctionEvaluator\n", + "\n", + "\n", + "def eval_func_sf_moo(input_params, output_params):\n", + " \"\"\"Example multi-objective function.\"\"\"\n", + " x1 = input_params[\"x1\"]\n", + " x2 = input_params[\"x2\"]\n", + " result = -(x1 + 10 * np.cos(x1)) * (x2 + 5 * np.cos(x2))\n", + " output_params[\"f1\"] = result\n", + " output_params[\"f2\"] = result * 2\n", + " output_params[\"p1\"] = np.sin(x1) + np.cos(x2)\n", + "\n", + "\n", + "var1 = VaryingParameter(\"x1\", 0.0, 5.0)\n", + "var2 = VaryingParameter(\"x2\", -5.0, 5.0)\n", + "par1 = Parameter(\"p1\")\n", + "obj1 = Objective(\"f1\", minimize=True)\n", + "obj2 = Objective(\"f2\", minimize=False)\n", + "\n", + "gen = RandomSamplingGenerator(\n", + " varying_parameters=[var1, var2],\n", + " objectives=[obj1, obj2],\n", + " analyzed_parameters=[par1],\n", + ")\n", + "ev = FunctionEvaluator(function=eval_func_sf_moo)\n", + "exploration = Exploration(\n", + " generator=gen,\n", + " evaluator=ev,\n", + " max_evals=50,\n", + " sim_workers=1,\n", + " exploration_dir_path=\"./exploration\",\n", + " libe_comms=\"threads\", #this is only needed to run on a Jupyter notebook.\n", + ")\n", + "\n", + "# Run exploration.\n", + "exploration.run()" + ] + }, + { + "cell_type": "raw", + "metadata": { + "raw_mimetype": "text/restructuredtext" + }, + "source": [ + "Initialize diagnostics\n", + "~~~~~~~~~~~~~~~~~~~~~~\n", + "\n", + "The diagnostics class only requires the path to the exploration directory\n", + "as input parameter, or directly the ``exploration`` instance." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from optimas.diagnostics import ExplorationDiagnostics\n", + "\n", + "diags = ExplorationDiagnostics(exploration)" + ] + }, + { + "cell_type": "raw", + "metadata": { + "raw_mimetype": "text/restructuredtext" + }, + "source": [ + "Building a GP model of each objective and analyzed parameter\n", + "~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n", + "\n", + "To build a GP model, simply call\n", + ":meth:`~optimas.diagnostics.Exploration.build_gp_model` on the diagnostics,\n", + "indicating the name of the variable to which the model should be fitted.\n", + "This variable can be any ``objective`` or ``analyzed_parameter`` of the\n", + "optimization.\n", + "\n", + "Note that when building a surrogate model of an analyzed parameter, it is\n", + "required to provide a value to the ``minimize`` argument. This parameter\n", + "should therefore be ``True`` is lower values of the analyzed parameter are\n", + "better than higher values. This information is necessary, e.g., for determining\n", + "the best point in the model." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Build one model for each objective and analyzed parameter.\n", + "f1_model = diags.build_gp_model(\"f1\")\n", + "f2_model = diags.build_gp_model(\"f2\")\n", + "p1_model = diags.build_gp_model(\"p1\", minimize=False)" + ] + }, + { + "cell_type": "raw", + "metadata": { + "raw_mimetype": "text/restructuredtext" + }, + "source": [ + "Visualizing the surrogate models\n", + "~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n", + "\n", + "The models provide some basic plotting methods for easy visualization, like\n", + ":meth:`~optimas.diagnostics.AxModelManager.plot_contour`\n", + "and :meth:`~optimas.diagnostics.AxModelManager.plot_slice`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# plot model for `f1`.\n", + "fig, ax1 = f1_model.plot_contour(mode=\"both\", figsize=(6, 3), dpi=300)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# plot 1D slice of `f1`.\n", + "fig, ax1 = f1_model.plot_slice(\"x1\", figsize=(6, 3), dpi=300)" + ] + }, + { + "cell_type": "raw", + "metadata": { + "raw_mimetype": "text/restructuredtext" + }, + "source": [ + "These methods also allow more complex plot compositions to be created,\n", + "such as in the example below, by providing a ``subplot_spec`` where the plot\n", + "should be drawn." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "from matplotlib.gridspec import GridSpec\n", + "\n", + "fig = plt.figure(figsize=(10, 3), dpi=300)\n", + "gs = GridSpec(1, 3, wspace=0.4)\n", + "\n", + "# plot model for `f1`.\n", + "fig, ax1 = f1_model.plot_contour(\n", + " pcolormesh_kw={\"cmap\": \"GnBu\"},\n", + " subplot_spec=gs[0, 0],\n", + ")\n", + "\n", + "# Get and draw top 3 evaluations for `f`\n", + "df_top = diags.get_best_evaluations(top=3, objective=\"f1\")\n", + "ax1.scatter(df_top[\"x1\"], df_top[\"x2\"], c=\"red\", marker=\"x\")\n", + "\n", + "# plot model for `f2`\n", + "fig, ax2 = f2_model.plot_contour(\n", + " pcolormesh_kw={\"cmap\": \"OrRd\"},\n", + " subplot_spec=gs[0, 1],\n", + ")\n", + "\n", + "# plot model for `p1`\n", + "fig, ax3 = p1_model.plot_contour(\n", + " pcolormesh_kw={\"cmap\": \"PuBu\"},\n", + " subplot_spec=gs[0, 2],\n", + ")" + ] + }, + { + "cell_type": "raw", + "metadata": { + "raw_mimetype": "text/restructuredtext" + }, + "source": [ + "Evaluating the surrogate model\n", + "~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n", + "\n", + "In addition to plotting, it is also possible to evaluate the model at any\n", + "point by using the :meth:`~optimas.diagnostics.AxModelManager.evaluate_model`\n", + "method.\n", + "\n", + "In the example below, this method is used to evaluate the model in all the\n", + "history points to create a cross-validation plot." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Evaluate model for each point in the history\n", + "mean, sem = f1_model.evaluate_model(diags.history)\n", + "min_f, max_f = np.min(diags.history[\"f1\"]), np.max(diags.history[\"f1\"])\n", + "\n", + "# Make plot\n", + "fig, ax = plt.subplots(figsize=(5, 4), dpi=300)\n", + "ax.errorbar(diags.history[\"f1\"], mean, yerr=sem, fmt=\"o\", ms=4, label=\"Data\")\n", + "ax.plot([min_f, max_f], [min_f, max_f], color=\"k\", ls=\"--\", label=\"Ideal correlation\")\n", + "ax.set_xlabel(\"Observations\")\n", + "ax.set_ylabel(\"Model predictions\")\n", + "ax.legend(frameon=False)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "optimas_env_py11", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.5" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/doc/source/user_guide/index.rst b/doc/source/user_guide/index.rst index fbf87f6e..da3cc7e3 100644 --- a/doc/source/user_guide/index.rst +++ b/doc/source/user_guide/index.rst @@ -22,6 +22,12 @@ User guide basic_usage/analyze_output basic_usage/exploration_diagnostics +.. toctree:: + :maxdepth: 2 + :caption: Advanced usage + + advanced_usage/build_gp_surrogates + .. toctree:: :maxdepth: 1 :caption: Citation diff --git a/optimas/__init__.py b/optimas/__init__.py index 3d26edf7..3d187266 100644 --- a/optimas/__init__.py +++ b/optimas/__init__.py @@ -1 +1 @@ -__version__ = "0.4.1" +__version__ = "0.5.0" diff --git a/optimas/diagnostics/__init__.py b/optimas/diagnostics/__init__.py index 88f7b9b8..a51566a3 100644 --- a/optimas/diagnostics/__init__.py +++ b/optimas/diagnostics/__init__.py @@ -1,3 +1,4 @@ from .exploration_diagnostics import ExplorationDiagnostics +from .ax_model_manager import AxModelManager -__all__ = ["ExplorationDiagnostics"] +__all__ = ["ExplorationDiagnostics", "AxModelManager"] diff --git a/optimas/diagnostics/ax_model_manager.py b/optimas/diagnostics/ax_model_manager.py new file mode 100644 index 00000000..2ea38398 --- /dev/null +++ b/optimas/diagnostics/ax_model_manager.py @@ -0,0 +1,651 @@ +"""Contains the definition of the ExplorationDiagnostics class.""" + +from typing import Optional, Union, List, Tuple, Dict, Any, Literal + +import numpy as np +from numpy.typing import NDArray +import pandas as pd +import matplotlib.pyplot as plt +from matplotlib.figure import Figure +from matplotlib.gridspec import GridSpec, GridSpecFromSubplotSpec, SubplotSpec +from matplotlib.axes import Axes + +# Ax utilities for model building +try: + from ax.service.ax_client import AxClient + from ax.modelbridge.generation_strategy import ( + GenerationStep, + GenerationStrategy, + ) + from ax.modelbridge.registry import Models + from ax.modelbridge.torch import TorchModelBridge + from ax.core.observation import ObservationFeatures + from ax.service.utils.instantiation import ObjectiveProperties + + ax_installed = True +except ImportError: + ax_installed = False + +from optimas.core import VaryingParameter, Objective +from optimas.utils.other import convert_to_dataframe + + +class AxModelManager: + """Class for building and exploring GP models using an ``AxClient``. + + Parameters + ---------- + source : AxClient, str or DataFrame + Source data for the model. It can be either an existing ``AxClient`` + with a GP model, a string with the path to a ``json`` file with a + serialized ``AxClient``, or a pandas ``DataFrame``. + When using a ``DataFrame``, a list of objectives and varying parameters + should also be provided. + objectives : list of `Objective`, optional + Only needed if ``source`` is a pandas ``DataFrame``. List of + objectives for which a GP model should be built. The names and data of + these objectives must be contained in the source ``DataFrame``. + varying_parameters : list of `VaryingParameter`, optional + Only needed if ``source`` is a pandas ``DataFrame``. List of + parameters that were varied to scan the value of the objectives. + The names and data of these parameters must be contained in the + source ``DataFrame``. + """ + + def __init__( + self, + source: Union[AxClient, str, pd.DataFrame], + varying_parameters: Optional[List[VaryingParameter]] = None, + objectives: Optional[List[Objective]] = None, + ) -> None: + if not ax_installed: + raise ImportError( + "`AxModelManager` requires Ax to be installed. " + "You can do so by running `pip install ax-platform`." + ) + if isinstance(source, AxClient): + self.ax_client = source + elif isinstance(source, str): + self.ax_client = AxClient.load_from_json_file(filepath=source) + elif isinstance(source, pd.DataFrame): + self.ax_client = self._build_ax_client_from_dataframe( + source, varying_parameters, objectives + ) + else: + raise ValueError( + f"Wrong source type: {type(source)}. " + "The source must be an `AxClient`, a path to an AxClient json " + "file, or a pandas `DataFrame`." + ) + self.ax_client.fit_model() + + @property + def _model(self) -> TorchModelBridge: + """Get the model from the AxClient instance.""" + return self.ax_client.generation_strategy.model + + def _build_ax_client_from_dataframe( + self, + df: pd.DataFrame, + varying_parameters: List[VaryingParameter], + objectives: List[Objective], + ) -> AxClient: + """Initialize the AxClient and the model using the given data. + + Parameters + ---------- + df : DataFrame + The source pandas ``DataFrame``. + objectives : list of `Objective`. + List of objectives for which a GP model should be built. + varying_parameters : list of `VaryingParameter`. + List of parameters that were varied to scan the value of the + objectives. + """ + # Define parameters for AxClient + axparameters = [] + for par in varying_parameters: + # Determine parameter type. + value_dtype = np.dtype(par.dtype) + if value_dtype.kind == "f": + value_type = "float" + elif value_dtype.kind == "i": + value_type = "int" + else: + raise ValueError( + "Ax range parameter can only be of type 'float'ot 'int', " + "not {var.dtype}." + ) + # Create parameter dict and append to list. + axparameters.append( + { + "name": par.name, + "type": "range", + "bounds": [par.lower_bound, par.upper_bound], + "is_fidelity": par.is_fidelity, + "target_value": par.fidelity_target_value, + "value_type": value_type, + } + ) + + # Define objectives for AxClient + axobjectives = { + obj.name: ObjectiveProperties(minimize=obj.minimize) + for obj in objectives + } + + # Create Ax client. + # We need to explicitly define a generation strategy because otherwise + # a random sampling step will be set up by Ax, and this step does not + # allow calling `model.predict`. Using MOO for multiobjective is + # needed because otherwise calls to `get_pareto_optimal_parameters` + # would fail. + model = Models.GPEI if len(objectives) == 1 else Models.MOO + gs = GenerationStrategy([GenerationStep(model=model, num_trials=-1)]) + ax_client = AxClient(generation_strategy=gs, verbose_logging=False) + ax_client.create_experiment( + parameters=axparameters, objectives=axobjectives + ) + + # Add trials from DataFrame + for _, row in df.iterrows(): + params = {vp.name: row[vp.name] for vp in varying_parameters} + _, trial_id = ax_client.attach_trial(params) + data = {obj.name: (row[obj.name], np.nan) for obj in objectives} + ax_client.complete_trial(trial_id, raw_data=data) + return ax_client + + def _get_best_point(self, metric_name: Optional[str] = None) -> Dict: + """Get the best point with the best predicted model value. + + Parameters + ---------- + metric_name: str, optional. + Name of the metric to evaluate. + If not specified, it will take the first first objective in + ``self.ax_client``. + + Returns + ------- + best_point : dict + A dictionary with the parameters of the best point. + """ + _, best_point = self.get_best_evaluation( + metric_name=metric_name, use_model_predictions=True + ) + return best_point + + def _get_mid_point(self) -> Dict: + """Get the middle point of the space of parameters. + + Returns + ------- + mid_point : dict + A dictionary with the parameters of the mid point. + """ + mid_point = {} + for key, par in self.ax_client.experiment.parameters.items(): + mid_point[key] = 0.5 * (par.lower + par.upper) + + return mid_point + + def _get_arm_index( + self, + arm_name: str, + ) -> int: + """Get the index of the arm by its name. + + Parameters + ---------- + arm_name : str + Name of the arm. If not given, the best arm is selected. + + Returns + ------- + index : int + Trial index of the arm. + """ + df = self.ax_client.get_trials_data_frame() + index = df.loc[df["arm_name"] == arm_name, "trial_index"].iloc[0] + return index + + def evaluate_model( + self, + sample: Union[pd.DataFrame, Dict, NDArray] = None, + metric_name: Optional[str] = None, + fixed_parameters: Optional[Dict] = None, + ) -> Tuple[NDArray]: + """Evaluate the model over the specified sample. + + Parameters + ---------- + sample : DataFrame, dict of NDArray or NDArray + containing the data sample where to evaluate the model. + If numpy array, it must contain the values of all the model + parameters. + If DataFrame or dict, it can contain only those parameters to vary. + The rest of parameters would be set to the model best point, + unless they are further specified using ``fixed_parameters``. + metric_name : str, optional. + Name of the metric to evaluate. + If not specified, it will take the first first objective in + ``self.ax_client``. + fixed_parameters : dict, optional. + A dictionary with structure ``{param_name: param_val}`` with the + values of the parameters to be fixed in the evaluation. If a given + parameter also exists in the ``sample``, the values in the + ``sample`` will be overwritten by the fixed value. + + Returns + ------- + NDArray, NDArray + Two numpy arrays containing the mean of the model + and the standard error of the mean (sem), respectively. + """ + if metric_name is None: + metric_name = self.ax_client.objective_names[0] + else: + metric_names = list(self.ax_client.experiment.metrics.keys()) + if metric_name not in metric_names: + raise ValueError( + f"Metric name {metric_name} does not match any of the " + f"metrics. Available metrics are: {metric_names}." + ) + + parnames = list(self.ax_client.experiment.parameters.keys()) + + sample = convert_to_dataframe(sample) + + if fixed_parameters is not None: + for key, val in fixed_parameters.items(): + sample[key] = val + + # check if labels of the dataframe match the parnames + for name in parnames: + if name not in sample.columns.values: + raise ValueError(f"Data for {name} is missing in the sample.") + # make list of `ObservationFeatures` + obsf_list = [] + for i in range(sample.shape[0]): + parameters = {} + for name in parnames: + parameters[name] = sample[name].iloc[i] + obsf_list.append(ObservationFeatures(parameters=parameters)) + + mu, cov = self._model.predict(obsf_list) + m_array = np.asarray(mu[metric_name]) + sem_array = np.sqrt(cov[metric_name][metric_name]) + return m_array, sem_array + + def get_best_evaluation( + self, + metric_name: Optional[str] = None, + use_model_predictions: Optional[bool] = True, + ) -> Tuple[int, Dict]: + """Get the best scoring point in the sample. + + Parameters + ---------- + metric_name : str, optional. + Name of the metric to evaluate. + If not specified, it will take the first first objective in + ``self.ax_client``. + use_model_predictions : bool, optional. + Whether to extract the best point using model predictions + or directly observed values. + + Returns + ------- + int, dict + The index of the best evaluation and a dictionary with its + parameters. + """ + # metric name + if metric_name is None: + metric_name = self.ax_client.objective_names[0] + + # get optimum + if len(self.ax_client.objective_names) > 1: + minimize = None + for obj in self.ax_client.objective.objectives: + if metric_name == obj.metric_names[0]: + minimize = obj.minimize + break + pp = self.ax_client.get_pareto_optimal_parameters( + use_model_predictions=use_model_predictions + ) + obj_vals, param_vals, trial_indices = [], [], [] + for index, (vals, (objs, covs)) in pp.items(): + trial_indices.append(index) + param_vals.append(vals) + obj_vals.append(objs[metric_name]) + i_best = np.argmin(obj_vals) if minimize else np.argmax(obj_vals) + best_point = param_vals[i_best] + index = trial_indices[i_best] + else: + if use_model_predictions is True: + best_arm, _ = self._model.model_best_point() + best_point = best_arm.parameters + index = self._get_arm_index(best_arm.name) + else: + # AxClient.get_best_parameters seems to always return the best + # point from the observed values, independently of the value + # of `use_model_predictions`. + index, best_point, _ = self.ax_client.get_best_trial( + use_model_predictions=use_model_predictions + ) + + return index, best_point + + def plot_contour( + self, + param_x: Optional[str] = None, + param_y: Optional[str] = None, + metric_name: Optional[str] = None, + slice_values: Optional[Union[Dict, Literal["best", "mid"]]] = "mid", + n_points: Optional[int] = 200, + range_x: Optional[List[float]] = None, + range_y: Optional[List[float]] = None, + mode: Optional[Literal["mean", "sem", "both"]] = "mean", + show_trials: Optional[bool] = True, + show_contour: Optional[bool] = True, + show_contour_labels: Optional[bool] = False, + subplot_spec: Optional[SubplotSpec] = None, + gridspec_kw: Optional[Dict[str, Any]] = None, + pcolormesh_kw: Optional[Dict[str, Any]] = None, + **figure_kw, + ) -> Tuple[Figure, Union[Axes, List[Axes]]]: + """Plot a 2D slice of the surrogate model. + + Parameters + ---------- + param_x : str + Name of the parameter to plot in x axis. + param_y : str + Name of the parameter to plot in y axis. + metric_name : str, optional. + Name of the metric to plot. + If not specified, it will take the first objective in + ``self.ax_client``. + slice_values : dict or str, optional. + The values along which to slice the model, if the model has more + than two dimensions. Possible values are: ``"best"`` (slice along + the best predicted point), ``"mid"`` (slice along the middle + point of the varying parameters), or a dictionary with structure + ``{param_name: param_val}`` that contains the slice values of each + parameter. By default, ``"mid"``. + n_points : int, optional + Number of points in each axis. By default, ``200``. + range_x, range_y : list of float, optional + Range of each axis. It not given, the lower and upper boundary + of each parameter will be used. + mode : str, optional. + Whether to plot the ``"mean"`` of the model, the standard error of + the mean ``"sem"``, or ``"both"``. By default, ``"mean"``. + show_trials : bool + Whether to show the trials used to build the model. By default, + ``True``. + show_contour : bool + Whether to show the contour lines. By default, ``True``. + show_contour_labels : bool + Whether to add labels to the contour lines. By default, ``False``. + subplot_spec : SubplotSpec, optional + A matplotlib ``SubplotSpec`` in which to draw the axis. + gridspec_kw : dict, optional + Dict with keywords passed to the ``GridSpec``. + pcolormesh_kw : dict, optional + Dict with keywords passed to ``ax.pcolormesh``. + **figure_kw + Additional keyword arguments to pass to ``pyplot.figure``. + Only used if no ``subplot_spec`` is given. + + Returns + ------- + Figure, Axes or list of Axes + A matplotlib figure and either a single ``Axes`` or a list of + ``Axes`` if ``mode="both"``. + """ + # get experiment info + experiment = self.ax_client.experiment + parnames = list(experiment.parameters.keys()) + + if len(parnames) < 2: + raise RuntimeError( + "Insufficient number of parameters in data for this plot " + "(minimum 2)." + ) + + # select the input variables + if param_x is None: + param_x = parnames[0] + if param_y is None: + param_y = parnames[1] + + # metric name + if metric_name is None: + metric_name = self.ax_client.objective_names[0] + + # set the plotting range + if range_x is None: + range_x = [None, None] + if range_y is None: + range_y = [None, None] + if range_x[0] is None: + range_x[0] = experiment.parameters[param_x].lower + if range_x[1] is None: + range_x[1] = experiment.parameters[param_x].upper + if range_y[0] is None: + range_y[0] = experiment.parameters[param_y].lower + if range_y[1] is None: + range_y[1] = experiment.parameters[param_y].upper + + # get grid sample of points where to evalutate the model + xaxis = np.linspace(range_x[0], range_x[1], n_points) + yaxis = np.linspace(range_y[0], range_y[1], n_points) + X, Y = np.meshgrid(xaxis, yaxis) + sample = {param_x: X.flatten(), param_y: Y.flatten()} + + if slice_values == "mid": + # Get mid point + slice_values = self._get_mid_point() + elif slice_values == "best": + # get best point + slice_values = self._get_best_point(metric_name=metric_name) + + fixed_parameters = {} + for name, val in slice_values.items(): + if name not in [param_x, param_y]: + fixed_parameters[name] = slice_values[name] + + # evaluate the model + f_plt, sd_plt = self.evaluate_model( + sample=sample, + metric_name=metric_name, + fixed_parameters=fixed_parameters, + ) + + # select quantities to plot and set the labels + f_plots = [] + labels = [] + if mode in ["mean", "both"]: + f_plots.append(f_plt.reshape(X.shape)) + labels.append(metric_name + ", mean") + if mode in ["sem", "both"]: + f_plots.append(sd_plt.reshape(X.shape)) + labels.append(metric_name + ", sem") + + # create figure + nplots = len(f_plots) + gridspec_kw = dict(gridspec_kw or {}) + if subplot_spec is None: + fig = plt.figure(**figure_kw) + gs = GridSpec(1, nplots, **gridspec_kw) + else: + fig = plt.gcf() + gs = GridSpecFromSubplotSpec(1, nplots, subplot_spec, **gridspec_kw) + + # draw plots + trials = self.ax_client.get_trials_data_frame() + axs = [] + for i, f in enumerate(f_plots): + ax = plt.subplot(gs[i]) + # colormesh + pcolormesh_kw = dict(pcolormesh_kw or {}) + im = ax.pcolormesh(xaxis, yaxis, f, shading="auto", **pcolormesh_kw) + cbar = plt.colorbar(im, ax=ax, location="top") + cbar.set_label(labels[i]) + ax.set(xlabel=param_x, ylabel=param_y) + # contour + if show_contour: + cset = ax.contour( + X, + Y, + f, + levels=20, + linewidths=0.5, + colors="black", + linestyles="solid", + ) + if show_contour_labels: + ax.clabel( + cset, inline=True, fmt="%1.1f", fontsize="xx-small" + ) + if show_trials: + ax.scatter( + trials[param_x], trials[param_y], s=8, c="black", marker="o" + ) + ax.set_xlim(range_x) + ax.set_ylim(range_y) + axs.append(ax) + + if nplots == 1: + return fig, axs[0] + else: + return fig, axs + + def plot_slice( + self, + param_name: Optional[str] = None, + metric_name: Optional[str] = None, + slice_values: Optional[Union[Dict, Literal["best", "mid"]]] = "mid", + n_points: Optional[int] = 200, + range: Optional[List[float]] = None, + show_legend: Optional[bool] = False, + subplot_spec: Optional[SubplotSpec] = None, + gridspec_kw: Optional[Dict[str, Any]] = None, + plot_kw: Optional[Dict[str, Any]] = None, + **figure_kw, + ) -> Tuple[Figure, Axes]: + """Plot a 1D slice of the surrogate model. + + Parameters + ---------- + param_name : str + Name of the parameter to plot in x axis. If not given, the first + varying parameter will be used. + metric_name : str, optional. + Name of the metric to plot. + If not specified, it will take the first objective in + ``self.ax_client``. + slice_values : dict or str, optional. + The values along which to slice the model, if the model has more + than one dimensions. Possible values are: ``"best"`` (slice along + the best predicted point), ``"mid"`` (slice along the middle + point of the varying parameters), or a dictionary with structure + ``{param_name: param_val}`` that contains the slice values of each + parameter. By default, ``"mid"``. + n_points : int, optional + Number of points along the x axis. By default, ``200``. + range : list of float, optional + Range of the x axis. It not given, the lower and upper boundary + of the x parameter will be used. + show_legend : bool + Whether to show a legend with the fixed slice values. By default, + ``False``. + subplot_spec : SubplotSpec, optional + A matplotlib ``SubplotSpec`` in which to draw the axis. + gridspec_kw : dict, optional + Dict with keywords passed to the ``GridSpec``. + plot_kw : dict, optional + Dict with keywords passed to ``ax.plot``. + **figure_kw + Additional keyword arguments to pass to ``pyplot.figure``. Only + used if no ``subplot_spec`` is given. + + Returns + ------- + Figure, Axes + """ + # get experiment info + experiment = self.ax_client.experiment + parnames = list(experiment.parameters.keys()) + + # select the input variables + if param_name is None: + param_name = parnames[0] + + # metric name + if metric_name is None: + metric_name = self.ax_client.objective_names[0] + + # set the plotting range + if range is None: + range = [None, None] + if range[0] is None: + range[0] = experiment.parameters[param_name].lower + if range[1] is None: + range[1] = experiment.parameters[param_name].upper + + # get sample of points where to evalutate the model + sample = {param_name: np.linspace(range[0], range[1], n_points)} + + if slice_values == "mid": + # Get mid point + slice_values = self._get_mid_point() + elif slice_values == "best": + # get best point + slice_values = self._get_best_point(metric_name=metric_name) + + fixed_parameters = {} + for name, val in slice_values.items(): + if name not in [param_name]: + fixed_parameters[name] = slice_values[name] + + # evaluate the model + mean, sem = self.evaluate_model( + sample=sample, + metric_name=metric_name, + fixed_parameters=fixed_parameters, + ) + + # create figure + gridspec_kw = dict(gridspec_kw or {}) + if subplot_spec is None: + fig = plt.figure(**figure_kw) + gs = GridSpec(1, 1, **gridspec_kw) + else: + fig = plt.gcf() + gs = GridSpecFromSubplotSpec(1, 1, subplot_spec, **gridspec_kw) + + # Make plot + plot_kw = dict(plot_kw or {}) + label = "" + for par, val in fixed_parameters.items(): + if label: + label += ", " + label += f"{par} = {val}" + ax = fig.add_subplot(gs[0]) + ax.plot(sample[param_name], mean, label=label, **plot_kw) + ax.fill_between( + x=sample[param_name], + y1=mean - sem, + y2=mean + sem, + color="lightgray", + alpha=0.5, + ) + ax.set_xlabel(param_name) + ax.set_ylabel(metric_name) + if show_legend: + ax.legend(frameon=False) + + return fig, ax diff --git a/optimas/diagnostics/exploration_diagnostics.py b/optimas/diagnostics/exploration_diagnostics.py index 3229692e..54a127da 100644 --- a/optimas/diagnostics/exploration_diagnostics.py +++ b/optimas/diagnostics/exploration_diagnostics.py @@ -18,6 +18,7 @@ from optimas.evaluators.base import Evaluator from optimas.explorations import Exploration from optimas.utils.other import get_df_with_selection +from optimas.diagnostics.ax_model_manager import AxModelManager class ExplorationDiagnostics: @@ -365,16 +366,7 @@ def get_best_evaluation( Objective to consider for determining the best evaluation. Only. needed if there is more than one objective. By default ``None``. """ - if objective is None: - objective = self.objectives[0] - elif isinstance(objective, str): - objective = self._get_objective(objective) - history = self.history[self.history.sim_ended] - if objective.minimize: - i_best = np.argmin(history[objective.name]) - else: - i_best = np.argmax(history[objective.name]) - return history.iloc[[i_best]] + return self.get_best_evaluations(objective, top=1) def get_pareto_front_evaluations( self, @@ -595,7 +587,7 @@ def plot_worker_timeline( def plot_history( self, - parnames: Optional[list] = None, + parnames: Optional[List] = None, xname: Optional[str] = None, select: Optional[Dict] = None, sort: Optional[Dict] = None, @@ -897,42 +889,112 @@ def print_evaluation(self, trial_index: int) -> None: print() - def print_best_evaluations( + def get_best_evaluations( self, + objective: Optional[Union[str, Objective]] = None, top: Optional[int] = 3, + ) -> pd.DataFrame: + """Get a list with the best evaluations. + + Parameters + ---------- + objective : Objective or str, optional + Objective, or name of the objective to plot. If `None`, the first + objective of the exploration is shown. + top : int, optional + Number of top evaluations to consider (3 by default). + e.g. top = 3 means that the three best evaluations will be shown. + + Returns + ------- + top_indices : List with the indices of the best evaluations. + """ + if objective is None: + objective = self.objectives[0] + elif isinstance(objective, str): + objective = self._get_objective(objective) + return self.history.sort_values( + by=objective.name, ascending=objective.minimize + ).iloc[:top] + + def print_best_evaluations( + self, objective: Optional[Union[str, Objective]] = None, + top: Optional[int] = 3, ) -> None: """Print top evaluations according to the given objective. Parameters ---------- + objective : Objective or str, optional + Objective, or name of the objective to plot. If `None`, the first + objective of the exploration is shown. top : int, optional Number of top evaluations to consider (3 by default). e.g. top = 3 means that the three best evaluations will be shown. - objective : str, optional - Objective, or name of the objective to plot. If `None`, the first - objective of the exploration is shown. """ if objective is None: objective = self.objectives[0] - if isinstance(objective, str): + elif isinstance(objective, str): objective = self._get_objective(objective) - top_indices = list( - self.history.sort_values( - by=objective.name, ascending=objective.minimize - ).index - )[:top] + best_evals = self.get_best_evaluations(top=top, objective=objective) objective_names = [obj.name for obj in self.objectives] varpar_names = [var.name for var in self.varying_parameters] anapar_names = [var.name for var in self.analyzed_parameters] print( "Top %i evaluations in metric %s (minimize = %s): " % (top, objective.name, objective.minimize), - top_indices, + best_evals.index.to_list(), ) print() - print( - self.history.loc[top_indices][ - objective_names + varpar_names + anapar_names - ] + print(best_evals[objective_names + varpar_names + anapar_names]) + + def build_gp_model( + self, parameter: str, minimize: Optional[bool] = None + ) -> AxModelManager: + """Build a GP model of the specified parameter. + + Parameters + ---------- + parameter: str + Name of an objective or analyzed parameter for which the model + will be built. + minimize: bool, optional + Required only if `parameter` is not an objective. + Use it to indicate whether lower or higher values of the parameter + are better. This is relevant, e.g. to determine the best point of + the model. + + Returns + ------- + AxModelManager + """ + # Select objective. + try: + objective = self._get_objective(parameter) + except ValueError: + try: + analyzed_parameter = self._get_analyzed_parameter(parameter) + except ValueError: + raise ValueError( + f"Parameter {parameter} not found. " + "It is not an objective nor an analyzed parameter of the " + "Exploration." + ) + if minimize is None: + raise ValueError( + f"Please specify whether the parameter {parameter} should " + "be minimized." + ) + objective = Objective( + analyzed_parameter.name, + minimize=minimize, + dtype=analyzed_parameter.dtype, + ) + + # Initialize `AxModelManager` with history dataframe. + return AxModelManager( + source=self.history, + varying_parameters=self.varying_parameters, + objectives=[objective], ) diff --git a/optimas/utils/other.py b/optimas/utils/other.py index be918558..ef366737 100644 --- a/optimas/utils/other.py +++ b/optimas/utils/other.py @@ -57,6 +57,12 @@ def convert_to_dataframe( elif isinstance(data, pd.DataFrame): return data elif isinstance(data, dict): + # Check whether the elements in the dictionary are arrays or not. + # If they are not, covert to 1-element arrays for DataFrame initialization. + element = data[list(data.keys())[0]] + if not hasattr(element, "__len__"): + for key, value in data.items(): + data[key] = np.ones(1, dtype=type(value)) * value return pd.DataFrame(data) elif isinstance(data, list): fields = list(data[0].keys()) diff --git a/tests/test_ax_model_manager.py b/tests/test_ax_model_manager.py new file mode 100644 index 00000000..1b5fc6eb --- /dev/null +++ b/tests/test_ax_model_manager.py @@ -0,0 +1,156 @@ +import os +import numpy as np +import matplotlib.pyplot as plt +from matplotlib.gridspec import GridSpec + +from optimas.explorations import Exploration +from optimas.core import VaryingParameter, Objective +from optimas.generators import AxSingleFidelityGenerator +from optimas.evaluators import FunctionEvaluator +from optimas.diagnostics import ExplorationDiagnostics, AxModelManager + + +def eval_func_sf_moo(input_params, output_params): + """Evaluation function for multi-objective single-fidelity test""" + x0 = input_params["x0"] + x1 = input_params["x1"] + result = -(x0 + 10 * np.cos(x0)) * (x1 + 5 * np.cos(x1)) + output_params["f"] = result + output_params["f2"] = result * 2 + + +def test_ax_model_manager(): + """ + Test that an exploration with a multi-objective single-fidelity generator + runs and that the generator and Ax client are updated after running. + """ + + var1 = VaryingParameter("x0", -50.0, 5.0) + var2 = VaryingParameter("x1", -5.0, 15.0) + var3 = VaryingParameter("x2", -5.0, 15.0) + obj = Objective("f", minimize=True) + obj2 = Objective("f2", minimize=False) + + gen = AxSingleFidelityGenerator( + varying_parameters=[var1, var2, var3], objectives=[obj, obj2] + ) + ev = FunctionEvaluator(function=eval_func_sf_moo) + exploration = Exploration( + generator=gen, + evaluator=ev, + max_evals=10, + sim_workers=2, + exploration_dir_path="./tests_output/test_ax_model_manager", + ) + + # Get reference to original AxClient. + ax_client = gen._ax_client + + # Run exploration. + exploration.run() + + # Open diagnostics and extract parameter sample `df` + diags = ExplorationDiagnostics(exploration) + varpar_names = [var.name for var in diags.varying_parameters] + df = diags.history[varpar_names + ["f"]] + + # Get model manager directly from the existing `AxClient` instance. + mm_axcl = AxModelManager(source=ax_client) + mean_axcl, sem_axcl = mm_axcl.evaluate_model(sample=df, metric_name="f") + + # Get model manager from the `AxClient` dumped json file. + max_evals = exploration.max_evals + exploration_dir_path = exploration.exploration_dir_path + path = os.path.join( + exploration_dir_path, + f"model_history/ax_client_at_eval_{max_evals}.json", + ) + mm_json = AxModelManager(source=path) + mean_json, sem_json = mm_json.evaluate_model(sample=df, metric_name="f") + + # Get model manager from diagnostics data. + mm_diag = diags.build_gp_model(parameter="f") + mean_diag, sem_diag = mm_diag.evaluate_model(sample=df, metric_name="f") + + # Add model evaluations to sample and print results. + df["f_mean_axcl"] = mean_axcl + df["f_mean_json"] = mean_json + df["f_mean_diag"] = mean_diag + print(df) + + # Check that different model initializations match within a 1% tolerance. + assert np.allclose(mean_axcl, mean_json, rtol=1e-2) + assert np.allclose(mean_axcl, mean_diag, rtol=1e-2) + + # Make example figure with two models in 2D. + fig = plt.figure(figsize=(8, 8)) + gs = GridSpec(2, 2, wspace=0.2, hspace=0.3) + + # center coordinates + x1_c = 0.5 * (var2.lower_bound + var2.upper_bound) + x2_c = 0.5 * (var3.lower_bound + var3.upper_bound) + + # plot model for `f` with custom slice value + fig, ax1 = mm_axcl.plot_contour( + metric_name="f", + slice_values={"x2": x2_c}, + pcolormesh_kw={"cmap": "GnBu"}, + subplot_spec=gs[0, 0], + ) + + # Get and draw top 3 evaluations for `f` + df_top = diags.get_best_evaluations(top=3, objective="f") + ax1.scatter(df_top["x0"], df_top["x1"], c="red", marker="x") + + # plot model for `f` with default settings (mid point) + fig, ax1 = mm_axcl.plot_contour( + metric_name="f", + subplot_spec=gs[0, 1], + ) + + # plot model for `f2` with custom slice value + fig, ax2 = mm_axcl.plot_contour( + metric_name="f2", + slice_values={"x2": x2_c}, + pcolormesh_kw={"cmap": "OrRd"}, + subplot_spec=gs[1, 0], + ) + + # plot model for `f2` along best slice + fig, ax2 = mm_axcl.plot_contour( + metric_name="f2", + slice_values="best", + pcolormesh_kw={"cmap": "OrRd"}, + subplot_spec=gs[1, 1], + ) + + # Get and draw top 3 evaluations for `f2` + df2_top = diags.get_best_evaluations(top=3, objective="f2") + ax2.scatter(df2_top["x0"], df2_top["x1"], c="blue", marker="x") + plt.savefig(os.path.join(exploration_dir_path, "models_2d.png")) + + fig, axs = mm_axcl.plot_contour(mode="both", figsize=(8, 4)) + fig.savefig("models_2d_both.png") + + # Make figure of the models in 1D with errors. + fig = plt.figure() + gs = GridSpec(2, 1, hspace=0.3) + fig, ax = mm_axcl.plot_slice( + "x0", + metric_name="f", + slice_values={"x1": x1_c, "x2": x2_c}, + subplot_spec=gs[0], + plot_kw={"color": "C0"}, + ) + fig, ax = mm_axcl.plot_slice( + "x0", + metric_name="f2", + slice_values={"x1": x1_c, "x2": x2_c}, + subplot_spec=gs[1], + plot_kw={"color": "C1"}, + ) + fig.savefig(os.path.join(exploration_dir_path, "models_1d.png")) + + +if __name__ == "__main__": + test_ax_model_manager()