From 73c85a640f114199e72d41d15446ebd2a68a647b Mon Sep 17 00:00:00 2001
From: delaossa <alberto.martinez.de.la.ossa@desy.de>
Date: Fri, 9 Feb 2024 14:22:52 +0100
Subject: [PATCH 01/54] Add surrogate models analysis capabilities through the
 `AxClient` class.

---
 optimas/diagnostics/__init__.py               |   3 +-
 optimas/diagnostics/ax_model_manager.py       | 347 ++++++++++++++++++
 .../diagnostics/exploration_diagnostics.py    | 121 +++++-
 3 files changed, 461 insertions(+), 10 deletions(-)
 create mode 100644 optimas/diagnostics/ax_model_manager.py

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..c7c2177c
--- /dev/null
+++ b/optimas/diagnostics/ax_model_manager.py
@@ -0,0 +1,347 @@
+"""Contains the definition of the ExplorationDiagnostics class."""
+
+from typing import Optional, Union, List, Tuple, Dict, Any, Literal
+import numpy as np
+import pandas as pd
+from copy import deepcopy
+import matplotlib.pyplot as plt
+from matplotlib.gridspec import GridSpec, GridSpecFromSubplotSpec, SubplotSpec
+from matplotlib.axes import Axes
+
+# Ax utilities for model building
+from ax.service.ax_client import AxClient
+from ax.modelbridge.generation_strategy import (
+    GenerationStep, GenerationStrategy)
+from ax.modelbridge.registry import Models
+from ax.modelbridge.factory import get_GPEI
+from ax.modelbridge.torch import TorchModelBridge
+from ax.core.observation import ObservationFeatures
+
+
+class AxModelManager(object):
+
+    def __init__(self, source: Union[AxClient, str, pd.DataFrame]) -> None:
+        """Utilities for building and exploring surrogate models using ``Ax``.
+
+        Parameter:
+        ----------
+        source: AxClient, str or DataFrame
+            Source of data from where to obtain the model.
+            If ``DataFrame``, the model has to be build using ``build_model``.
+            If ``AxClient``, it uses the model in there.
+            If ``str``, it should contain the path to an ``AxClient`` json file.
+        """
+
+        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.df = source
+            self.ax_client = None
+        else:
+            raise RuntimeError("Wrong source.")
+
+        if self.ax_client:
+            # fit GP model
+            self.ax_client.generation_strategy._fit_current_model(data=None)
+
+    @property
+    def model(self) -> TorchModelBridge:
+        """Get the model from the AxClient instance."""
+        return self.ax_client.generation_strategy.model
+
+    def build_model(
+        self, 
+        parnames: List[str], 
+        objname: str, 
+        minimize: Optional[bool] = True
+    ) -> None:
+        """Initialize the AxClient using the given data, 
+            the model parameters and the metric.
+            Then it fits a Gaussian Process model to the data.
+
+        Parameter:
+        ----------
+        parnames: list of string
+            List with the names of the parameters of the model
+        objname: string
+            Name of the objective
+        minimize: bool
+            Whether to minimize or maximize the objective.
+            Only relevant to establish the best point.
+        """
+        parameters = [{"name": p_name,
+                       "type": "range",
+                       "bounds": [self.df[p_name].min(), self.df[p_name].max()],
+                       "value_type": "float"
+                       } for p_name in parnames]
+
+        # create Ax client
+        gs = GenerationStrategy([GenerationStep(model=Models.GPEI, num_trials=-1)])
+        self.ax_client = AxClient(generation_strategy=gs, verbose_logging=False)
+        self.ax_client.create_experiment(
+            name="optimas_data",
+            parameters=parameters,
+            objective_name=objname,
+            minimize=minimize,
+        )
+
+        # adds data
+        metric_name = list(self.ax_client.experiment.metrics.keys())[0]
+        for index, row in self.df.iterrows():
+            params = {p_name: row[p_name] for p_name in parnames}
+            _, trial_id = self.ax_client.attach_trial(params)
+            self.ax_client.complete_trial(trial_id, 
+                                          {metric_name: (row[metric_name], np.nan)})
+
+        # fit GP model
+        self.ax_client.generation_strategy._fit_current_model(data=None)
+
+    def evaluate_model(
+        self, 
+        sample: Union[pd.DataFrame, Dict, np.ndarray] = None, 
+        p0: Dict = None,
+    ) -> Tuple[np.ndarray]:
+        """Evaluate the model over the specified sample.
+
+        Parameter:
+        ----------
+        sample: DataFrame, dict of arrays or numpy array,
+            containing the data sample where to evaluate the model.
+            If numpy array, it must contain the values of all the model parameres.
+            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 ``p0``.
+        p0: dictionary
+            Particular values of parameters to be fixed for the evaluation 
+            over the givensample.
+
+        Returns
+        -------
+        f_array, sd_array : Two numpy arrays containing the evaluation of the model
+            value and its uncertainty, respectively.
+        """
+        if self.model is None:
+            raise RuntimeError("Model not present. Run ``build_model`` first.")
+        
+        # Get optimum
+        best_arm, best_point_predictions = self.model.model_best_point()
+        parameters = best_arm.parameters
+        parnames = list(parameters.keys())
+
+        # user specific point
+        if p0 is not None:
+            for key in p0.keys():
+                if key in parameters.keys():
+                    parameters[key] = p0[key]
+
+        obsf_list = []
+        obsf_0 = ObservationFeatures(parameters=parameters)
+        if isinstance(sample, np.ndarray):
+            # check the shape of the array
+            if sample.shape[1] != len(parnames):
+                raise RuntimeError(
+                    "Second dimension of the sample array should match " 
+                    "the number of parameters of the model."
+                    )
+            for i in range(sample.shape[0]):
+                predf = deepcopy(obsf_0)
+                for j, parname in enumerate(parameters.keys()):
+                    predf.parameters[parname] = sample[i][j]
+                obsf_list.append(predf)
+        elif isinstance(sample, pd.DataFrame):
+            # check if labels of the dataframe match the parnames
+            for col in sample.columns:
+                if col not in parnames:
+                    raise RuntimeError("Column %s does not match any of the parameter names" % col)
+            for i in range(sample.shape[0]):
+                predf = deepcopy(obsf_0)
+                for col in sample.columns:
+                    predf.parameters[col] = sample[col].iloc[i]
+                obsf_list.append(predf)
+        elif isinstance(sample, dict):
+            # check if the keys of the dictionary match the parnames
+            for key in sample.keys():
+                if key not in parnames:
+                    raise RuntimeError("Key %s does not match any of the parameter names" % col)
+            for i in range(sample[list(sample.keys())[0]].shape[0]):
+                predf = deepcopy(obsf_0)
+                for key in sample.keys():
+                    predf.parameters[key] = sample[key][i]
+                obsf_list.append(predf)
+        elif sample is None:
+            predf = deepcopy(obsf_0)
+            obsf_list.append(predf)
+        else:
+            raise RuntimeError("Wrong data type")
+
+        mu, cov = self.model.predict(obsf_list)
+        metric_name = list(self.ax_client.experiment.metrics.keys())[0]
+        f_array = np.asarray(mu[metric_name])
+        sd_array = np.sqrt(cov[metric_name][metric_name])
+
+        return f_array, sd_array
+
+    def plot_model(
+        self, 
+        xname: Optional[str] = None, 
+        yname: Optional[str] = None, 
+        p0: Optional[Dict] = None, 
+        npoints: Optional[int] = 200,
+        xrange: Optional[List[float]] = None, 
+        yrange: Optional[List[float]] = None, 
+        mode: Optional[Literal["value", "error", "both"]] = "value",
+        clabel: Optional[bool] = False,
+        subplot_spec: Optional[SubplotSpec] = None,
+        gridspec_kw: Dict[str, Any] = None,
+        pcolormesh_kw: Dict[str, Any] = None,
+        **figure_kw,
+    ) -> Union[Axes, List[Axes]]:
+        """Plot model in the two selected variables, while others are fixed to the optimum.
+
+        Parameter:
+        ----------
+        xname: string
+            Name of the variable to plot in x axis.
+        yname: string
+            Name of the variable to plot in y axis.
+        p0: dictionary
+            Particular values of parameters to be fixed for the evaluation over the sample.
+        npoints: int, optional
+            Number of points in each axis.
+        mode: string, optional.
+            ``value`` plots the model value, ``error`` its standard deviation,
+            ``both`` both.
+        clabel: bool,
+            when true labels are shown along the contour lines.
+        gridspec_kw : dict, optional
+            Dict with keywords passed to the `GridSpec`.
+        pcolormesh_kw : dict, optional
+            Dict with keywords passed to `pcolormesh`.
+        **figure_kw
+            Additional keyword arguments to pass to `pyplot.figure`. Only used
+            if no ``subplot_spec`` is given.
+
+        Returns
+        -------
+        `~.axes.Axes` or array of Axes 
+            Either a single `~matplotlib.axes.Axes` object or a list of Axes
+            objects if more than one subplot was created.
+        """
+
+        if self.ax_client is None:
+            raise RuntimeError("AxClient not present. Run `build_model_ax` first.")
+
+        if self.model is None:
+            raise RuntimeError("Model not present. Run `build_model_ax` first.")
+
+        # get experiment info
+        experiment = self.ax_client.experiment
+        parnames = list(experiment.parameters.keys())
+        # minimize = experiment.optimization_config.objective.minimize
+
+        if len(parnames) < 2:
+            raise RuntimeError(
+                "Insufficient number of parameters in data for this plot " 
+                "(minimum 2)."
+            )
+
+        # Make a parameter scan in two of the input dimensions
+        if xname is None:
+            xname = parnames[0]        
+        if yname is None:
+            yname = parnames[1]
+
+        # Set the plotting range
+        if xrange is None:
+            xrange = [None, None]
+        if yrange is None:
+            yrange = [None, None]
+        if xrange[0] is None:
+            xrange[0] = experiment.parameters[xname].lower
+        if xrange[1] is None:
+            xrange[1] = experiment.parameters[xname].upper
+        if yrange[0] is None:
+            yrange[0] = experiment.parameters[yname].lower
+        if yrange[1] is None:
+            yrange[1] = experiment.parameters[yname].upper
+
+        # Get grid sample of points where to evalutate the model
+        xaxis = np.linspace(xrange[0], xrange[1], npoints)
+        yaxis = np.linspace(yrange[0], yrange[1], npoints)
+        X, Y = np.meshgrid(xaxis, yaxis)
+        xarray = X.flatten()
+        yarray = Y.flatten()
+
+        sample = pd.DataFrame({xname: xarray, yname: yarray})
+        f_plt, sd_plt = self.evaluate_model(sample, p0=p0)
+        metric_name = list(self.ax_client.experiment.metrics.keys())[0]
+        
+        # get numpy arrays with experiment parameters
+        xtrials = np.zeros(experiment.num_trials)
+        ytrials = np.zeros(experiment.num_trials)
+        for i in range(experiment.num_trials):
+            xtrials[i] = experiment.trials[i].arm.parameters[xname]
+            ytrials[i] = experiment.trials[i].arm.parameters[yname]
+
+        f_plots = []
+        labels = []
+        if mode in ["value", "both"]: 
+            f_plots.append(f_plt)
+            labels.append(metric_name)
+        if mode in ["error", "both"]: 
+            f_plots.append(sd_plt)
+            labels.append(metric_name + ", error")
+
+        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)
+
+        axs = []
+        for i, f in enumerate(f_plots):
+            ax = plt.subplot(gs[i])
+            pcolormesh_kw = dict(pcolormesh_kw or {})
+            im = ax.pcolormesh(xaxis, yaxis, f.reshape(X.shape), 
+                               shading="auto", **pcolormesh_kw)
+            cbar = plt.colorbar(im, ax=ax)
+            cbar.set_label(labels[i])
+            ax.set(xlabel=xname, ylabel=yname)
+            # adding contour lines
+            cset = ax.contour(X, Y, f.reshape(X.shape), levels=20,
+                       linewidths=0.5, colors="black", linestyles="solid")
+            if clabel:
+                plt.clabel(cset, inline=True, fmt="%1.1f", fontsize="xx-small")
+            ax.scatter(xtrials, ytrials, s=2, c="black", marker="o")
+            ax.set_xlim(xrange)
+            ax.set_ylim(yrange)
+            axs.append(ax)
+
+        if nplots == 1:
+            return axs[0]
+        else:
+            return axs
+
+    def get_arm_index(
+        self, 
+        arm_name: Optional[str] = None,
+    ) -> int:
+        """Get the index of the arm by its name.
+        
+        Parameter:
+        ----------
+        arm_name: string, optional.
+            Name of the arm. If not given, the best arm is selected.
+
+        """
+        if arm_name is None:
+            best_arm, best_point_predictions = self.model.model_best_point()
+            arm_name = best_arm.name
+
+        df = self.ax_client.experiment.fetch_data().df
+        return df.loc[df["arm_name"] == arm_name, "trial_index"].iloc[0]
diff --git a/optimas/diagnostics/exploration_diagnostics.py b/optimas/diagnostics/exploration_diagnostics.py
index 3229692e..a5b24bab 100644
--- a/optimas/diagnostics/exploration_diagnostics.py
+++ b/optimas/diagnostics/exploration_diagnostics.py
@@ -18,6 +18,8 @@
 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
+from ax.service.ax_client import AxClient
 
 
 class ExplorationDiagnostics:
@@ -595,7 +597,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,31 +899,59 @@ def print_evaluation(self, trial_index: int) -> None:
 
         print()
 
-    def print_best_evaluations(
+    def get_best_evaluations_index(
         self,
-        top: Optional[int] = 3,
         objective: Optional[Union[str, Objective]] = None,
-    ) -> None:
-        """Print top evaluations according to the given objective.
+        top: Optional[int] = 3,
+        ) -> List[int]:
+        """Get a list with the indices of the best 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.
+        
+        Returns
+        -------
+        top_indices : List with the indices of the best evaluations.
         """
         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]
+        return top_indices
+
+    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.
+        """
+        if objective is None:
+            objective = self.objectives[0]
+        elif isinstance(objective, str):
+            objective = self._get_objective(objective)
+        top_indices = self.get_best_evaluations_index(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]
@@ -936,3 +966,76 @@ def print_best_evaluations(
                 objective_names + varpar_names + anapar_names
             ]
         )
+
+    def get_model_manager_from_ax_client(
+            self,
+            source: Union[AxClient, str],
+    ) ->  AxModelManager:
+        """Initialize AxModelManager from an existing ``AxClient``.
+            
+        Parameter:
+        ----------
+        source: AxClient or str,
+            Source of data from where to obtain the model.
+            It can be an existing ``AxClient`` or the path to
+            a json file.
+        
+        Returns:
+        --------
+        An instance of AxModelManager            
+        """
+        self.model_manager = AxModelManager(source)
+        return self.model_manager
+
+    def build_model(
+            self, 
+            objname: Optional[str] = '', 
+            parnames: Optional[List[str]] = [], 
+            minimize: Optional[bool] = True
+    ) -> AxModelManager:
+        """Initialize AxModelManager and builds a GP model.
+
+        Parameter:
+        ----------
+        objname: string, optional
+            Name of the objective (or metric).
+            If not given, it takes the first of the objectives.
+        parnames: list of string, optional
+            List with the names of the parameters of the model.
+            If not given, it assumes the varying parameters.
+        minimize: bool, optional
+            Whether to minimize or maximize the objective.
+            Only relevant to establish the best point of the model,
+            but not to build the model.
+
+        Returns:
+        --------
+        An instance of AxModelManager
+        """
+        varpar_names = [var.name for var in self.varying_parameters]
+        objective_names = [obj.name for obj in self.objectives]
+
+        if len(parnames) == 0:
+            parnames = varpar_names
+        if objname == '':
+            objname = objective_names[0]
+
+        if objname in objective_names:
+            minimize = self._get_objective(objname).minimize
+
+        # Copy the history DataFrame
+        df = self.history.copy()
+        self.model_manager = AxModelManager(df)
+        self.model_manager.build_model(parnames=parnames, 
+                                       objname=objname, 
+                                       minimize=minimize)
+
+        return self.model_manager
+
+    def get_model_manager(self):
+        """Get the associated AxModelManager or build from scratch.
+        """
+        if self.model_manager is None:
+            return self.build_model()
+        else:
+            return self.model_manager
\ No newline at end of file

From 721a42ac3d8448030cb0f620034b46988bc1b7aa Mon Sep 17 00:00:00 2001
From: delaossa <alberto.martinez.de.la.ossa@desy.de>
Date: Fri, 9 Feb 2024 16:40:04 +0100
Subject: [PATCH 02/54] Some fixes

---
 optimas/diagnostics/ax_model_manager.py | 53 +++++++++++++++----------
 1 file changed, 32 insertions(+), 21 deletions(-)

diff --git a/optimas/diagnostics/ax_model_manager.py b/optimas/diagnostics/ax_model_manager.py
index c7c2177c..9e51ba78 100644
--- a/optimas/diagnostics/ax_model_manager.py
+++ b/optimas/diagnostics/ax_model_manager.py
@@ -31,7 +31,6 @@ def __init__(self, source: Union[AxClient, str, pd.DataFrame]) -> None:
             If ``AxClient``, it uses the model in there.
             If ``str``, it should contain the path to an ``AxClient`` json file.
         """
-
         if isinstance(source, AxClient):
             self.ax_client = source
         elif isinstance(source, str):
@@ -53,8 +52,8 @@ def model(self) -> TorchModelBridge:
 
     def build_model(
         self, 
+        objname: str,
         parnames: List[str], 
-        objname: str, 
         minimize: Optional[bool] = True
     ) -> None:
         """Initialize the AxClient using the given data, 
@@ -119,8 +118,8 @@ def evaluate_model(
 
         Returns
         -------
-        f_array, sd_array : Two numpy arrays containing the evaluation of the model
-            value and its uncertainty, respectively.
+        m_array, sem_array : Two numpy arrays containing the mean of the model
+            and the standard error of the mean (sem), respectively.
         """
         if self.model is None:
             raise RuntimeError("Model not present. Run ``build_model`` first.")
@@ -154,7 +153,9 @@ def evaluate_model(
             # check if labels of the dataframe match the parnames
             for col in sample.columns:
                 if col not in parnames:
-                    raise RuntimeError("Column %s does not match any of the parameter names" % col)
+                    raise RuntimeError(
+                        "Column %s does not match any of the parameter names" % col
+                    )
             for i in range(sample.shape[0]):
                 predf = deepcopy(obsf_0)
                 for col in sample.columns:
@@ -164,12 +165,22 @@ def evaluate_model(
             # check if the keys of the dictionary match the parnames
             for key in sample.keys():
                 if key not in parnames:
-                    raise RuntimeError("Key %s does not match any of the parameter names" % col)
-            for i in range(sample[list(sample.keys())[0]].shape[0]):
+                    raise RuntimeError(
+                        "Key %s does not match any of the parameter names" % col
+                    )
+            element = sample[list(sample.keys())[0]]
+            if hasattr(element, "__len__"):
+                for i in range(len(element)):
+                    predf = deepcopy(obsf_0)
+                    for key in sample.keys():
+                        predf.parameters[key] = sample[key][i]
+                    obsf_list.append(predf)
+            else:
                 predf = deepcopy(obsf_0)
                 for key in sample.keys():
-                    predf.parameters[key] = sample[key][i]
+                    predf.parameters[key] = sample[key]
                 obsf_list.append(predf)
+
         elif sample is None:
             predf = deepcopy(obsf_0)
             obsf_list.append(predf)
@@ -178,10 +189,9 @@ def evaluate_model(
 
         mu, cov = self.model.predict(obsf_list)
         metric_name = list(self.ax_client.experiment.metrics.keys())[0]
-        f_array = np.asarray(mu[metric_name])
-        sd_array = np.sqrt(cov[metric_name][metric_name])
-
-        return f_array, sd_array
+        m_array = np.asarray(mu[metric_name])
+        sem_array = np.sqrt(cov[metric_name][metric_name])
+        return m_array, sem_array
 
     def plot_model(
         self, 
@@ -191,7 +201,7 @@ def plot_model(
         npoints: Optional[int] = 200,
         xrange: Optional[List[float]] = None, 
         yrange: Optional[List[float]] = None, 
-        mode: Optional[Literal["value", "error", "both"]] = "value",
+        mode: Optional[Literal["mean", "sem", "both"]] = "mean",
         clabel: Optional[bool] = False,
         subplot_spec: Optional[SubplotSpec] = None,
         gridspec_kw: Dict[str, Any] = None,
@@ -211,8 +221,8 @@ def plot_model(
         npoints: int, optional
             Number of points in each axis.
         mode: string, optional.
-            ``value`` plots the model value, ``error`` its standard deviation,
-            ``both`` both.
+            ``mean`` plots the model mean, ``sem`` the standard error of the mean,
+            ``both`` plots both.
         clabel: bool,
             when true labels are shown along the contour lines.
         gridspec_kw : dict, optional
@@ -229,7 +239,6 @@ def plot_model(
             Either a single `~matplotlib.axes.Axes` object or a list of Axes
             objects if more than one subplot was created.
         """
-
         if self.ax_client is None:
             raise RuntimeError("AxClient not present. Run `build_model_ax` first.")
 
@@ -287,12 +296,12 @@ def plot_model(
 
         f_plots = []
         labels = []
-        if mode in ["value", "both"]: 
+        if mode in ["mean", "both"]: 
             f_plots.append(f_plt)
-            labels.append(metric_name)
-        if mode in ["error", "both"]: 
+            labels.append(metric_name + ", mean")
+        if mode in ["sem", "both"]: 
             f_plots.append(sd_plt)
-            labels.append(metric_name + ", error")
+            labels.append(metric_name + ", sem")
 
         nplots = len(f_plots)
         gridspec_kw = dict(gridspec_kw or {})
@@ -320,6 +329,9 @@ def plot_model(
             ax.scatter(xtrials, ytrials, s=2, c="black", marker="o")
             ax.set_xlim(xrange)
             ax.set_ylim(yrange)
+            if i > 0:
+                ax.set_ylabel('')
+                ax.set_yticklabels([])
             axs.append(ax)
 
         if nplots == 1:
@@ -337,7 +349,6 @@ def get_arm_index(
         ----------
         arm_name: string, optional.
             Name of the arm. If not given, the best arm is selected.
-
         """
         if arm_name is None:
             best_arm, best_point_predictions = self.model.model_best_point()

From a56c7042db38545873fc42266b89f03e0f30ca5b Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Fri, 9 Feb 2024 16:10:53 +0000
Subject: [PATCH 03/54] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 optimas/diagnostics/ax_model_manager.py       | 103 +++++++++++-------
 .../diagnostics/exploration_diagnostics.py    |  43 ++++----
 2 files changed, 84 insertions(+), 62 deletions(-)

diff --git a/optimas/diagnostics/ax_model_manager.py b/optimas/diagnostics/ax_model_manager.py
index 9e51ba78..276a4ced 100644
--- a/optimas/diagnostics/ax_model_manager.py
+++ b/optimas/diagnostics/ax_model_manager.py
@@ -11,7 +11,9 @@
 # Ax utilities for model building
 from ax.service.ax_client import AxClient
 from ax.modelbridge.generation_strategy import (
-    GenerationStep, GenerationStrategy)
+    GenerationStep,
+    GenerationStrategy,
+)
 from ax.modelbridge.registry import Models
 from ax.modelbridge.factory import get_GPEI
 from ax.modelbridge.torch import TorchModelBridge
@@ -51,12 +53,9 @@ def model(self) -> TorchModelBridge:
         return self.ax_client.generation_strategy.model
 
     def build_model(
-        self, 
-        objname: str,
-        parnames: List[str], 
-        minimize: Optional[bool] = True
+        self, objname: str, parnames: List[str], minimize: Optional[bool] = True
     ) -> None:
-        """Initialize the AxClient using the given data, 
+        """Initialize the AxClient using the given data,
             the model parameters and the metric.
             Then it fits a Gaussian Process model to the data.
 
@@ -70,14 +69,20 @@ def build_model(
             Whether to minimize or maximize the objective.
             Only relevant to establish the best point.
         """
-        parameters = [{"name": p_name,
-                       "type": "range",
-                       "bounds": [self.df[p_name].min(), self.df[p_name].max()],
-                       "value_type": "float"
-                       } for p_name in parnames]
+        parameters = [
+            {
+                "name": p_name,
+                "type": "range",
+                "bounds": [self.df[p_name].min(), self.df[p_name].max()],
+                "value_type": "float",
+            }
+            for p_name in parnames
+        ]
 
         # create Ax client
-        gs = GenerationStrategy([GenerationStep(model=Models.GPEI, num_trials=-1)])
+        gs = GenerationStrategy(
+            [GenerationStep(model=Models.GPEI, num_trials=-1)]
+        )
         self.ax_client = AxClient(generation_strategy=gs, verbose_logging=False)
         self.ax_client.create_experiment(
             name="optimas_data",
@@ -91,15 +96,16 @@ def build_model(
         for index, row in self.df.iterrows():
             params = {p_name: row[p_name] for p_name in parnames}
             _, trial_id = self.ax_client.attach_trial(params)
-            self.ax_client.complete_trial(trial_id, 
-                                          {metric_name: (row[metric_name], np.nan)})
+            self.ax_client.complete_trial(
+                trial_id, {metric_name: (row[metric_name], np.nan)}
+            )
 
         # fit GP model
         self.ax_client.generation_strategy._fit_current_model(data=None)
 
     def evaluate_model(
-        self, 
-        sample: Union[pd.DataFrame, Dict, np.ndarray] = None, 
+        self,
+        sample: Union[pd.DataFrame, Dict, np.ndarray] = None,
         p0: Dict = None,
     ) -> Tuple[np.ndarray]:
         """Evaluate the model over the specified sample.
@@ -113,7 +119,7 @@ def evaluate_model(
             The rest of parameters would be set to the model best point,
             unless they are further specified using ``p0``.
         p0: dictionary
-            Particular values of parameters to be fixed for the evaluation 
+            Particular values of parameters to be fixed for the evaluation
             over the givensample.
 
         Returns
@@ -123,7 +129,7 @@ def evaluate_model(
         """
         if self.model is None:
             raise RuntimeError("Model not present. Run ``build_model`` first.")
-        
+
         # Get optimum
         best_arm, best_point_predictions = self.model.model_best_point()
         parameters = best_arm.parameters
@@ -141,9 +147,9 @@ def evaluate_model(
             # check the shape of the array
             if sample.shape[1] != len(parnames):
                 raise RuntimeError(
-                    "Second dimension of the sample array should match " 
+                    "Second dimension of the sample array should match "
                     "the number of parameters of the model."
-                    )
+                )
             for i in range(sample.shape[0]):
                 predf = deepcopy(obsf_0)
                 for j, parname in enumerate(parameters.keys()):
@@ -154,7 +160,8 @@ def evaluate_model(
             for col in sample.columns:
                 if col not in parnames:
                     raise RuntimeError(
-                        "Column %s does not match any of the parameter names" % col
+                        "Column %s does not match any of the parameter names"
+                        % col
                     )
             for i in range(sample.shape[0]):
                 predf = deepcopy(obsf_0)
@@ -194,13 +201,13 @@ def evaluate_model(
         return m_array, sem_array
 
     def plot_model(
-        self, 
-        xname: Optional[str] = None, 
-        yname: Optional[str] = None, 
-        p0: Optional[Dict] = None, 
+        self,
+        xname: Optional[str] = None,
+        yname: Optional[str] = None,
+        p0: Optional[Dict] = None,
         npoints: Optional[int] = 200,
-        xrange: Optional[List[float]] = None, 
-        yrange: Optional[List[float]] = None, 
+        xrange: Optional[List[float]] = None,
+        yrange: Optional[List[float]] = None,
         mode: Optional[Literal["mean", "sem", "both"]] = "mean",
         clabel: Optional[bool] = False,
         subplot_spec: Optional[SubplotSpec] = None,
@@ -235,12 +242,14 @@ def plot_model(
 
         Returns
         -------
-        `~.axes.Axes` or array of Axes 
+        `~.axes.Axes` or array of Axes
             Either a single `~matplotlib.axes.Axes` object or a list of Axes
             objects if more than one subplot was created.
         """
         if self.ax_client is None:
-            raise RuntimeError("AxClient not present. Run `build_model_ax` first.")
+            raise RuntimeError(
+                "AxClient not present. Run `build_model_ax` first."
+            )
 
         if self.model is None:
             raise RuntimeError("Model not present. Run `build_model_ax` first.")
@@ -252,13 +261,13 @@ def plot_model(
 
         if len(parnames) < 2:
             raise RuntimeError(
-                "Insufficient number of parameters in data for this plot " 
+                "Insufficient number of parameters in data for this plot "
                 "(minimum 2)."
             )
 
         # Make a parameter scan in two of the input dimensions
         if xname is None:
-            xname = parnames[0]        
+            xname = parnames[0]
         if yname is None:
             yname = parnames[1]
 
@@ -286,7 +295,7 @@ def plot_model(
         sample = pd.DataFrame({xname: xarray, yname: yarray})
         f_plt, sd_plt = self.evaluate_model(sample, p0=p0)
         metric_name = list(self.ax_client.experiment.metrics.keys())[0]
-        
+
         # get numpy arrays with experiment parameters
         xtrials = np.zeros(experiment.num_trials)
         ytrials = np.zeros(experiment.num_trials)
@@ -296,10 +305,10 @@ def plot_model(
 
         f_plots = []
         labels = []
-        if mode in ["mean", "both"]: 
+        if mode in ["mean", "both"]:
             f_plots.append(f_plt)
             labels.append(metric_name + ", mean")
-        if mode in ["sem", "both"]: 
+        if mode in ["sem", "both"]:
             f_plots.append(sd_plt)
             labels.append(metric_name + ", sem")
 
@@ -316,21 +325,33 @@ def plot_model(
         for i, f in enumerate(f_plots):
             ax = plt.subplot(gs[i])
             pcolormesh_kw = dict(pcolormesh_kw or {})
-            im = ax.pcolormesh(xaxis, yaxis, f.reshape(X.shape), 
-                               shading="auto", **pcolormesh_kw)
+            im = ax.pcolormesh(
+                xaxis,
+                yaxis,
+                f.reshape(X.shape),
+                shading="auto",
+                **pcolormesh_kw,
+            )
             cbar = plt.colorbar(im, ax=ax)
             cbar.set_label(labels[i])
             ax.set(xlabel=xname, ylabel=yname)
             # adding contour lines
-            cset = ax.contour(X, Y, f.reshape(X.shape), levels=20,
-                       linewidths=0.5, colors="black", linestyles="solid")
+            cset = ax.contour(
+                X,
+                Y,
+                f.reshape(X.shape),
+                levels=20,
+                linewidths=0.5,
+                colors="black",
+                linestyles="solid",
+            )
             if clabel:
                 plt.clabel(cset, inline=True, fmt="%1.1f", fontsize="xx-small")
             ax.scatter(xtrials, ytrials, s=2, c="black", marker="o")
             ax.set_xlim(xrange)
             ax.set_ylim(yrange)
             if i > 0:
-                ax.set_ylabel('')
+                ax.set_ylabel("")
                 ax.set_yticklabels([])
             axs.append(ax)
 
@@ -340,11 +361,11 @@ def plot_model(
             return axs
 
     def get_arm_index(
-        self, 
+        self,
         arm_name: Optional[str] = None,
     ) -> int:
         """Get the index of the arm by its name.
-        
+
         Parameter:
         ----------
         arm_name: string, optional.
diff --git a/optimas/diagnostics/exploration_diagnostics.py b/optimas/diagnostics/exploration_diagnostics.py
index a5b24bab..253dbb7b 100644
--- a/optimas/diagnostics/exploration_diagnostics.py
+++ b/optimas/diagnostics/exploration_diagnostics.py
@@ -903,8 +903,8 @@ def get_best_evaluations_index(
         self,
         objective: Optional[Union[str, Objective]] = None,
         top: Optional[int] = 3,
-        ) -> List[int]:
-        """Get a list with the indices of the best evaluations 
+    ) -> List[int]:
+        """Get a list with the indices of the best evaluations
             according to the given objective.
 
         Parameters
@@ -915,7 +915,7 @@ def get_best_evaluations_index(
         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.
@@ -951,7 +951,9 @@ def print_best_evaluations(
             objective = self.objectives[0]
         elif isinstance(objective, str):
             objective = self._get_objective(objective)
-        top_indices = self.get_best_evaluations_index(top=top, objective=objective)
+        top_indices = self.get_best_evaluations_index(
+            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]
@@ -968,30 +970,30 @@ def print_best_evaluations(
         )
 
     def get_model_manager_from_ax_client(
-            self,
-            source: Union[AxClient, str],
-    ) ->  AxModelManager:
+        self,
+        source: Union[AxClient, str],
+    ) -> AxModelManager:
         """Initialize AxModelManager from an existing ``AxClient``.
-            
+
         Parameter:
         ----------
         source: AxClient or str,
             Source of data from where to obtain the model.
             It can be an existing ``AxClient`` or the path to
             a json file.
-        
+
         Returns:
         --------
-        An instance of AxModelManager            
+        An instance of AxModelManager
         """
         self.model_manager = AxModelManager(source)
         return self.model_manager
 
     def build_model(
-            self, 
-            objname: Optional[str] = '', 
-            parnames: Optional[List[str]] = [], 
-            minimize: Optional[bool] = True
+        self,
+        objname: Optional[str] = "",
+        parnames: Optional[List[str]] = [],
+        minimize: Optional[bool] = True,
     ) -> AxModelManager:
         """Initialize AxModelManager and builds a GP model.
 
@@ -1017,7 +1019,7 @@ def build_model(
 
         if len(parnames) == 0:
             parnames = varpar_names
-        if objname == '':
+        if objname == "":
             objname = objective_names[0]
 
         if objname in objective_names:
@@ -1026,16 +1028,15 @@ def build_model(
         # Copy the history DataFrame
         df = self.history.copy()
         self.model_manager = AxModelManager(df)
-        self.model_manager.build_model(parnames=parnames, 
-                                       objname=objname, 
-                                       minimize=minimize)
+        self.model_manager.build_model(
+            parnames=parnames, objname=objname, minimize=minimize
+        )
 
         return self.model_manager
 
     def get_model_manager(self):
-        """Get the associated AxModelManager or build from scratch.
-        """
+        """Get the associated AxModelManager or build from scratch."""
         if self.model_manager is None:
             return self.build_model()
         else:
-            return self.model_manager
\ No newline at end of file
+            return self.model_manager

From d8213b705c55c3ab6379a783f4ca2767914193c9 Mon Sep 17 00:00:00 2001
From: delaossa <alberto.martinez.de.la.ossa@desy.de>
Date: Sat, 10 Feb 2024 17:51:10 +0100
Subject: [PATCH 04/54] More fixes.

---
 optimas/diagnostics/ax_model_manager.py       | 74 ++++++++++---------
 .../diagnostics/exploration_diagnostics.py    | 17 ++---
 2 files changed, 48 insertions(+), 43 deletions(-)

diff --git a/optimas/diagnostics/ax_model_manager.py b/optimas/diagnostics/ax_model_manager.py
index 276a4ced..d5cd9b41 100644
--- a/optimas/diagnostics/ax_model_manager.py
+++ b/optimas/diagnostics/ax_model_manager.py
@@ -2,6 +2,7 @@
 
 from typing import Optional, Union, List, Tuple, Dict, Any, Literal
 import numpy as np
+import numpy.typing as npt
 import pandas as pd
 from copy import deepcopy
 import matplotlib.pyplot as plt
@@ -21,18 +22,18 @@
 
 
 class AxModelManager(object):
+    """Utilities for building and exploring surrogate models using ``Ax``.
 
-    def __init__(self, source: Union[AxClient, str, pd.DataFrame]) -> None:
-        """Utilities for building and exploring surrogate models using ``Ax``.
+    Parameters
+    ----------
+    source: AxClient, str or DataFrame
+        Source data for the model.
+        If ``DataFrame``, the model has to be build using ``build_model``.
+        If ``AxClient``, it uses the data in there to build a model.
+        If ``str``, it should be the path to an ``AxClient`` json file.
+    """
 
-        Parameter:
-        ----------
-        source: AxClient, str or DataFrame
-            Source of data from where to obtain the model.
-            If ``DataFrame``, the model has to be build using ``build_model``.
-            If ``AxClient``, it uses the model in there.
-            If ``str``, it should contain the path to an ``AxClient`` json file.
-        """
+    def __init__(self, source: Union[AxClient, str, pd.DataFrame]) -> None:
         if isinstance(source, AxClient):
             self.ax_client = source
         elif isinstance(source, str):
@@ -44,8 +45,8 @@ def __init__(self, source: Union[AxClient, str, pd.DataFrame]) -> None:
             raise RuntimeError("Wrong source.")
 
         if self.ax_client:
-            # fit GP model
-            self.ax_client.generation_strategy._fit_current_model(data=None)
+            if self.ax_client.generation_strategy.model is None:
+                self.ax_client.generation_strategy._fit_current_model(data=None)
 
     @property
     def model(self) -> TorchModelBridge:
@@ -55,16 +56,14 @@ def model(self) -> TorchModelBridge:
     def build_model(
         self, objname: str, parnames: List[str], minimize: Optional[bool] = True
     ) -> None:
-        """Initialize the AxClient using the given data,
-            the model parameters and the metric.
-            Then it fits a Gaussian Process model to the data.
+        """Initialize the AxClient and the model using the given data.
 
-        Parameter:
+        Parameters
         ----------
         parnames: list of string
-            List with the names of the parameters of the model
+            List with the names of the parameters of the model.
         objname: string
-            Name of the objective
+            Name of the objective.
         minimize: bool
             Whether to minimize or maximize the objective.
             Only relevant to establish the best point.
@@ -105,9 +104,9 @@ def build_model(
 
     def evaluate_model(
         self,
-        sample: Union[pd.DataFrame, Dict, np.ndarray] = None,
+        sample: Union[pd.DataFrame, Dict, npt.NDArray] = None,
         p0: Dict = None,
-    ) -> Tuple[np.ndarray]:
+    ) -> Tuple[npt.NDArray]:
         """Evaluate the model over the specified sample.
 
         Parameter:
@@ -130,7 +129,7 @@ def evaluate_model(
         if self.model is None:
             raise RuntimeError("Model not present. Run ``build_model`` first.")
 
-        # Get optimum
+        # get optimum
         best_arm, best_point_predictions = self.model.model_best_point()
         parameters = best_arm.parameters
         parnames = list(parameters.keys())
@@ -230,7 +229,7 @@ def plot_model(
         mode: string, optional.
             ``mean`` plots the model mean, ``sem`` the standard error of the mean,
             ``both`` plots both.
-        clabel: bool,
+        clabel: bool
             when true labels are shown along the contour lines.
         gridspec_kw : dict, optional
             Dict with keywords passed to the `GridSpec`.
@@ -265,13 +264,13 @@ def plot_model(
                 "(minimum 2)."
             )
 
-        # Make a parameter scan in two of the input dimensions
+        # select the input variables
         if xname is None:
             xname = parnames[0]
         if yname is None:
             yname = parnames[1]
 
-        # Set the plotting range
+        # set the plotting range
         if xrange is None:
             xrange = [None, None]
         if yrange is None:
@@ -285,16 +284,14 @@ def plot_model(
         if yrange[1] is None:
             yrange[1] = experiment.parameters[yname].upper
 
-        # Get grid sample of points where to evalutate the model
+        # get grid sample of points where to evalutate the model
         xaxis = np.linspace(xrange[0], xrange[1], npoints)
         yaxis = np.linspace(yrange[0], yrange[1], npoints)
         X, Y = np.meshgrid(xaxis, yaxis)
         xarray = X.flatten()
         yarray = Y.flatten()
-
         sample = pd.DataFrame({xname: xarray, yname: yarray})
         f_plt, sd_plt = self.evaluate_model(sample, p0=p0)
-        metric_name = list(self.ax_client.experiment.metrics.keys())[0]
 
         # get numpy arrays with experiment parameters
         xtrials = np.zeros(experiment.num_trials)
@@ -303,8 +300,10 @@ def plot_model(
             xtrials[i] = experiment.trials[i].arm.parameters[xname]
             ytrials[i] = experiment.trials[i].arm.parameters[yname]
 
+        # select quantities to plot and set the labels
         f_plots = []
         labels = []
+        metric_name = list(experiment.metrics.keys())[0]
         if mode in ["mean", "both"]:
             f_plots.append(f_plt)
             labels.append(metric_name + ", mean")
@@ -312,6 +311,7 @@ def plot_model(
             f_plots.append(sd_plt)
             labels.append(metric_name + ", sem")
 
+        # create figure
         nplots = len(f_plots)
         gridspec_kw = dict(gridspec_kw or {})
         if subplot_spec is None:
@@ -321,9 +321,11 @@ def plot_model(
             fig = plt.gcf()
             gs = GridSpecFromSubplotSpec(1, nplots, subplot_spec, **gridspec_kw)
 
+        # draw plots
         axs = []
         for i, f in enumerate(f_plots):
             ax = plt.subplot(gs[i])
+            # colormesh
             pcolormesh_kw = dict(pcolormesh_kw or {})
             im = ax.pcolormesh(
                 xaxis,
@@ -332,10 +334,10 @@ def plot_model(
                 shading="auto",
                 **pcolormesh_kw,
             )
-            cbar = plt.colorbar(im, ax=ax)
+            cbar = plt.colorbar(im, ax=ax, location='top')
             cbar.set_label(labels[i])
             ax.set(xlabel=xname, ylabel=yname)
-            # adding contour lines
+            # contour lines
             cset = ax.contour(
                 X,
                 Y,
@@ -347,12 +349,10 @@ def plot_model(
             )
             if clabel:
                 plt.clabel(cset, inline=True, fmt="%1.1f", fontsize="xx-small")
+            # Draw trials
             ax.scatter(xtrials, ytrials, s=2, c="black", marker="o")
             ax.set_xlim(xrange)
             ax.set_ylim(yrange)
-            if i > 0:
-                ax.set_ylabel("")
-                ax.set_yticklabels([])
             axs.append(ax)
 
         if nplots == 1:
@@ -366,14 +366,20 @@ def get_arm_index(
     ) -> int:
         """Get the index of the arm by its name.
 
-        Parameter:
+        Parameters
         ----------
         arm_name: string, optional.
             Name of the arm. If not given, the best arm is selected.
+
+        Returns
+        -------
+        index: int
+            Trial index of the arm.
         """
         if arm_name is None:
             best_arm, best_point_predictions = self.model.model_best_point()
             arm_name = best_arm.name
 
         df = self.ax_client.experiment.fetch_data().df
-        return df.loc[df["arm_name"] == arm_name, "trial_index"].iloc[0]
+        index = df.loc[df["arm_name"] == arm_name, "trial_index"].iloc[0]
+        return index
\ No newline at end of file
diff --git a/optimas/diagnostics/exploration_diagnostics.py b/optimas/diagnostics/exploration_diagnostics.py
index 253dbb7b..9914a1cb 100644
--- a/optimas/diagnostics/exploration_diagnostics.py
+++ b/optimas/diagnostics/exploration_diagnostics.py
@@ -904,8 +904,7 @@ def get_best_evaluations_index(
         objective: Optional[Union[str, Objective]] = None,
         top: Optional[int] = 3,
     ) -> List[int]:
-        """Get a list with the indices of the best evaluations
-            according to the given objective.
+        """Get a list with the indices of the best evaluations.
 
         Parameters
         ----------
@@ -975,16 +974,16 @@ def get_model_manager_from_ax_client(
     ) -> AxModelManager:
         """Initialize AxModelManager from an existing ``AxClient``.
 
-        Parameter:
+        Parameters
         ----------
         source: AxClient or str,
             Source of data from where to obtain the model.
             It can be an existing ``AxClient`` or the path to
             a json file.
 
-        Returns:
-        --------
-        An instance of AxModelManager
+        Returns
+        -------
+        An instance of AxModelManager.
         """
         self.model_manager = AxModelManager(source)
         return self.model_manager
@@ -997,7 +996,7 @@ def build_model(
     ) -> AxModelManager:
         """Initialize AxModelManager and builds a GP model.
 
-        Parameter:
+        Parameters
         ----------
         objname: string, optional
             Name of the objective (or metric).
@@ -1010,8 +1009,8 @@ def build_model(
             Only relevant to establish the best point of the model,
             but not to build the model.
 
-        Returns:
-        --------
+        Returns
+        -------
         An instance of AxModelManager
         """
         varpar_names = [var.name for var in self.varying_parameters]

From 3019124ddc9c55761d7250be1344254c83e44be3 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Sat, 10 Feb 2024 16:51:29 +0000
Subject: [PATCH 05/54] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 optimas/diagnostics/ax_model_manager.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/optimas/diagnostics/ax_model_manager.py b/optimas/diagnostics/ax_model_manager.py
index d5cd9b41..507b97b0 100644
--- a/optimas/diagnostics/ax_model_manager.py
+++ b/optimas/diagnostics/ax_model_manager.py
@@ -334,7 +334,7 @@ def plot_model(
                 shading="auto",
                 **pcolormesh_kw,
             )
-            cbar = plt.colorbar(im, ax=ax, location='top')
+            cbar = plt.colorbar(im, ax=ax, location="top")
             cbar.set_label(labels[i])
             ax.set(xlabel=xname, ylabel=yname)
             # contour lines
@@ -382,4 +382,4 @@ def get_arm_index(
 
         df = self.ax_client.experiment.fetch_data().df
         index = df.loc[df["arm_name"] == arm_name, "trial_index"].iloc[0]
-        return index
\ No newline at end of file
+        return index

From 17b58b551b901a5693aa230bbdd6d36222d14baa Mon Sep 17 00:00:00 2001
From: delaossa <alberto.martinez.de.la.ossa@desy.de>
Date: Mon, 12 Feb 2024 13:07:53 +0100
Subject: [PATCH 06/54] Change callable for fitting model.

---
 optimas/diagnostics/ax_model_manager.py | 26 ++++++++++++-------------
 1 file changed, 13 insertions(+), 13 deletions(-)

diff --git a/optimas/diagnostics/ax_model_manager.py b/optimas/diagnostics/ax_model_manager.py
index d5cd9b41..c34538ab 100644
--- a/optimas/diagnostics/ax_model_manager.py
+++ b/optimas/diagnostics/ax_model_manager.py
@@ -2,8 +2,8 @@
 
 from typing import Optional, Union, List, Tuple, Dict, Any, Literal
 import numpy as np
-import numpy.typing as npt
-import pandas as pd
+from numpy.typing import NDArray
+from pandas import DataFrame
 from copy import deepcopy
 import matplotlib.pyplot as plt
 from matplotlib.gridspec import GridSpec, GridSpecFromSubplotSpec, SubplotSpec
@@ -33,12 +33,12 @@ class AxModelManager(object):
         If ``str``, it should be the path to an ``AxClient`` json file.
     """
 
-    def __init__(self, source: Union[AxClient, str, pd.DataFrame]) -> None:
+    def __init__(self, source: Union[AxClient, str, DataFrame]) -> None:
         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):
+        elif isinstance(source, DataFrame):
             self.df = source
             self.ax_client = None
         else:
@@ -46,7 +46,7 @@ def __init__(self, source: Union[AxClient, str, pd.DataFrame]) -> None:
 
         if self.ax_client:
             if self.ax_client.generation_strategy.model is None:
-                self.ax_client.generation_strategy._fit_current_model(data=None)
+                self.ax_client.fit_model()
 
     @property
     def model(self) -> TorchModelBridge:
@@ -100,13 +100,13 @@ def build_model(
             )
 
         # fit GP model
-        self.ax_client.generation_strategy._fit_current_model(data=None)
+        self.ax_client.fit_model()
 
     def evaluate_model(
         self,
-        sample: Union[pd.DataFrame, Dict, npt.NDArray] = None,
+        sample: Union[DataFrame, Dict, NDArray] = None,
         p0: Dict = None,
-    ) -> Tuple[npt.NDArray]:
+    ) -> Tuple[NDArray]:
         """Evaluate the model over the specified sample.
 
         Parameter:
@@ -154,7 +154,7 @@ def evaluate_model(
                 for j, parname in enumerate(parameters.keys()):
                     predf.parameters[parname] = sample[i][j]
                 obsf_list.append(predf)
-        elif isinstance(sample, pd.DataFrame):
+        elif isinstance(sample, DataFrame):
             # check if labels of the dataframe match the parnames
             for col in sample.columns:
                 if col not in parnames:
@@ -247,11 +247,11 @@ def plot_model(
         """
         if self.ax_client is None:
             raise RuntimeError(
-                "AxClient not present. Run `build_model_ax` first."
+                "AxClient not present. Run `build_model` first."
             )
 
         if self.model is None:
-            raise RuntimeError("Model not present. Run `build_model_ax` first.")
+            raise RuntimeError("Model not present. Run `build_model` first.")
 
         # get experiment info
         experiment = self.ax_client.experiment
@@ -290,7 +290,7 @@ def plot_model(
         X, Y = np.meshgrid(xaxis, yaxis)
         xarray = X.flatten()
         yarray = Y.flatten()
-        sample = pd.DataFrame({xname: xarray, yname: yarray})
+        sample = DataFrame({xname: xarray, yname: yarray})
         f_plt, sd_plt = self.evaluate_model(sample, p0=p0)
 
         # get numpy arrays with experiment parameters
@@ -349,7 +349,7 @@ def plot_model(
             )
             if clabel:
                 plt.clabel(cset, inline=True, fmt="%1.1f", fontsize="xx-small")
-            # Draw trials
+            # draw trials
             ax.scatter(xtrials, ytrials, s=2, c="black", marker="o")
             ax.set_xlim(xrange)
             ax.set_ylim(yrange)

From 3fdc0b40d51af8465e3ebaa8d9663704cb26475a Mon Sep 17 00:00:00 2001
From: delaossa <alberto.martinez.de.la.ossa@desy.de>
Date: Mon, 12 Feb 2024 16:28:35 +0100
Subject: [PATCH 07/54] Make it work for multi-objective cases.

---
 optimas/diagnostics/ax_model_manager.py | 59 +++++++++++++++++++++----
 1 file changed, 51 insertions(+), 8 deletions(-)

diff --git a/optimas/diagnostics/ax_model_manager.py b/optimas/diagnostics/ax_model_manager.py
index c34538ab..e277309a 100644
--- a/optimas/diagnostics/ax_model_manager.py
+++ b/optimas/diagnostics/ax_model_manager.py
@@ -16,7 +16,6 @@
     GenerationStrategy,
 )
 from ax.modelbridge.registry import Models
-from ax.modelbridge.factory import get_GPEI
 from ax.modelbridge.torch import TorchModelBridge
 from ax.core.observation import ObservationFeatures
 
@@ -105,6 +104,7 @@ def build_model(
     def evaluate_model(
         self,
         sample: Union[DataFrame, Dict, NDArray] = None,
+        metric_name: Optional[str] = None, 
         p0: Dict = None,
     ) -> Tuple[NDArray]:
         """Evaluate the model over the specified sample.
@@ -117,6 +117,9 @@ def evaluate_model(
             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 ``p0``.
+        metric_name: str.
+            Name of the metric to evaluate.
+            If not specified, it will take the first of the list.
         p0: dictionary
             Particular values of parameters to be fixed for the evaluation
             over the givensample.
@@ -129,9 +132,35 @@ def evaluate_model(
         if self.model is None:
             raise RuntimeError("Model not present. Run ``build_model`` first.")
 
+        if metric_name is None:
+            metric_name = list(self.ax_client.experiment.metrics.keys())[0]
+        else:
+            if metric_name not in list(self.ax_client.experiment.metrics.keys()):
+                raise RuntimeError(
+                    "Metric name %s does not match any of the metrics" % metric_name
+                )
+
         # get optimum
-        best_arm, best_point_predictions = self.model.model_best_point()
-        parameters = best_arm.parameters
+        if len(self.ax_client.experiment.metrics) > 1:
+            minimize = None
+            for obj in self.ax_client.objective.objectives:
+                if obj.metric_names[0] == metric_name:
+                    minimize = obj.minimize
+                    break
+            pp = self.ax_client.get_pareto_optimal_parameters()
+            obj_vals = [
+                objs[metric_name] for i, (vals, (objs, covs)) in pp.items()
+            ]
+            param_vals = [vals for i, (vals, (objs, covs)) in pp.items()]
+            if minimize:
+                best_obj_i = np.argmin(obj_vals)
+            else:
+                best_obj_i = np.argmax(obj_vals)
+            best_pars = param_vals[best_obj_i]
+        else:
+            best_pars = self.ax_client.get_best_parameters()[0]
+
+        parameters = best_pars
         parnames = list(parameters.keys())
 
         # user specific point
@@ -140,6 +169,7 @@ def evaluate_model(
                 if key in parameters.keys():
                     parameters[key] = p0[key]
 
+        # make list of `ObservationFeatures`
         obsf_list = []
         obsf_0 = ObservationFeatures(parameters=parameters)
         if isinstance(sample, np.ndarray):
@@ -194,7 +224,6 @@ def evaluate_model(
             raise RuntimeError("Wrong data type")
 
         mu, cov = self.model.predict(obsf_list)
-        metric_name = list(self.ax_client.experiment.metrics.keys())[0]
         m_array = np.asarray(mu[metric_name])
         sem_array = np.sqrt(cov[metric_name][metric_name])
         return m_array, sem_array
@@ -203,7 +232,9 @@ def plot_model(
         self,
         xname: Optional[str] = None,
         yname: Optional[str] = None,
+        mname: Optional[str] = None,
         p0: Optional[Dict] = None,
+        new: Optional[bool] = False,
         npoints: Optional[int] = 200,
         xrange: Optional[List[float]] = None,
         yrange: Optional[List[float]] = None,
@@ -222,6 +253,9 @@ def plot_model(
             Name of the variable to plot in x axis.
         yname: string
             Name of the variable to plot in y axis.
+        mname: string, optional.
+            Name of the metric to plot.
+            If not specified, it will take the first of the list in the ``AxClient``.
         p0: dictionary
             Particular values of parameters to be fixed for the evaluation over the sample.
         npoints: int, optional
@@ -291,7 +325,17 @@ def plot_model(
         xarray = X.flatten()
         yarray = Y.flatten()
         sample = DataFrame({xname: xarray, yname: yarray})
-        f_plt, sd_plt = self.evaluate_model(sample, p0=p0)
+
+        # metric name
+        if mname is None:
+            mname = list(experiment.metrics.keys())[0]
+
+        # evaluate the model
+        f_plt, sd_plt = self.evaluate_model(
+            sample=sample, 
+            metric_name=mname, 
+            p0=p0,
+        )
 
         # get numpy arrays with experiment parameters
         xtrials = np.zeros(experiment.num_trials)
@@ -303,13 +347,12 @@ def plot_model(
         # select quantities to plot and set the labels
         f_plots = []
         labels = []
-        metric_name = list(experiment.metrics.keys())[0]
         if mode in ["mean", "both"]:
             f_plots.append(f_plt)
-            labels.append(metric_name + ", mean")
+            labels.append(mname + ", mean")
         if mode in ["sem", "both"]:
             f_plots.append(sd_plt)
-            labels.append(metric_name + ", sem")
+            labels.append(mname + ", sem")
 
         # create figure
         nplots = len(f_plots)

From 1348e40f77c0074dea54fc3c6e32cb6bfad85c40 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 12 Feb 2024 15:29:18 +0000
Subject: [PATCH 08/54] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 optimas/diagnostics/ax_model_manager.py | 17 +++++++++--------
 1 file changed, 9 insertions(+), 8 deletions(-)

diff --git a/optimas/diagnostics/ax_model_manager.py b/optimas/diagnostics/ax_model_manager.py
index c2997050..33ba4706 100644
--- a/optimas/diagnostics/ax_model_manager.py
+++ b/optimas/diagnostics/ax_model_manager.py
@@ -104,7 +104,7 @@ def build_model(
     def evaluate_model(
         self,
         sample: Union[DataFrame, Dict, NDArray] = None,
-        metric_name: Optional[str] = None, 
+        metric_name: Optional[str] = None,
         p0: Dict = None,
     ) -> Tuple[NDArray]:
         """Evaluate the model over the specified sample.
@@ -135,9 +135,12 @@ def evaluate_model(
         if metric_name is None:
             metric_name = list(self.ax_client.experiment.metrics.keys())[0]
         else:
-            if metric_name not in list(self.ax_client.experiment.metrics.keys()):
+            if metric_name not in list(
+                self.ax_client.experiment.metrics.keys()
+            ):
                 raise RuntimeError(
-                    "Metric name %s does not match any of the metrics" % metric_name
+                    "Metric name %s does not match any of the metrics"
+                    % metric_name
                 )
 
         # get optimum
@@ -280,9 +283,7 @@ def plot_model(
             objects if more than one subplot was created.
         """
         if self.ax_client is None:
-            raise RuntimeError(
-                "AxClient not present. Run `build_model` first."
-            )
+            raise RuntimeError("AxClient not present. Run `build_model` first.")
 
         if self.model is None:
             raise RuntimeError("Model not present. Run `build_model` first.")
@@ -332,8 +333,8 @@ def plot_model(
 
         # evaluate the model
         f_plt, sd_plt = self.evaluate_model(
-            sample=sample, 
-            metric_name=mname, 
+            sample=sample,
+            metric_name=mname,
             p0=p0,
         )
 

From ce73dcabf819eaf8e21ecc7988ce26e121d31a35 Mon Sep 17 00:00:00 2001
From: delaossa <alberto.martinez.de.la.ossa@desy.de>
Date: Tue, 13 Feb 2024 13:39:37 +0100
Subject: [PATCH 09/54] Some more fixes.

---
 optimas/diagnostics/ax_model_manager.py       | 37 ++++++++++---------
 .../diagnostics/exploration_diagnostics.py    |  7 ----
 2 files changed, 19 insertions(+), 25 deletions(-)

diff --git a/optimas/diagnostics/ax_model_manager.py b/optimas/diagnostics/ax_model_manager.py
index c2997050..b80d30f8 100644
--- a/optimas/diagnostics/ax_model_manager.py
+++ b/optimas/diagnostics/ax_model_manager.py
@@ -44,8 +44,7 @@ def __init__(self, source: Union[AxClient, str, DataFrame]) -> None:
             raise RuntimeError("Wrong source.")
 
         if self.ax_client:
-            if self.ax_client.generation_strategy.model is None:
-                self.ax_client.fit_model()
+            self.ax_client.fit_model()
 
     @property
     def model(self) -> TorchModelBridge:
@@ -56,6 +55,7 @@ def build_model(
         self, objname: str, parnames: List[str], minimize: Optional[bool] = True
     ) -> None:
         """Initialize the AxClient and the model using the given data.
+            This method only works with one objective.
 
         Parameters
         ----------
@@ -105,7 +105,7 @@ def evaluate_model(
         self,
         sample: Union[DataFrame, Dict, NDArray] = None,
         metric_name: Optional[str] = None, 
-        p0: Dict = None,
+        p0: Optional[Dict] = None,
     ) -> Tuple[NDArray]:
         """Evaluate the model over the specified sample.
 
@@ -117,12 +117,11 @@ def evaluate_model(
             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 ``p0``.
-        metric_name: str.
+        metric_name: str, optional.
             Name of the metric to evaluate.
-            If not specified, it will take the first of the list.
-        p0: dictionary
-            Particular values of parameters to be fixed for the evaluation
-            over the givensample.
+            If not specified, it will take the first first objective in ``self.ax_client``.
+        p0: dictionary, optional.
+            Particular values of parameters to be fixed for the evaluation.
 
         Returns
         -------
@@ -133,18 +132,21 @@ def evaluate_model(
             raise RuntimeError("Model not present. Run ``build_model`` first.")
 
         if metric_name is None:
-            metric_name = list(self.ax_client.experiment.metrics.keys())[0]
+            metric_name = self.ax_client.objective_names[0]
         else:
-            if metric_name not in list(self.ax_client.experiment.metrics.keys()):
+            metric_names = list(self.ax_client.experiment.metrics.keys())
+            if metric_name not in metric_names:
                 raise RuntimeError(
-                    "Metric name %s does not match any of the metrics" % metric_name
+                    f"Metric name {metric_name} does not match any of the metrics. "
+                    f"Available metrics are: {metric_names}."
                 )
 
         # get optimum
-        if len(self.ax_client.experiment.metrics) > 1:
+        try:
+            objectives = self.ax_client.objective.objectives
             minimize = None
-            for obj in self.ax_client.objective.objectives:
-                if obj.metric_names[0] == metric_name:
+            for obj in objectives:
+                if metric_name == obj.metric_names[0]:
                     minimize = obj.minimize
                     break
             pp = self.ax_client.get_pareto_optimal_parameters()
@@ -157,7 +159,7 @@ def evaluate_model(
             else:
                 best_obj_i = np.argmax(obj_vals)
             best_pars = param_vals[best_obj_i]
-        else:
+        except AttributeError:
             best_pars = self.ax_client.get_best_parameters()[0]
 
         parameters = best_pars
@@ -234,7 +236,6 @@ def plot_model(
         yname: Optional[str] = None,
         mname: Optional[str] = None,
         p0: Optional[Dict] = None,
-        new: Optional[bool] = False,
         npoints: Optional[int] = 200,
         xrange: Optional[List[float]] = None,
         yrange: Optional[List[float]] = None,
@@ -255,7 +256,7 @@ def plot_model(
             Name of the variable to plot in y axis.
         mname: string, optional.
             Name of the metric to plot.
-            If not specified, it will take the first of the list in the ``AxClient``.
+            If not specified, it will take the first objective in ``self.ax_client``.
         p0: dictionary
             Particular values of parameters to be fixed for the evaluation over the sample.
         npoints: int, optional
@@ -328,7 +329,7 @@ def plot_model(
 
         # metric name
         if mname is None:
-            mname = list(experiment.metrics.keys())[0]
+            mname = self.ax_client.objective_names[0]
 
         # evaluate the model
         f_plt, sd_plt = self.evaluate_model(
diff --git a/optimas/diagnostics/exploration_diagnostics.py b/optimas/diagnostics/exploration_diagnostics.py
index 9914a1cb..f30c938a 100644
--- a/optimas/diagnostics/exploration_diagnostics.py
+++ b/optimas/diagnostics/exploration_diagnostics.py
@@ -1032,10 +1032,3 @@ def build_model(
         )
 
         return self.model_manager
-
-    def get_model_manager(self):
-        """Get the associated AxModelManager or build from scratch."""
-        if self.model_manager is None:
-            return self.build_model()
-        else:
-            return self.model_manager

From dd02114d74d0ca174ea4041d4a8b5c756829b257 Mon Sep 17 00:00:00 2001
From: delaossa <alberto.martinez.de.la.ossa@desy.de>
Date: Tue, 13 Feb 2024 13:40:50 +0100
Subject: [PATCH 10/54] Add test for `ax_model_manager`.

---
 tests/test_ax_model_manager.py | 131 +++++++++++++++++++++++++++++++++
 1 file changed, 131 insertions(+)
 create mode 100644 tests/test_ax_model_manager.py

diff --git a/tests/test_ax_model_manager.py b/tests/test_ax_model_manager.py
new file mode 100644
index 00000000..364a7c8e
--- /dev/null
+++ b/tests/test_ax_model_manager.py
@@ -0,0 +1,131 @@
+import os
+import numpy as np
+import matplotlib.pyplot as plt
+from matplotlib.gridspec import GridSpec
+import pytest
+
+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)
+    obj = Objective("f", minimize=True)
+    obj2 = Objective("f2", minimize=False)
+
+    gen = AxSingleFidelityGenerator(
+        varying_parameters=[var1, var2], 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]
+
+    # 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_model(objname='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) \
+         and np.allclose(mean_axcl, mean_diag, rtol=1e-2)
+    
+    # Make example figure with two models in 2D.
+    fig = plt.figure(figsize=(10, 4.8))
+    gs = GridSpec(1, 2, wspace=0.2, hspace=0.3)
+
+    # plot model for `f`
+    ax1 = mm_axcl.plot_model(mname='f', 
+                            pcolormesh_kw={'cmap': 'GnBu'}, 
+                            subplot_spec=gs[0, 0])
+    
+    # Get and draw top 3 evaluations for `f`
+    top_f = diags.get_best_evaluations_index(top=3, objective='f')
+    df_top = diags.history.loc[top_f][varpar_names]
+    ax1.scatter(df_top['x0'], df_top['x1'], c='red', marker='x')
+
+    # plot model for `f2`
+    ax2 = mm_axcl.plot_model(mname='f2', 
+                           pcolormesh_kw={'cmap': 'OrRd'}, 
+                           subplot_spec=gs[0, 1])
+    
+    # Get and draw top 3 evaluations for `f`
+    top_f2 = diags.get_best_evaluations_index(top=3, objective='f2')
+    df2_top = diags.history.loc[top_f2][varpar_names]
+    ax2.scatter(df2_top['x0'], df2_top['x1'], c='blue', marker='x')
+    plt.savefig(os.path.join(exploration_dir_path, "models.png"))
+
+    # Make example figure of the models in 1D with errors.
+    parx0 = mm_axcl.ax_client.experiment.parameters['x0']
+    parx1 = mm_axcl.ax_client.experiment.parameters['x1']
+    x1_mid = 0.5 * (parx1.lower + parx1.upper)
+    x0 = np.linspace(parx0.lower, parx0.upper, 100)
+    metric_names = mm_axcl.ax_client.objective_names
+    fig, axs = plt.subplots(len(metric_names), 1,
+                            sharex=True)
+    for i, (ax, metric_name) in enumerate(zip(axs, metric_names)):
+        mean, sed = mm_axcl.evaluate_model(sample={'x0': x0},
+                                        p0={'x1': x1_mid},
+                                        metric_name=metric_name)
+        ax.plot(x0, mean, color=f'C{i}', label=f'x1 = {x1_mid}')
+        ax.fill_between(x0, mean - sed, mean + sed, color='lightgray', alpha=0.5)
+        ax.set_ylabel(metric_name)
+        ax.legend(frameon=False)
+    plt.xlabel('x0')
+    plt.savefig(os.path.join(exploration_dir_path, "models_1d.png"))
+
+
+if __name__ == "__main__":
+    test_ax_model_manager()

From 4ea48c134953ea2cab824f1a083f3787d151d4ca Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Tue, 13 Feb 2024 12:43:47 +0000
Subject: [PATCH 11/54] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 optimas/diagnostics/ax_model_manager.py |  2 +-
 tests/test_ax_model_manager.py          | 72 +++++++++++++------------
 2 files changed, 39 insertions(+), 35 deletions(-)

diff --git a/optimas/diagnostics/ax_model_manager.py b/optimas/diagnostics/ax_model_manager.py
index 10fdecba..48a19da7 100644
--- a/optimas/diagnostics/ax_model_manager.py
+++ b/optimas/diagnostics/ax_model_manager.py
@@ -104,7 +104,7 @@ def build_model(
     def evaluate_model(
         self,
         sample: Union[DataFrame, Dict, NDArray] = None,
-        metric_name: Optional[str] = None, 
+        metric_name: Optional[str] = None,
         p0: Optional[Dict] = None,
     ) -> Tuple[NDArray]:
         """Evaluate the model over the specified sample.
diff --git a/tests/test_ax_model_manager.py b/tests/test_ax_model_manager.py
index 364a7c8e..1d501dce 100644
--- a/tests/test_ax_model_manager.py
+++ b/tests/test_ax_model_manager.py
@@ -58,72 +58,76 @@ def test_ax_model_manager():
 
     # 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')
+    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')
+    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')
+    mean_json, sem_json = mm_json.evaluate_model(sample=df, metric_name="f")
 
     # Get model manager from diagnostics data.
-    mm_diag = diags.build_model(objname='f')
-    mean_diag, sem_diag = mm_diag.evaluate_model(sample=df, metric_name='f')
+    mm_diag = diags.build_model(objname="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
+    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) \
-         and np.allclose(mean_axcl, mean_diag, rtol=1e-2)
-    
+    assert np.allclose(mean_axcl, mean_json, rtol=1e-2) and np.allclose(
+        mean_axcl, mean_diag, rtol=1e-2
+    )
+
     # Make example figure with two models in 2D.
     fig = plt.figure(figsize=(10, 4.8))
     gs = GridSpec(1, 2, wspace=0.2, hspace=0.3)
 
     # plot model for `f`
-    ax1 = mm_axcl.plot_model(mname='f', 
-                            pcolormesh_kw={'cmap': 'GnBu'}, 
-                            subplot_spec=gs[0, 0])
-    
+    ax1 = mm_axcl.plot_model(
+        mname="f", pcolormesh_kw={"cmap": "GnBu"}, subplot_spec=gs[0, 0]
+    )
+
     # Get and draw top 3 evaluations for `f`
-    top_f = diags.get_best_evaluations_index(top=3, objective='f')
+    top_f = diags.get_best_evaluations_index(top=3, objective="f")
     df_top = diags.history.loc[top_f][varpar_names]
-    ax1.scatter(df_top['x0'], df_top['x1'], c='red', marker='x')
+    ax1.scatter(df_top["x0"], df_top["x1"], c="red", marker="x")
 
     # plot model for `f2`
-    ax2 = mm_axcl.plot_model(mname='f2', 
-                           pcolormesh_kw={'cmap': 'OrRd'}, 
-                           subplot_spec=gs[0, 1])
-    
+    ax2 = mm_axcl.plot_model(
+        mname="f2", pcolormesh_kw={"cmap": "OrRd"}, subplot_spec=gs[0, 1]
+    )
+
     # Get and draw top 3 evaluations for `f`
-    top_f2 = diags.get_best_evaluations_index(top=3, objective='f2')
+    top_f2 = diags.get_best_evaluations_index(top=3, objective="f2")
     df2_top = diags.history.loc[top_f2][varpar_names]
-    ax2.scatter(df2_top['x0'], df2_top['x1'], c='blue', marker='x')
+    ax2.scatter(df2_top["x0"], df2_top["x1"], c="blue", marker="x")
     plt.savefig(os.path.join(exploration_dir_path, "models.png"))
 
     # Make example figure of the models in 1D with errors.
-    parx0 = mm_axcl.ax_client.experiment.parameters['x0']
-    parx1 = mm_axcl.ax_client.experiment.parameters['x1']
+    parx0 = mm_axcl.ax_client.experiment.parameters["x0"]
+    parx1 = mm_axcl.ax_client.experiment.parameters["x1"]
     x1_mid = 0.5 * (parx1.lower + parx1.upper)
     x0 = np.linspace(parx0.lower, parx0.upper, 100)
     metric_names = mm_axcl.ax_client.objective_names
-    fig, axs = plt.subplots(len(metric_names), 1,
-                            sharex=True)
+    fig, axs = plt.subplots(len(metric_names), 1, sharex=True)
     for i, (ax, metric_name) in enumerate(zip(axs, metric_names)):
-        mean, sed = mm_axcl.evaluate_model(sample={'x0': x0},
-                                        p0={'x1': x1_mid},
-                                        metric_name=metric_name)
-        ax.plot(x0, mean, color=f'C{i}', label=f'x1 = {x1_mid}')
-        ax.fill_between(x0, mean - sed, mean + sed, color='lightgray', alpha=0.5)
+        mean, sed = mm_axcl.evaluate_model(
+            sample={"x0": x0}, p0={"x1": x1_mid}, metric_name=metric_name
+        )
+        ax.plot(x0, mean, color=f"C{i}", label=f"x1 = {x1_mid}")
+        ax.fill_between(
+            x0, mean - sed, mean + sed, color="lightgray", alpha=0.5
+        )
         ax.set_ylabel(metric_name)
         ax.legend(frameon=False)
-    plt.xlabel('x0')
+    plt.xlabel("x0")
     plt.savefig(os.path.join(exploration_dir_path, "models_1d.png"))
 
 

From 13e1aed2af917dc0b0f46f5f621ba6d91eeda3a3 Mon Sep 17 00:00:00 2001
From: delaossa <alberto.martinez.de.la.ossa@desy.de>
Date: Tue, 13 Feb 2024 18:10:24 +0100
Subject: [PATCH 12/54] Fix line in documentation.

---
 optimas/diagnostics/ax_model_manager.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/optimas/diagnostics/ax_model_manager.py b/optimas/diagnostics/ax_model_manager.py
index 48a19da7..47b42dd8 100644
--- a/optimas/diagnostics/ax_model_manager.py
+++ b/optimas/diagnostics/ax_model_manager.py
@@ -55,7 +55,6 @@ def build_model(
         self, objname: str, parnames: List[str], minimize: Optional[bool] = True
     ) -> None:
         """Initialize the AxClient and the model using the given data.
-            This method only works with one objective.
 
         Parameters
         ----------

From 6d0b3d6b8e19e5f3ff0e186c5fec342d565280de Mon Sep 17 00:00:00 2001
From: delaossa <alberto.martinez.de.la.ossa@desy.de>
Date: Tue, 13 Feb 2024 21:17:17 +0100
Subject: [PATCH 13/54] Allow to initialize multi-objective models.

---
 optimas/diagnostics/ax_model_manager.py       | 92 +++++++++++++------
 .../diagnostics/exploration_diagnostics.py    | 56 ++++++-----
 2 files changed, 99 insertions(+), 49 deletions(-)

diff --git a/optimas/diagnostics/ax_model_manager.py b/optimas/diagnostics/ax_model_manager.py
index 47b42dd8..9738b734 100644
--- a/optimas/diagnostics/ax_model_manager.py
+++ b/optimas/diagnostics/ax_model_manager.py
@@ -8,6 +8,7 @@
 import matplotlib.pyplot as plt
 from matplotlib.gridspec import GridSpec, GridSpecFromSubplotSpec, SubplotSpec
 from matplotlib.axes import Axes
+from optimas.core import VaryingParameter, Objective
 
 # Ax utilities for model building
 from ax.service.ax_client import AxClient
@@ -18,6 +19,7 @@
 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
 
 
 class AxModelManager(object):
@@ -52,7 +54,12 @@ def model(self) -> TorchModelBridge:
         return self.ax_client.generation_strategy.model
 
     def build_model(
-        self, objname: str, parnames: List[str], minimize: Optional[bool] = True
+        self, 
+        parnames: Optional[List[str]] = None, 
+        objname: Optional[str] = None, 
+        minimize: Optional[bool] = True,
+        parameters: Optional[List[VaryingParameter]] = None,
+        objectives: Optional[List[Objective]] = None,
     ) -> None:
         """Initialize the AxClient and the model using the given data.
 
@@ -60,41 +67,74 @@ def build_model(
         ----------
         parnames: list of string
             List with the names of the parameters of the model.
-        objname: string
+        objname: string, optional.
             Name of the objective.
-        minimize: bool
+        minimize: bool, optional.
             Whether to minimize or maximize the objective.
             Only relevant to establish the best point.
+        objectives: list of `Objective`.
+            This is the way to build multi-objective models.
+            Useful to initialize from `ExplorationDiagnostics`.
+        parameters: list of `VaryingParameter`.
+            Useful to initialize from `ExplorationDiagnostics`.
         """
-        parameters = [
-            {
-                "name": p_name,
-                "type": "range",
-                "bounds": [self.df[p_name].min(), self.df[p_name].max()],
-                "value_type": "float",
-            }
-            for p_name in parnames
-        ]
-
-        # create Ax client
-        gs = GenerationStrategy(
-            [GenerationStep(model=Models.GPEI, num_trials=-1)]
-        )
-        self.ax_client = AxClient(generation_strategy=gs, verbose_logging=False)
-        self.ax_client.create_experiment(
-            name="optimas_data",
-            parameters=parameters,
-            objective_name=objname,
-            minimize=minimize,
-        )
+        if parameters is None:
+            if parnames is None:
+                raise RuntimeError(
+                    "Either `parameters` or `parnames` should be provided."
+                )
+            parameters = [
+                {
+                    "name": p_name,
+                    "type": "range",
+                    "bounds": [self.df[p_name].min(), self.df[p_name].max()],
+                    "value_type": "float",
+                }
+                for p_name in parnames
+            ]
+        else:  
+            parnames = [par['name'] for par in parameters]
+    
+        if objectives is not None:
+            objectives_dict = {}
+            for obj in objectives:
+                objectives_dict[obj.name] = ObjectiveProperties(minimize=obj.minimize)
+            # create Ax client
+            gs = GenerationStrategy(
+                [GenerationStep(model=Models.MOO, num_trials=-1)]
+            )
+            self.ax_client = AxClient(generation_strategy=gs, verbose_logging=False)
+            self.ax_client.create_experiment(
+                name="optimas_data",
+                parameters=parameters,
+                objectives=objectives_dict,
+            )
+        elif objname is not None:
+            # create Ax client
+            gs = GenerationStrategy(
+                [GenerationStep(model=Models.GPEI, num_trials=-1)]
+            )
+            self.ax_client = AxClient(generation_strategy=gs, verbose_logging=False)
+            self.ax_client.create_experiment(
+                name="optimas_data",
+                parameters=parameters,
+                objective_name=objname,
+                minimize=minimize,
+            )
+        else:
+            raise RuntimeError(
+                "Either `objectives` or `objname` should be provided."
+            )
 
         # adds data
-        metric_name = list(self.ax_client.experiment.metrics.keys())[0]
         for index, row in self.df.iterrows():
             params = {p_name: row[p_name] for p_name in parnames}
             _, trial_id = self.ax_client.attach_trial(params)
+            data = {}
+            for mname in list(self.ax_client.experiment.metrics.keys()):
+                data[mname] = (row[mname], np.nan)
             self.ax_client.complete_trial(
-                trial_id, {metric_name: (row[metric_name], np.nan)}
+                trial_id, raw_data=data
             )
 
         # fit GP model
diff --git a/optimas/diagnostics/exploration_diagnostics.py b/optimas/diagnostics/exploration_diagnostics.py
index f30c938a..c682d5b3 100644
--- a/optimas/diagnostics/exploration_diagnostics.py
+++ b/optimas/diagnostics/exploration_diagnostics.py
@@ -20,6 +20,7 @@
 from optimas.utils.other import get_df_with_selection
 from optimas.diagnostics.ax_model_manager import AxModelManager
 from ax.service.ax_client import AxClient
+from ax.service.utils.instantiation import ObjectiveProperties
 
 
 class ExplorationDiagnostics:
@@ -990,8 +991,7 @@ def get_model_manager_from_ax_client(
 
     def build_model(
         self,
-        objname: Optional[str] = "",
-        parnames: Optional[List[str]] = [],
+        objname: Optional[str] = None,
         minimize: Optional[bool] = True,
     ) -> AxModelManager:
         """Initialize AxModelManager and builds a GP model.
@@ -1000,35 +1000,45 @@ def build_model(
         ----------
         objname: string, optional
             Name of the objective (or metric).
-            If not given, it takes the first of the objectives.
-        parnames: list of string, optional
-            List with the names of the parameters of the model.
-            If not given, it assumes the varying parameters.
+            If not given, it takes the list of objectives.
         minimize: bool, optional
             Whether to minimize or maximize the objective.
-            Only relevant to establish the best point of the model,
-            but not to build the model.
+            It is only used when `objname` is given.
+            Only relevant to establish the best point of the model.
 
         Returns
         -------
         An instance of AxModelManager
         """
-        varpar_names = [var.name for var in self.varying_parameters]
-        objective_names = [obj.name for obj in self.objectives]
-
-        if len(parnames) == 0:
-            parnames = varpar_names
-        if objname == "":
-            objname = objective_names[0]
-
-        if objname in objective_names:
-            minimize = self._get_objective(objname).minimize
-
-        # Copy the history DataFrame
+        # Initialize `AxModelManager` with history data
         df = self.history.copy()
         self.model_manager = AxModelManager(df)
-        self.model_manager.build_model(
-            parnames=parnames, objname=objname, minimize=minimize
-        )
+
+        # Get parameters for the model
+        parameters = []
+        for var in self.varying_parameters:
+            parameters.append(
+                {
+                    "name": var.name,
+                    "type": "range",
+                    "bounds": [var.lower_bound, var.upper_bound],
+                    "is_fidelity": var.is_fidelity,
+                    "target_value": var.fidelity_target_value,
+                    "value_type": "float",
+                }
+            )
+
+        # Get objectives
+        objective_names = [obj.name for obj in self.objectives]
+        if objname is not None:
+            if objname in objective_names:
+                minimize = self._get_objective(objname).minimize
+            self.model_manager.build_model(
+                parameters=parameters, objname=objname, minimize=minimize
+            )
+        else:
+            self.model_manager.build_model(
+                parameters=parameters, objectives=self.objectives
+            )
 
         return self.model_manager

From fd59406900fde5cd74b990efc8365d4adc4701ff Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Tue, 13 Feb 2024 20:17:38 +0000
Subject: [PATCH 14/54] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 optimas/diagnostics/ax_model_manager.py | 28 ++++++++++++++-----------
 1 file changed, 16 insertions(+), 12 deletions(-)

diff --git a/optimas/diagnostics/ax_model_manager.py b/optimas/diagnostics/ax_model_manager.py
index 9738b734..42e2acc8 100644
--- a/optimas/diagnostics/ax_model_manager.py
+++ b/optimas/diagnostics/ax_model_manager.py
@@ -54,9 +54,9 @@ def model(self) -> TorchModelBridge:
         return self.ax_client.generation_strategy.model
 
     def build_model(
-        self, 
-        parnames: Optional[List[str]] = None, 
-        objname: Optional[str] = None, 
+        self,
+        parnames: Optional[List[str]] = None,
+        objname: Optional[str] = None,
         minimize: Optional[bool] = True,
         parameters: Optional[List[VaryingParameter]] = None,
         objectives: Optional[List[Objective]] = None,
@@ -92,18 +92,22 @@ def build_model(
                 }
                 for p_name in parnames
             ]
-        else:  
-            parnames = [par['name'] for par in parameters]
-    
+        else:
+            parnames = [par["name"] for par in parameters]
+
         if objectives is not None:
             objectives_dict = {}
             for obj in objectives:
-                objectives_dict[obj.name] = ObjectiveProperties(minimize=obj.minimize)
+                objectives_dict[obj.name] = ObjectiveProperties(
+                    minimize=obj.minimize
+                )
             # create Ax client
             gs = GenerationStrategy(
                 [GenerationStep(model=Models.MOO, num_trials=-1)]
             )
-            self.ax_client = AxClient(generation_strategy=gs, verbose_logging=False)
+            self.ax_client = AxClient(
+                generation_strategy=gs, verbose_logging=False
+            )
             self.ax_client.create_experiment(
                 name="optimas_data",
                 parameters=parameters,
@@ -114,7 +118,9 @@ def build_model(
             gs = GenerationStrategy(
                 [GenerationStep(model=Models.GPEI, num_trials=-1)]
             )
-            self.ax_client = AxClient(generation_strategy=gs, verbose_logging=False)
+            self.ax_client = AxClient(
+                generation_strategy=gs, verbose_logging=False
+            )
             self.ax_client.create_experiment(
                 name="optimas_data",
                 parameters=parameters,
@@ -133,9 +139,7 @@ def build_model(
             data = {}
             for mname in list(self.ax_client.experiment.metrics.keys()):
                 data[mname] = (row[mname], np.nan)
-            self.ax_client.complete_trial(
-                trial_id, raw_data=data
-            )
+            self.ax_client.complete_trial(trial_id, raw_data=data)
 
         # fit GP model
         self.ax_client.fit_model()

From d40fd6d5f65665bd6bb4b027de0ff7f0e28a2bc0 Mon Sep 17 00:00:00 2001
From: delaossa <alberto.martinez.de.la.ossa@desy.de>
Date: Wed, 14 Feb 2024 12:12:01 +0100
Subject: [PATCH 15/54] Change the logix of .

---
 optimas/diagnostics/ax_model_manager.py | 38 +++++++++++++------------
 1 file changed, 20 insertions(+), 18 deletions(-)

diff --git a/optimas/diagnostics/ax_model_manager.py b/optimas/diagnostics/ax_model_manager.py
index 42e2acc8..b98ed17f 100644
--- a/optimas/diagnostics/ax_model_manager.py
+++ b/optimas/diagnostics/ax_model_manager.py
@@ -78,6 +78,7 @@ def build_model(
         parameters: list of `VaryingParameter`.
             Useful to initialize from `ExplorationDiagnostics`.
         """
+        # Define parameters
         if parameters is None:
             if parnames is None:
                 raise RuntimeError(
@@ -95,15 +96,15 @@ def build_model(
         else:
             parnames = [par["name"] for par in parameters]
 
-        if objectives is not None:
-            objectives_dict = {}
-            for obj in objectives:
-                objectives_dict[obj.name] = ObjectiveProperties(
-                    minimize=obj.minimize
+        # Define objectives
+        if objectives is None:
+            if objname is None:
+                raise RuntimeError(
+                    "Either `objectives` or `objname` should be provided."
                 )
-            # create Ax client
+            # Create single objective AxClient
             gs = GenerationStrategy(
-                [GenerationStep(model=Models.MOO, num_trials=-1)]
+                [GenerationStep(model=Models.GPEI, num_trials=-1)]
             )
             self.ax_client = AxClient(
                 generation_strategy=gs, verbose_logging=False
@@ -111,12 +112,18 @@ def build_model(
             self.ax_client.create_experiment(
                 name="optimas_data",
                 parameters=parameters,
-                objectives=objectives_dict,
+                objective_name=objname,
+                minimize=minimize,
             )
-        elif objname is not None:
-            # create Ax client
+        else:
+            # Create multi-objective Axclient
+            objectives_dict = {}
+            for obj in objectives:
+                objectives_dict[obj.name] = ObjectiveProperties(
+                    minimize=obj.minimize
+                )
             gs = GenerationStrategy(
-                [GenerationStep(model=Models.GPEI, num_trials=-1)]
+                [GenerationStep(model=Models.MOO, num_trials=-1)]
             )
             self.ax_client = AxClient(
                 generation_strategy=gs, verbose_logging=False
@@ -124,15 +131,10 @@ def build_model(
             self.ax_client.create_experiment(
                 name="optimas_data",
                 parameters=parameters,
-                objective_name=objname,
-                minimize=minimize,
-            )
-        else:
-            raise RuntimeError(
-                "Either `objectives` or `objname` should be provided."
+                objectives=objectives_dict,
             )
 
-        # adds data
+        # Add trials from DataFrame
         for index, row in self.df.iterrows():
             params = {p_name: row[p_name] for p_name in parnames}
             _, trial_id = self.ax_client.attach_trial(params)

From ad85ed47d1772d4675ed0dd6577e769f17be9944 Mon Sep 17 00:00:00 2001
From: delaossa <alberto.martinez.de.la.ossa@desy.de>
Date: Wed, 14 Feb 2024 12:12:01 +0100
Subject: [PATCH 16/54] Change the logic of .

---
 optimas/diagnostics/ax_model_manager.py | 38 +++++++++++++------------
 1 file changed, 20 insertions(+), 18 deletions(-)

diff --git a/optimas/diagnostics/ax_model_manager.py b/optimas/diagnostics/ax_model_manager.py
index 42e2acc8..b98ed17f 100644
--- a/optimas/diagnostics/ax_model_manager.py
+++ b/optimas/diagnostics/ax_model_manager.py
@@ -78,6 +78,7 @@ def build_model(
         parameters: list of `VaryingParameter`.
             Useful to initialize from `ExplorationDiagnostics`.
         """
+        # Define parameters
         if parameters is None:
             if parnames is None:
                 raise RuntimeError(
@@ -95,15 +96,15 @@ def build_model(
         else:
             parnames = [par["name"] for par in parameters]
 
-        if objectives is not None:
-            objectives_dict = {}
-            for obj in objectives:
-                objectives_dict[obj.name] = ObjectiveProperties(
-                    minimize=obj.minimize
+        # Define objectives
+        if objectives is None:
+            if objname is None:
+                raise RuntimeError(
+                    "Either `objectives` or `objname` should be provided."
                 )
-            # create Ax client
+            # Create single objective AxClient
             gs = GenerationStrategy(
-                [GenerationStep(model=Models.MOO, num_trials=-1)]
+                [GenerationStep(model=Models.GPEI, num_trials=-1)]
             )
             self.ax_client = AxClient(
                 generation_strategy=gs, verbose_logging=False
@@ -111,12 +112,18 @@ def build_model(
             self.ax_client.create_experiment(
                 name="optimas_data",
                 parameters=parameters,
-                objectives=objectives_dict,
+                objective_name=objname,
+                minimize=minimize,
             )
-        elif objname is not None:
-            # create Ax client
+        else:
+            # Create multi-objective Axclient
+            objectives_dict = {}
+            for obj in objectives:
+                objectives_dict[obj.name] = ObjectiveProperties(
+                    minimize=obj.minimize
+                )
             gs = GenerationStrategy(
-                [GenerationStep(model=Models.GPEI, num_trials=-1)]
+                [GenerationStep(model=Models.MOO, num_trials=-1)]
             )
             self.ax_client = AxClient(
                 generation_strategy=gs, verbose_logging=False
@@ -124,15 +131,10 @@ def build_model(
             self.ax_client.create_experiment(
                 name="optimas_data",
                 parameters=parameters,
-                objective_name=objname,
-                minimize=minimize,
-            )
-        else:
-            raise RuntimeError(
-                "Either `objectives` or `objname` should be provided."
+                objectives=objectives_dict,
             )
 
-        # adds data
+        # Add trials from DataFrame
         for index, row in self.df.iterrows():
             params = {p_name: row[p_name] for p_name in parnames}
             _, trial_id = self.ax_client.attach_trial(params)

From e0927047e8954a116b21aeebcf94d29f0291b8cc Mon Sep 17 00:00:00 2001
From: delaossa <alberto.martinez.de.la.ossa@desy.de>
Date: Wed, 14 Feb 2024 12:55:44 +0100
Subject: [PATCH 17/54] Convert parameter type properly.

---
 optimas/diagnostics/exploration_diagnostics.py | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/optimas/diagnostics/exploration_diagnostics.py b/optimas/diagnostics/exploration_diagnostics.py
index c682d5b3..4c18bf74 100644
--- a/optimas/diagnostics/exploration_diagnostics.py
+++ b/optimas/diagnostics/exploration_diagnostics.py
@@ -1014,9 +1014,21 @@ def build_model(
         df = self.history.copy()
         self.model_manager = AxModelManager(df)
 
-        # Get parameters for the model
+        # Get parameters for AxClient.
         parameters = []
         for var in self.varying_parameters:
+            # Determine parameter type.
+            value_dtype = np.dtype(var.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.
             parameters.append(
                 {
                     "name": var.name,
@@ -1024,7 +1036,7 @@ def build_model(
                     "bounds": [var.lower_bound, var.upper_bound],
                     "is_fidelity": var.is_fidelity,
                     "target_value": var.fidelity_target_value,
-                    "value_type": "float",
+                    "value_type": value_type,
                 }
             )
 

From faa44b812ec47612815dd2fc1b74a95a19aa2130 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=81ngel=20Ferran=20Pousa?= <angel.ferran.pousa@desy.de>
Date: Thu, 15 Feb 2024 22:13:30 +0100
Subject: [PATCH 18/54] Refactor class methods

---
 optimas/diagnostics/ax_model_manager.py | 165 +++++++-----------------
 1 file changed, 50 insertions(+), 115 deletions(-)

diff --git a/optimas/diagnostics/ax_model_manager.py b/optimas/diagnostics/ax_model_manager.py
index b98ed17f..76e98d62 100644
--- a/optimas/diagnostics/ax_model_manager.py
+++ b/optimas/diagnostics/ax_model_manager.py
@@ -4,11 +4,11 @@
 import numpy as np
 from numpy.typing import NDArray
 from pandas import DataFrame
-from copy import deepcopy
 import matplotlib.pyplot as plt
 from matplotlib.gridspec import GridSpec, GridSpecFromSubplotSpec, SubplotSpec
 from matplotlib.axes import Axes
 from optimas.core import VaryingParameter, Objective
+from optimas.utils.other import convert_to_dataframe
 
 # Ax utilities for model building
 from ax.service.ax_client import AxClient
@@ -150,7 +150,6 @@ def evaluate_model(
         self,
         sample: Union[DataFrame, Dict, NDArray] = None,
         metric_name: Optional[str] = None,
-        p0: Optional[Dict] = None,
     ) -> Tuple[NDArray]:
         """Evaluate the model over the specified sample.
 
@@ -165,8 +164,6 @@ def evaluate_model(
         metric_name: str, optional.
             Name of the metric to evaluate.
             If not specified, it will take the first first objective in ``self.ax_client``.
-        p0: dictionary, optional.
-            Particular values of parameters to be fixed for the evaluation.
 
         Returns
         -------
@@ -186,89 +183,21 @@ def evaluate_model(
                     f"Available metrics are: {metric_names}."
                 )
 
-        # get optimum
-        try:
-            objectives = self.ax_client.objective.objectives
-            minimize = None
-            for obj in objectives:
-                if metric_name == obj.metric_names[0]:
-                    minimize = obj.minimize
-                    break
-            pp = self.ax_client.get_pareto_optimal_parameters()
-            obj_vals = [
-                objs[metric_name] for i, (vals, (objs, covs)) in pp.items()
-            ]
-            param_vals = [vals for i, (vals, (objs, covs)) in pp.items()]
-            if minimize:
-                best_obj_i = np.argmin(obj_vals)
-            else:
-                best_obj_i = np.argmax(obj_vals)
-            best_pars = param_vals[best_obj_i]
-        except AttributeError:
-            best_pars = self.ax_client.get_best_parameters()[0]
-
-        parameters = best_pars
-        parnames = list(parameters.keys())
+        parnames = list(self.ax_client.experiment.parameters.keys())
 
-        # user specific point
-        if p0 is not None:
-            for key in p0.keys():
-                if key in parameters.keys():
-                    parameters[key] = p0[key]
+        sample = convert_to_dataframe(sample)
 
+        # 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 = []
-        obsf_0 = ObservationFeatures(parameters=parameters)
-        if isinstance(sample, np.ndarray):
-            # check the shape of the array
-            if sample.shape[1] != len(parnames):
-                raise RuntimeError(
-                    "Second dimension of the sample array should match "
-                    "the number of parameters of the model."
-                )
-            for i in range(sample.shape[0]):
-                predf = deepcopy(obsf_0)
-                for j, parname in enumerate(parameters.keys()):
-                    predf.parameters[parname] = sample[i][j]
-                obsf_list.append(predf)
-        elif isinstance(sample, DataFrame):
-            # check if labels of the dataframe match the parnames
-            for col in sample.columns:
-                if col not in parnames:
-                    raise RuntimeError(
-                        "Column %s does not match any of the parameter names"
-                        % col
-                    )
-            for i in range(sample.shape[0]):
-                predf = deepcopy(obsf_0)
-                for col in sample.columns:
-                    predf.parameters[col] = sample[col].iloc[i]
-                obsf_list.append(predf)
-        elif isinstance(sample, dict):
-            # check if the keys of the dictionary match the parnames
-            for key in sample.keys():
-                if key not in parnames:
-                    raise RuntimeError(
-                        "Key %s does not match any of the parameter names" % col
-                    )
-            element = sample[list(sample.keys())[0]]
-            if hasattr(element, "__len__"):
-                for i in range(len(element)):
-                    predf = deepcopy(obsf_0)
-                    for key in sample.keys():
-                        predf.parameters[key] = sample[key][i]
-                    obsf_list.append(predf)
-            else:
-                predf = deepcopy(obsf_0)
-                for key in sample.keys():
-                    predf.parameters[key] = sample[key]
-                obsf_list.append(predf)
-
-        elif sample is None:
-            predf = deepcopy(obsf_0)
-            obsf_list.append(predf)
-        else:
-            raise RuntimeError("Wrong data type")
+        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])
@@ -287,8 +216,8 @@ def plot_model(
         mode: Optional[Literal["mean", "sem", "both"]] = "mean",
         clabel: Optional[bool] = False,
         subplot_spec: Optional[SubplotSpec] = None,
-        gridspec_kw: Dict[str, Any] = None,
-        pcolormesh_kw: Dict[str, Any] = None,
+        gridspec_kw: Optional[Dict[str, Any]] = None,
+        pcolormesh_kw: Optional[Dict[str, Any]] = None,
         **figure_kw,
     ) -> Union[Axes, List[Axes]]:
         """Plot model in the two selected variables, while others are fixed to the optimum.
@@ -334,7 +263,6 @@ def plot_model(
         # get experiment info
         experiment = self.ax_client.experiment
         parnames = list(experiment.parameters.keys())
-        # minimize = experiment.optimization_config.objective.minimize
 
         if len(parnames) < 2:
             raise RuntimeError(
@@ -348,6 +276,10 @@ def plot_model(
         if yname is None:
             yname = parnames[1]
 
+        # metric name
+        if mname is None:
+            mname = self.ax_client.objective_names[0]
+
         # set the plotting range
         if xrange is None:
             xrange = [None, None]
@@ -366,36 +298,44 @@ def plot_model(
         xaxis = np.linspace(xrange[0], xrange[1], npoints)
         yaxis = np.linspace(yrange[0], yrange[1], npoints)
         X, Y = np.meshgrid(xaxis, yaxis)
-        xarray = X.flatten()
-        yarray = Y.flatten()
-        sample = DataFrame({xname: xarray, yname: yarray})
+        sample = {xname: X.flatten(), yname: Y.flatten()}
+
+        if p0 is None:
+            # get optimum
+            if len(self.ax_client.objective_names) > 1:
+                minimize = None
+                for obj in self.ax_client.objective.objectives:
+                    if mname == obj.metric_names[0]:
+                        minimize = obj.minimize
+                        break
+                pp = self.ax_client.get_pareto_optimal_parameters()
+                obj_vals = [
+                    objs[mname] for i, (vals, (objs, covs)) in pp.items()
+                ]
+                param_vals = [vals for i, (vals, (objs, covs)) in pp.items()]
+                if minimize:
+                    best_obj_i = np.argmin(obj_vals)
+                else:
+                    best_obj_i = np.argmax(obj_vals)
+                p0 = param_vals[best_obj_i]
+            else:
+                p0 = self.ax_client.get_best_parameters()[0]
 
-        # metric name
-        if mname is None:
-            mname = self.ax_client.objective_names[0]
+        for name, val in p0.items():
+            if name not in [xname, yname]:
+                sample[name] = np.ones(npoints**2) * val
 
         # evaluate the model
-        f_plt, sd_plt = self.evaluate_model(
-            sample=sample,
-            metric_name=mname,
-            p0=p0,
-        )
-
-        # get numpy arrays with experiment parameters
-        xtrials = np.zeros(experiment.num_trials)
-        ytrials = np.zeros(experiment.num_trials)
-        for i in range(experiment.num_trials):
-            xtrials[i] = experiment.trials[i].arm.parameters[xname]
-            ytrials[i] = experiment.trials[i].arm.parameters[yname]
+        f_plt, sd_plt = self.evaluate_model(sample=sample, metric_name=mname)
 
         # select quantities to plot and set the labels
         f_plots = []
         labels = []
         if mode in ["mean", "both"]:
-            f_plots.append(f_plt)
+            f_plots.append(f_plt.reshape(X.shape))
             labels.append(mname + ", mean")
         if mode in ["sem", "both"]:
-            f_plots.append(sd_plt)
+            f_plots.append(sd_plt.reshape(X.shape))
             labels.append(mname + ", sem")
 
         # create figure
@@ -409,18 +349,13 @@ def plot_model(
             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.reshape(X.shape),
-                shading="auto",
-                **pcolormesh_kw,
-            )
+            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=xname, ylabel=yname)
@@ -428,7 +363,7 @@ def plot_model(
             cset = ax.contour(
                 X,
                 Y,
-                f.reshape(X.shape),
+                f,
                 levels=20,
                 linewidths=0.5,
                 colors="black",
@@ -437,7 +372,7 @@ def plot_model(
             if clabel:
                 plt.clabel(cset, inline=True, fmt="%1.1f", fontsize="xx-small")
             # draw trials
-            ax.scatter(xtrials, ytrials, s=2, c="black", marker="o")
+            ax.scatter(trials[xname], trials[yname], s=2, c="black", marker="o")
             ax.set_xlim(xrange)
             ax.set_ylim(yrange)
             axs.append(ax)

From 4ffd81cecea49b72f8152bdef93522a070bb0863 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=81ngel=20Ferran=20Pousa?= <angel.ferran.pousa@desy.de>
Date: Thu, 15 Feb 2024 22:13:47 +0100
Subject: [PATCH 19/54] Update test

---
 tests/test_ax_model_manager.py | 19 +++++++++----------
 1 file changed, 9 insertions(+), 10 deletions(-)

diff --git a/tests/test_ax_model_manager.py b/tests/test_ax_model_manager.py
index 1d501dce..60e7f950 100644
--- a/tests/test_ax_model_manager.py
+++ b/tests/test_ax_model_manager.py
@@ -30,11 +30,12 @@ def test_ax_model_manager():
 
     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], objectives=[obj, obj2]
+        varying_parameters=[var1, var2, var3], objectives=[obj, obj2]
     )
     ev = FunctionEvaluator(function=eval_func_sf_moo)
     exploration = Exploration(
@@ -81,9 +82,8 @@ def test_ax_model_manager():
     print(df)
 
     # Check that different model initializations match within a 1% tolerance.
-    assert np.allclose(mean_axcl, mean_json, rtol=1e-2) and np.allclose(
-        mean_axcl, mean_diag, rtol=1e-2
-    )
+    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=(10, 4.8))
@@ -111,17 +111,16 @@ def test_ax_model_manager():
     plt.savefig(os.path.join(exploration_dir_path, "models.png"))
 
     # Make example figure of the models in 1D with errors.
-    parx0 = mm_axcl.ax_client.experiment.parameters["x0"]
-    parx1 = mm_axcl.ax_client.experiment.parameters["x1"]
-    x1_mid = 0.5 * (parx1.lower + parx1.upper)
-    x0 = np.linspace(parx0.lower, parx0.upper, 100)
+    x1 = np.ones(100) * 0.5 * (var2.lower_bound + var2.upper_bound)
+    x2 = np.ones(100) * 0.5 * (var3.lower_bound + var3.upper_bound)
+    x0 = np.linspace(var1.lower_bound, var1.upper_bound, 100)
     metric_names = mm_axcl.ax_client.objective_names
     fig, axs = plt.subplots(len(metric_names), 1, sharex=True)
     for i, (ax, metric_name) in enumerate(zip(axs, metric_names)):
         mean, sed = mm_axcl.evaluate_model(
-            sample={"x0": x0}, p0={"x1": x1_mid}, metric_name=metric_name
+            sample={"x0": x0, "x1": x1, "x2": x2}, metric_name=metric_name
         )
-        ax.plot(x0, mean, color=f"C{i}", label=f"x1 = {x1_mid}")
+        ax.plot(x0, mean, color=f"C{i}", label=f"x1 = {x1[0]}")
         ax.fill_between(
             x0, mean - sed, mean + sed, color="lightgray", alpha=0.5
         )

From 3781d6a9da6380099317d9acb55e323470d4463a Mon Sep 17 00:00:00 2001
From: delaossa <alberto.martinez.de.la.ossa@desy.de>
Date: Tue, 27 Feb 2024 10:36:33 +0100
Subject: [PATCH 20/54] Allow converting 1-element dictionary data to
 `DataFrame`.

---
 optimas/utils/other.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/optimas/utils/other.py b/optimas/utils/other.py
index be918558..3cce37a4 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())

From 6357f13c3c6d950a5aee060d33aae33ff6329ac9 Mon Sep 17 00:00:00 2001
From: delaossa <alberto.martinez.de.la.ossa@desy.de>
Date: Tue, 27 Feb 2024 21:20:23 +0100
Subject: [PATCH 21/54] Add function to get the best point parameters.

---
 optimas/diagnostics/ax_model_manager.py | 82 +++++++++++++++++--------
 1 file changed, 56 insertions(+), 26 deletions(-)

diff --git a/optimas/diagnostics/ax_model_manager.py b/optimas/diagnostics/ax_model_manager.py
index 76e98d62..22b69876 100644
--- a/optimas/diagnostics/ax_model_manager.py
+++ b/optimas/diagnostics/ax_model_manager.py
@@ -204,6 +204,57 @@ def evaluate_model(
         sem_array = np.sqrt(cov[metric_name][metric_name])
         return m_array, sem_array
 
+    def get_best_point(
+        self, 
+        metric_name: Optional[str] = None,
+        use_model_predictions: Optional[bool] = True,
+    ) -> Tuple[NDArray]:
+        """Get the best scoring point in the sample.
+
+         Parameter:
+        ----------
+        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
+        -------
+        best_point : dict
+            A dictionary with the parameters of the best point.       
+        """
+        # 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 = [objs[metric_name] for index, (vals, (objs, covs)) in pp.items()]
+            param_vals = [vals for index, (vals, (objs, covs)) in pp.items()]
+            if minimize:
+                best_obj_i = np.argmin(obj_vals)
+            else:
+                best_obj_i = np.argmax(obj_vals)
+            best_point = param_vals[best_obj_i]
+        else:
+            if use_model_predictions is True:
+                best_arm, best_point_predictions = self.model.model_best_point()
+                best_point = best_arm.parameters
+            else:
+                # AxClient.get_best_parameters seems to always return the best point
+                # from the observed values, independently of the value of `use_model_predictions`.
+                best_point = self.ax_client.get_best_parameters(use_model_predictions=use_model_predictions)[0]
+
+        return best_point
+
     def plot_model(
         self,
         xname: Optional[str] = None,
@@ -301,25 +352,8 @@ def plot_model(
         sample = {xname: X.flatten(), yname: Y.flatten()}
 
         if p0 is None:
-            # get optimum
-            if len(self.ax_client.objective_names) > 1:
-                minimize = None
-                for obj in self.ax_client.objective.objectives:
-                    if mname == obj.metric_names[0]:
-                        minimize = obj.minimize
-                        break
-                pp = self.ax_client.get_pareto_optimal_parameters()
-                obj_vals = [
-                    objs[mname] for i, (vals, (objs, covs)) in pp.items()
-                ]
-                param_vals = [vals for i, (vals, (objs, covs)) in pp.items()]
-                if minimize:
-                    best_obj_i = np.argmin(obj_vals)
-                else:
-                    best_obj_i = np.argmax(obj_vals)
-                p0 = param_vals[best_obj_i]
-            else:
-                p0 = self.ax_client.get_best_parameters()[0]
+            # get best point
+            p0 = self.get_best_point(metric_name=mname, use_model_predictions=True)
 
         for name, val in p0.items():
             if name not in [xname, yname]:
@@ -384,13 +418,13 @@ def plot_model(
 
     def get_arm_index(
         self,
-        arm_name: Optional[str] = None,
+        arm_name: str,
     ) -> int:
         """Get the index of the arm by its name.
 
         Parameters
         ----------
-        arm_name: string, optional.
+        arm_name: string.
             Name of the arm. If not given, the best arm is selected.
 
         Returns
@@ -398,10 +432,6 @@ def get_arm_index(
         index: int
             Trial index of the arm.
         """
-        if arm_name is None:
-            best_arm, best_point_predictions = self.model.model_best_point()
-            arm_name = best_arm.name
-
-        df = self.ax_client.experiment.fetch_data().df
+        df = self.ax_client.get_trials_data_frame()
         index = df.loc[df["arm_name"] == arm_name, "trial_index"].iloc[0]
         return index

From 88980525e3ff4b65dcc22fe11a77da6fb751ca01 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Tue, 27 Feb 2024 20:21:32 +0000
Subject: [PATCH 22/54] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 optimas/diagnostics/ax_model_manager.py | 24 ++++++++++++++++--------
 optimas/utils/other.py                  |  2 +-
 2 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/optimas/diagnostics/ax_model_manager.py b/optimas/diagnostics/ax_model_manager.py
index 22b69876..ba37bd76 100644
--- a/optimas/diagnostics/ax_model_manager.py
+++ b/optimas/diagnostics/ax_model_manager.py
@@ -205,7 +205,7 @@ def evaluate_model(
         return m_array, sem_array
 
     def get_best_point(
-        self, 
+        self,
         metric_name: Optional[str] = None,
         use_model_predictions: Optional[bool] = True,
     ) -> Tuple[NDArray]:
@@ -217,13 +217,13 @@ def get_best_point(
             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.    
+            Whether to extract the best point using model predictions
+            or directly observed values.
 
         Returns
         -------
         best_point : dict
-            A dictionary with the parameters of the best point.       
+            A dictionary with the parameters of the best point.
         """
         # metric name
         if metric_name is None:
@@ -236,8 +236,12 @@ def get_best_point(
                 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 = [objs[metric_name] for index, (vals, (objs, covs)) in pp.items()]
+            pp = self.ax_client.get_pareto_optimal_parameters(
+                use_model_predictions=use_model_predictions
+            )
+            obj_vals = [
+                objs[metric_name] for index, (vals, (objs, covs)) in pp.items()
+            ]
             param_vals = [vals for index, (vals, (objs, covs)) in pp.items()]
             if minimize:
                 best_obj_i = np.argmin(obj_vals)
@@ -251,7 +255,9 @@ def get_best_point(
             else:
                 # AxClient.get_best_parameters seems to always return the best point
                 # from the observed values, independently of the value of `use_model_predictions`.
-                best_point = self.ax_client.get_best_parameters(use_model_predictions=use_model_predictions)[0]
+                best_point = self.ax_client.get_best_parameters(
+                    use_model_predictions=use_model_predictions
+                )[0]
 
         return best_point
 
@@ -353,7 +359,9 @@ def plot_model(
 
         if p0 is None:
             # get best point
-            p0 = self.get_best_point(metric_name=mname, use_model_predictions=True)
+            p0 = self.get_best_point(
+                metric_name=mname, use_model_predictions=True
+            )
 
         for name, val in p0.items():
             if name not in [xname, yname]:
diff --git a/optimas/utils/other.py b/optimas/utils/other.py
index 3cce37a4..ef366737 100644
--- a/optimas/utils/other.py
+++ b/optimas/utils/other.py
@@ -62,7 +62,7 @@ def convert_to_dataframe(
         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 
+                data[key] = np.ones(1, dtype=type(value)) * value
         return pd.DataFrame(data)
     elif isinstance(data, list):
         fields = list(data[0].keys())

From d1317c2fe9927b49c600be31f10acddab36d7d0d Mon Sep 17 00:00:00 2001
From: delaossa <alberto.martinez.de.la.ossa@desy.de>
Date: Wed, 28 Feb 2024 16:38:20 +0100
Subject: [PATCH 23/54] Pass optimas parameters to
 `AxModelManager.build_model`.

---
 optimas/diagnostics/ax_model_manager.py       | 44 ++++++++++++++-----
 .../diagnostics/exploration_diagnostics.py    | 41 ++++-------------
 tests/test_ax_model_manager.py                | 20 ++++++---
 3 files changed, 55 insertions(+), 50 deletions(-)

diff --git a/optimas/diagnostics/ax_model_manager.py b/optimas/diagnostics/ax_model_manager.py
index ba37bd76..ebf89d1b 100644
--- a/optimas/diagnostics/ax_model_manager.py
+++ b/optimas/diagnostics/ax_model_manager.py
@@ -78,13 +78,13 @@ def build_model(
         parameters: list of `VaryingParameter`.
             Useful to initialize from `ExplorationDiagnostics`.
         """
-        # Define parameters
+        # Define parameters for AxClient
         if parameters is None:
             if parnames is None:
                 raise RuntimeError(
                     "Either `parameters` or `parnames` should be provided."
                 )
-            parameters = [
+            axparameters = [
                 {
                     "name": p_name,
                     "type": "range",
@@ -94,9 +94,33 @@ def build_model(
                 for p_name in parnames
             ]
         else:
-            parnames = [par["name"] for par in parameters]
+            parnames = [par.name for par in parameters]
+            axparameters = []
+            for par in 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
+        # Define objectives for AxClient
         if objectives is None:
             if objname is None:
                 raise RuntimeError(
@@ -111,15 +135,15 @@ def build_model(
             )
             self.ax_client.create_experiment(
                 name="optimas_data",
-                parameters=parameters,
+                parameters=axparameters,
                 objective_name=objname,
                 minimize=minimize,
             )
         else:
             # Create multi-objective Axclient
-            objectives_dict = {}
+            axobjectives = {}
             for obj in objectives:
-                objectives_dict[obj.name] = ObjectiveProperties(
+                axobjectives[obj.name] = ObjectiveProperties(
                     minimize=obj.minimize
                 )
             gs = GenerationStrategy(
@@ -130,8 +154,8 @@ def build_model(
             )
             self.ax_client.create_experiment(
                 name="optimas_data",
-                parameters=parameters,
-                objectives=objectives_dict,
+                parameters=axparameters,
+                objectives=axobjectives,
             )
 
         # Add trials from DataFrame
@@ -414,7 +438,7 @@ def plot_model(
             if clabel:
                 plt.clabel(cset, inline=True, fmt="%1.1f", fontsize="xx-small")
             # draw trials
-            ax.scatter(trials[xname], trials[yname], s=2, c="black", marker="o")
+            ax.scatter(trials[xname], trials[yname], s=8, c="black", marker="o")
             ax.set_xlim(xrange)
             ax.set_ylim(yrange)
             axs.append(ax)
diff --git a/optimas/diagnostics/exploration_diagnostics.py b/optimas/diagnostics/exploration_diagnostics.py
index 4c18bf74..d60658c0 100644
--- a/optimas/diagnostics/exploration_diagnostics.py
+++ b/optimas/diagnostics/exploration_diagnostics.py
@@ -1000,7 +1000,8 @@ def build_model(
         ----------
         objname: string, optional
             Name of the objective (or metric).
-            If not given, it takes the list of objectives.
+            If not given, it takes the list of objectives
+            in the diagnostics.
         minimize: bool, optional
             Whether to minimize or maximize the objective.
             It is only used when `objname` is given.
@@ -1010,47 +1011,21 @@ def build_model(
         -------
         An instance of AxModelManager
         """
-        # Initialize `AxModelManager` with history data
+        # Initialize `AxModelManager` with history dataframe.
         df = self.history.copy()
         self.model_manager = AxModelManager(df)
 
-        # Get parameters for AxClient.
-        parameters = []
-        for var in self.varying_parameters:
-            # Determine parameter type.
-            value_dtype = np.dtype(var.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.
-            parameters.append(
-                {
-                    "name": var.name,
-                    "type": "range",
-                    "bounds": [var.lower_bound, var.upper_bound],
-                    "is_fidelity": var.is_fidelity,
-                    "target_value": var.fidelity_target_value,
-                    "value_type": value_type,
-                }
-            )
-
-        # Get objectives
-        objective_names = [obj.name for obj in self.objectives]
+        # Select objective.
         if objname is not None:
+            objective_names = [obj.name for obj in self.objectives]
             if objname in objective_names:
                 minimize = self._get_objective(objname).minimize
             self.model_manager.build_model(
-                parameters=parameters, objname=objname, minimize=minimize
-            )
+                parameters=self.varying_parameters, objname=objname, minimize=minimize
+            )   
         else:
             self.model_manager.build_model(
-                parameters=parameters, objectives=self.objectives
+                parameters=self.varying_parameters, objectives=self.objectives
             )
 
         return self.model_manager
diff --git a/tests/test_ax_model_manager.py b/tests/test_ax_model_manager.py
index 60e7f950..7ed57384 100644
--- a/tests/test_ax_model_manager.py
+++ b/tests/test_ax_model_manager.py
@@ -89,9 +89,14 @@ def test_ax_model_manager():
     fig = plt.figure(figsize=(10, 4.8))
     gs = GridSpec(1, 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`
     ax1 = mm_axcl.plot_model(
-        mname="f", pcolormesh_kw={"cmap": "GnBu"}, subplot_spec=gs[0, 0]
+        mname="f", p0={'x2': x2_c},
+        pcolormesh_kw={"cmap": "GnBu"}, subplot_spec=gs[0, 0]
     )
 
     # Get and draw top 3 evaluations for `f`
@@ -101,26 +106,27 @@ def test_ax_model_manager():
 
     # plot model for `f2`
     ax2 = mm_axcl.plot_model(
-        mname="f2", pcolormesh_kw={"cmap": "OrRd"}, subplot_spec=gs[0, 1]
+        mname="f2", p0={'x2': x2_c},
+        pcolormesh_kw={"cmap": "OrRd"}, subplot_spec=gs[0, 1]
     )
 
-    # Get and draw top 3 evaluations for `f`
+    # Get and draw top 3 evaluations for `f2`
     top_f2 = diags.get_best_evaluations_index(top=3, objective="f2")
     df2_top = diags.history.loc[top_f2][varpar_names]
     ax2.scatter(df2_top["x0"], df2_top["x1"], c="blue", marker="x")
-    plt.savefig(os.path.join(exploration_dir_path, "models.png"))
+    plt.savefig(os.path.join(exploration_dir_path, "models_2d.png"))
 
     # Make example figure of the models in 1D with errors.
-    x1 = np.ones(100) * 0.5 * (var2.lower_bound + var2.upper_bound)
-    x2 = np.ones(100) * 0.5 * (var3.lower_bound + var3.upper_bound)
     x0 = np.linspace(var1.lower_bound, var1.upper_bound, 100)
+    x1 = np.ones_like(x0) * x1_c
+    x2 = np.ones_like(x0) * x2_c
     metric_names = mm_axcl.ax_client.objective_names
     fig, axs = plt.subplots(len(metric_names), 1, sharex=True)
     for i, (ax, metric_name) in enumerate(zip(axs, metric_names)):
         mean, sed = mm_axcl.evaluate_model(
             sample={"x0": x0, "x1": x1, "x2": x2}, metric_name=metric_name
         )
-        ax.plot(x0, mean, color=f"C{i}", label=f"x1 = {x1[0]}")
+        ax.plot(x0, mean, color=f"C{i}", label=f"x1 = {x1_c}, x2 = {x2_c}")
         ax.fill_between(
             x0, mean - sed, mean + sed, color="lightgray", alpha=0.5
         )

From 7122d7f1729b2776a993a65c57a16d2ccd9801ff Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Wed, 28 Feb 2024 15:38:56 +0000
Subject: [PATCH 24/54] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 optimas/diagnostics/exploration_diagnostics.py |  6 ++++--
 tests/test_ax_model_manager.py                 | 14 +++++++++-----
 2 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/optimas/diagnostics/exploration_diagnostics.py b/optimas/diagnostics/exploration_diagnostics.py
index d60658c0..bc6ef178 100644
--- a/optimas/diagnostics/exploration_diagnostics.py
+++ b/optimas/diagnostics/exploration_diagnostics.py
@@ -1021,8 +1021,10 @@ def build_model(
             if objname in objective_names:
                 minimize = self._get_objective(objname).minimize
             self.model_manager.build_model(
-                parameters=self.varying_parameters, objname=objname, minimize=minimize
-            )   
+                parameters=self.varying_parameters,
+                objname=objname,
+                minimize=minimize,
+            )
         else:
             self.model_manager.build_model(
                 parameters=self.varying_parameters, objectives=self.objectives
diff --git a/tests/test_ax_model_manager.py b/tests/test_ax_model_manager.py
index 7ed57384..6ce85083 100644
--- a/tests/test_ax_model_manager.py
+++ b/tests/test_ax_model_manager.py
@@ -92,11 +92,13 @@ def test_ax_model_manager():
     # 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`
     ax1 = mm_axcl.plot_model(
-        mname="f", p0={'x2': x2_c},
-        pcolormesh_kw={"cmap": "GnBu"}, subplot_spec=gs[0, 0]
+        mname="f",
+        p0={"x2": x2_c},
+        pcolormesh_kw={"cmap": "GnBu"},
+        subplot_spec=gs[0, 0],
     )
 
     # Get and draw top 3 evaluations for `f`
@@ -106,8 +108,10 @@ def test_ax_model_manager():
 
     # plot model for `f2`
     ax2 = mm_axcl.plot_model(
-        mname="f2", p0={'x2': x2_c},
-        pcolormesh_kw={"cmap": "OrRd"}, subplot_spec=gs[0, 1]
+        mname="f2",
+        p0={"x2": x2_c},
+        pcolormesh_kw={"cmap": "OrRd"},
+        subplot_spec=gs[0, 1],
     )
 
     # Get and draw top 3 evaluations for `f2`

From a7f24b2c123a7e3472b99aa5f11e34b9b50414de Mon Sep 17 00:00:00 2001
From: delaossa <alberto.martinez.de.la.ossa@desy.de>
Date: Thu, 29 Feb 2024 11:10:08 +0100
Subject: [PATCH 25/54] Add `get_mid_point` method.

---
 optimas/diagnostics/ax_model_manager.py | 24 +++++++++++++++++++++---
 1 file changed, 21 insertions(+), 3 deletions(-)

diff --git a/optimas/diagnostics/ax_model_manager.py b/optimas/diagnostics/ax_model_manager.py
index ebf89d1b..e0145187 100644
--- a/optimas/diagnostics/ax_model_manager.py
+++ b/optimas/diagnostics/ax_model_manager.py
@@ -232,7 +232,7 @@ def get_best_point(
         self,
         metric_name: Optional[str] = None,
         use_model_predictions: Optional[bool] = True,
-    ) -> Tuple[NDArray]:
+    ) -> Dict:
         """Get the best scoring point in the sample.
 
          Parameter:
@@ -285,6 +285,22 @@ def get_best_point(
 
         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 plot_model(
         self,
         xname: Optional[str] = None,
@@ -312,8 +328,10 @@ def plot_model(
         mname: string, optional.
             Name of the metric to plot.
             If not specified, it will take the first objective in ``self.ax_client``.
-        p0: dictionary
-            Particular values of parameters to be fixed for the evaluation over the sample.
+        p0: dictionary, optional.
+            A dictionary ``{name: val}`` for the fixed values of the other
+            parameters. If not provided, then the values of the best predicted
+            parametrization will be used.
         npoints: int, optional
             Number of points in each axis.
         mode: string, optional.

From 2142dce0bb87815429a8a7ff2e1cea226bac1a2c Mon Sep 17 00:00:00 2001
From: Angel Ferran Pousa <angel.ferran.pousa@desy.de>
Date: Wed, 6 Mar 2024 15:00:13 +0100
Subject: [PATCH 26/54] Simplify implementation

---
 optimas/diagnostics/ax_model_manager.py       | 225 ++++++++----------
 .../diagnostics/exploration_diagnostics.py    | 123 ++++------
 tests/test_ax_model_manager.py                |  10 +-
 3 files changed, 145 insertions(+), 213 deletions(-)

diff --git a/optimas/diagnostics/ax_model_manager.py b/optimas/diagnostics/ax_model_manager.py
index e0145187..cc77ff28 100644
--- a/optimas/diagnostics/ax_model_manager.py
+++ b/optimas/diagnostics/ax_model_manager.py
@@ -3,7 +3,7 @@
 from typing import Optional, Union, List, Tuple, Dict, Any, Literal
 import numpy as np
 from numpy.typing import NDArray
-from pandas import DataFrame
+import pandas as pd
 import matplotlib.pyplot as plt
 from matplotlib.gridspec import GridSpec, GridSpecFromSubplotSpec, SubplotSpec
 from matplotlib.axes import Axes
@@ -22,8 +22,8 @@
 from ax.service.utils.instantiation import ObjectiveProperties
 
 
-class AxModelManager(object):
-    """Utilities for building and exploring surrogate models using ``Ax``.
+class AxModelManager:
+    """Manager for building and exploring GP surrogate models using ``Ax``.
 
     Parameters
     ----------
@@ -32,147 +32,120 @@ class AxModelManager(object):
         If ``DataFrame``, the model has to be build using ``build_model``.
         If ``AxClient``, it uses the data in there to build a model.
         If ``str``, it should be the path to an ``AxClient`` json file.
+    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, DataFrame]) -> None:
+    def __init__(
+        self,
+        source: Union[AxClient, str, pd.DataFrame],
+        varying_parameters: Optional[List[VaryingParameter]] = None,
+        objectives: Optional[List[Objective]] = None,
+    ) -> None:
         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, DataFrame):
-            self.df = source
-            self.ax_client = None
+        elif isinstance(source, pd.DataFrame):
+            self.ax_client = self._build_ax_client_from_dataframe(
+                source, varying_parameters, objectives
+            )
         else:
-            raise RuntimeError("Wrong source.")
-
-        if self.ax_client:
-            self.ax_client.fit_model()
+            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:
+    def _model(self) -> TorchModelBridge:
         """Get the model from the AxClient instance."""
         return self.ax_client.generation_strategy.model
 
-    def build_model(
+    def _build_ax_client_from_dataframe(
         self,
-        parnames: Optional[List[str]] = None,
-        objname: Optional[str] = None,
-        minimize: Optional[bool] = True,
-        parameters: Optional[List[VaryingParameter]] = None,
-        objectives: Optional[List[Objective]] = None,
-    ) -> None:
+        df: pd.DataFrame,
+        varying_parameters: List[VaryingParameter],
+        objectives: List[Objective],
+    ) -> AxClient:
         """Initialize the AxClient and the model using the given data.
 
         Parameters
         ----------
-        parnames: list of string
-            List with the names of the parameters of the model.
-        objname: string, optional.
-            Name of the objective.
-        minimize: bool, optional.
-            Whether to minimize or maximize the objective.
-            Only relevant to establish the best point.
+        df : DataFrame
+            The source pandas ``DataFrame``.
         objectives: list of `Objective`.
-            This is the way to build multi-objective models.
-            Useful to initialize from `ExplorationDiagnostics`.
-        parameters: list of `VaryingParameter`.
-            Useful to initialize from `ExplorationDiagnostics`.
+            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
-        if parameters is None:
-            if parnames is None:
-                raise RuntimeError(
-                    "Either `parameters` or `parnames` should be provided."
+        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}."
                 )
-            axparameters = [
+            # Create parameter dict and append to list.
+            axparameters.append(
                 {
-                    "name": p_name,
+                    "name": par.name,
                     "type": "range",
-                    "bounds": [self.df[p_name].min(), self.df[p_name].max()],
-                    "value_type": "float",
+                    "bounds": [par.lower_bound, par.upper_bound],
+                    "is_fidelity": par.is_fidelity,
+                    "target_value": par.fidelity_target_value,
+                    "value_type": value_type,
                 }
-                for p_name in parnames
-            ]
-        else:
-            parnames = [par.name for par in parameters]
-            axparameters = []
-            for par in 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
-        if objectives is None:
-            if objname is None:
-                raise RuntimeError(
-                    "Either `objectives` or `objname` should be provided."
-                )
-            # Create single objective AxClient
-            gs = GenerationStrategy(
+        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`. Any strategy that uses a GP surrogate
+        # should work.
+        ax_client = AxClient(
+            generation_strategy=GenerationStrategy(
                 [GenerationStep(model=Models.GPEI, num_trials=-1)]
-            )
-            self.ax_client = AxClient(
-                generation_strategy=gs, verbose_logging=False
-            )
-            self.ax_client.create_experiment(
-                name="optimas_data",
-                parameters=axparameters,
-                objective_name=objname,
-                minimize=minimize,
-            )
-        else:
-            # Create multi-objective Axclient
-            axobjectives = {}
-            for obj in objectives:
-                axobjectives[obj.name] = ObjectiveProperties(
-                    minimize=obj.minimize
-                )
-            gs = GenerationStrategy(
-                [GenerationStep(model=Models.MOO, num_trials=-1)]
-            )
-            self.ax_client = AxClient(
-                generation_strategy=gs, verbose_logging=False
-            )
-            self.ax_client.create_experiment(
-                name="optimas_data",
-                parameters=axparameters,
-                objectives=axobjectives,
-            )
+            ),
+            verbose_logging=False,
+        )
+        ax_client.create_experiment(
+            parameters=axparameters, objectives=axobjectives
+        )
 
         # Add trials from DataFrame
-        for index, row in self.df.iterrows():
-            params = {p_name: row[p_name] for p_name in parnames}
-            _, trial_id = self.ax_client.attach_trial(params)
-            data = {}
-            for mname in list(self.ax_client.experiment.metrics.keys()):
-                data[mname] = (row[mname], np.nan)
-            self.ax_client.complete_trial(trial_id, raw_data=data)
-
-        # fit GP model
-        self.ax_client.fit_model()
+        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 evaluate_model(
         self,
-        sample: Union[DataFrame, Dict, NDArray] = None,
+        sample: Union[pd.DataFrame, Dict, NDArray] = None,
         metric_name: Optional[str] = None,
     ) -> Tuple[NDArray]:
         """Evaluate the model over the specified sample.
@@ -194,9 +167,6 @@ def evaluate_model(
         m_array, sem_array : Two numpy arrays containing the mean of the model
             and the standard error of the mean (sem), respectively.
         """
-        if self.model is None:
-            raise RuntimeError("Model not present. Run ``build_model`` first.")
-
         if metric_name is None:
             metric_name = self.ax_client.objective_names[0]
         else:
@@ -223,7 +193,7 @@ def evaluate_model(
                 parameters[name] = sample[name].iloc[i]
             obsf_list.append(ObservationFeatures(parameters=parameters))
 
-        mu, cov = self.model.predict(obsf_list)
+        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
@@ -273,15 +243,18 @@ def get_best_point(
                 best_obj_i = np.argmax(obj_vals)
             best_point = param_vals[best_obj_i]
         else:
-            if use_model_predictions is True:
-                best_arm, best_point_predictions = self.model.model_best_point()
-                best_point = best_arm.parameters
-            else:
+            # Somehow `use_model_predictions` does not seem to make any
+            # difference when calling `get_best_trial`. We should check that.
+            # 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`.
-                best_point = self.ax_client.get_best_parameters(
-                    use_model_predictions=use_model_predictions
-                )[0]
+            index, best_point, _ = self.ax_client.get_best_trial(
+                use_model_predictions=use_model_predictions
+            )
 
         return best_point
 
@@ -353,12 +326,6 @@ def plot_model(
             Either a single `~matplotlib.axes.Axes` object or a list of Axes
             objects if more than one subplot was created.
         """
-        if self.ax_client is None:
-            raise RuntimeError("AxClient not present. Run `build_model` first.")
-
-        if self.model is None:
-            raise RuntimeError("Model not present. Run `build_model` first.")
-
         # get experiment info
         experiment = self.ax_client.experiment
         parnames = list(experiment.parameters.keys())
diff --git a/optimas/diagnostics/exploration_diagnostics.py b/optimas/diagnostics/exploration_diagnostics.py
index bc6ef178..0873f277 100644
--- a/optimas/diagnostics/exploration_diagnostics.py
+++ b/optimas/diagnostics/exploration_diagnostics.py
@@ -19,8 +19,6 @@
 from optimas.explorations import Exploration
 from optimas.utils.other import get_df_with_selection
 from optimas.diagnostics.ax_model_manager import AxModelManager
-from ax.service.ax_client import AxClient
-from ax.service.utils.instantiation import ObjectiveProperties
 
 
 class ExplorationDiagnostics:
@@ -368,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,
@@ -900,12 +889,12 @@ def print_evaluation(self, trial_index: int) -> None:
 
         print()
 
-    def get_best_evaluations_index(
+    def get_best_evaluations(
         self,
         objective: Optional[Union[str, Objective]] = None,
         top: Optional[int] = 3,
-    ) -> List[int]:
-        """Get a list with the indices of the best evaluations.
+    ) -> pd.DataFrame:
+        """Get a list with the best evaluations.
 
         Parameters
         ----------
@@ -924,12 +913,9 @@ def get_best_evaluations_index(
             objective = self.objectives[0]
         elif isinstance(objective, str):
             objective = self._get_objective(objective)
-        top_indices = list(
-            self.history.sort_values(
+        return self.history.sort_values(
                 by=objective.name, ascending=objective.minimize
-            ).index
-        )[:top]
-        return top_indices
+            ).iloc[:top]
 
     def print_best_evaluations(
         self,
@@ -951,83 +937,64 @@ def print_best_evaluations(
             objective = self.objectives[0]
         elif isinstance(objective, str):
             objective = self._get_objective(objective)
-        top_indices = self.get_best_evaluations_index(
-            top=top, objective=objective
-        )
+        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 get_model_manager_from_ax_client(
-        self,
-        source: Union[AxClient, str],
+    def build_gp_model(
+        self, parameter: str, minimize: Optional[bool] = None
     ) -> AxModelManager:
-        """Initialize AxModelManager from an existing ``AxClient``.
+        """Build a GP model of the specified parameter.
 
         Parameters
         ----------
-        source: AxClient or str,
-            Source of data from where to obtain the model.
-            It can be an existing ``AxClient`` or the path to
-            a json file.
-
-        Returns
-        -------
-        An instance of AxModelManager.
-        """
-        self.model_manager = AxModelManager(source)
-        return self.model_manager
-
-    def build_model(
-        self,
-        objname: Optional[str] = None,
-        minimize: Optional[bool] = True,
-    ) -> AxModelManager:
-        """Initialize AxModelManager and builds a GP model.
-
-        Parameters
-        ----------
-        objname: string, optional
-            Name of the objective (or metric).
-            If not given, it takes the list of objectives
-            in the diagnostics.
+        parameter: str
+            Name of an objective or analyzed parameter for which the model
+            will be built.
         minimize: bool, optional
-            Whether to minimize or maximize the objective.
-            It is only used when `objname` is given.
-            Only relevant to establish the best point of the model.
+            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
         -------
-        An instance of AxModelManager
+        AxModelManager
         """
-        # Initialize `AxModelManager` with history dataframe.
-        df = self.history.copy()
-        self.model_manager = AxModelManager(df)
-
         # Select objective.
-        if objname is not None:
-            objective_names = [obj.name for obj in self.objectives]
-            if objname in objective_names:
-                minimize = self._get_objective(objname).minimize
-            self.model_manager.build_model(
-                parameters=self.varying_parameters,
-                objname=objname,
+        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,
-            )
-        else:
-            self.model_manager.build_model(
-                parameters=self.varying_parameters, objectives=self.objectives
+                dtype=analyzed_parameter.dtype,
             )
 
-        return self.model_manager
+        # Initialize `AxModelManager` with history dataframe.
+        return AxModelManager(
+            source=self.history,
+            varying_parameters=self.varying_parameters,
+            objectives=[objective],
+        )
diff --git a/tests/test_ax_model_manager.py b/tests/test_ax_model_manager.py
index 6ce85083..96fd025b 100644
--- a/tests/test_ax_model_manager.py
+++ b/tests/test_ax_model_manager.py
@@ -55,7 +55,7 @@ def test_ax_model_manager():
     # 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]
+    df = diags.history[varpar_names + ["f"]]
 
     # Get model manager directly from the existing `AxClient` instance.
     mm_axcl = AxModelManager(source=ax_client)
@@ -72,7 +72,7 @@ def test_ax_model_manager():
     mean_json, sem_json = mm_json.evaluate_model(sample=df, metric_name="f")
 
     # Get model manager from diagnostics data.
-    mm_diag = diags.build_model(objname="f")
+    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.
@@ -102,8 +102,7 @@ def test_ax_model_manager():
     )
 
     # Get and draw top 3 evaluations for `f`
-    top_f = diags.get_best_evaluations_index(top=3, objective="f")
-    df_top = diags.history.loc[top_f][varpar_names]
+    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 `f2`
@@ -115,8 +114,7 @@ def test_ax_model_manager():
     )
 
     # Get and draw top 3 evaluations for `f2`
-    top_f2 = diags.get_best_evaluations_index(top=3, objective="f2")
-    df2_top = diags.history.loc[top_f2][varpar_names]
+    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"))
 

From 02199f63f444ff1f060506ed6b332f23615d5dce Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Wed, 6 Mar 2024 14:01:24 +0000
Subject: [PATCH 27/54] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 optimas/diagnostics/ax_model_manager.py        | 4 ++--
 optimas/diagnostics/exploration_diagnostics.py | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/optimas/diagnostics/ax_model_manager.py b/optimas/diagnostics/ax_model_manager.py
index cc77ff28..7cbd6672 100644
--- a/optimas/diagnostics/ax_model_manager.py
+++ b/optimas/diagnostics/ax_model_manager.py
@@ -250,8 +250,8 @@ def get_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`.
+            # 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
             )
diff --git a/optimas/diagnostics/exploration_diagnostics.py b/optimas/diagnostics/exploration_diagnostics.py
index 0873f277..54a127da 100644
--- a/optimas/diagnostics/exploration_diagnostics.py
+++ b/optimas/diagnostics/exploration_diagnostics.py
@@ -914,8 +914,8 @@ def get_best_evaluations(
         elif isinstance(objective, str):
             objective = self._get_objective(objective)
         return self.history.sort_values(
-                by=objective.name, ascending=objective.minimize
-            ).iloc[:top]
+            by=objective.name, ascending=objective.minimize
+        ).iloc[:top]
 
     def print_best_evaluations(
         self,

From 5ef261e6eaa1abed00d72277028fc87933054675 Mon Sep 17 00:00:00 2001
From: delaossa <alberto.martinez.de.la.ossa@desy.de>
Date: Wed, 6 Mar 2024 15:56:45 +0100
Subject: [PATCH 28/54] Use `Models.MOO` for multi-objective cases.

---
 optimas/diagnostics/ax_model_manager.py | 20 +++++++++-----------
 1 file changed, 9 insertions(+), 11 deletions(-)

diff --git a/optimas/diagnostics/ax_model_manager.py b/optimas/diagnostics/ax_model_manager.py
index 7cbd6672..a0608216 100644
--- a/optimas/diagnostics/ax_model_manager.py
+++ b/optimas/diagnostics/ax_model_manager.py
@@ -127,7 +127,7 @@ def _build_ax_client_from_dataframe(
         # should work.
         ax_client = AxClient(
             generation_strategy=GenerationStrategy(
-                [GenerationStep(model=Models.GPEI, num_trials=-1)]
+                [GenerationStep(model=Models.GPEI if len(objectives) == 1 else Models.MOO, num_trials=-1)]
             ),
             verbose_logging=False,
         )
@@ -243,18 +243,16 @@ def get_best_point(
                 best_obj_i = np.argmax(obj_vals)
             best_point = param_vals[best_obj_i]
         else:
-            # Somehow `use_model_predictions` does not seem to make any
-            # difference when calling `get_best_trial`. We should check that.
-            # 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:
+            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
-            )
+                index, best_point, _ = self.ax_client.get_best_trial(
+                    use_model_predictions=use_model_predictions
+                )
 
         return best_point
 

From dc2d9aaec17dbd4c1f458a5e945b1864bb8d0980 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Wed, 6 Mar 2024 15:01:25 +0000
Subject: [PATCH 29/54] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 optimas/diagnostics/ax_model_manager.py | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/optimas/diagnostics/ax_model_manager.py b/optimas/diagnostics/ax_model_manager.py
index a0608216..32a77fec 100644
--- a/optimas/diagnostics/ax_model_manager.py
+++ b/optimas/diagnostics/ax_model_manager.py
@@ -127,7 +127,14 @@ def _build_ax_client_from_dataframe(
         # should work.
         ax_client = AxClient(
             generation_strategy=GenerationStrategy(
-                [GenerationStep(model=Models.GPEI if len(objectives) == 1 else Models.MOO, num_trials=-1)]
+                [
+                    GenerationStep(
+                        model=(
+                            Models.GPEI if len(objectives) == 1 else Models.MOO
+                        ),
+                        num_trials=-1,
+                    )
+                ]
             ),
             verbose_logging=False,
         )
@@ -248,8 +255,8 @@ def get_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`.
+                # 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
                 )

From 5b1880d5d7076383fa14df60e965f04d8c5ca3e2 Mon Sep 17 00:00:00 2001
From: Angel Ferran Pousa <angel.ferran.pousa@desy.de>
Date: Wed, 6 Mar 2024 16:22:39 +0100
Subject: [PATCH 30/54] Reduce number of lines

---
 optimas/diagnostics/ax_model_manager.py | 16 +++-------------
 1 file changed, 3 insertions(+), 13 deletions(-)

diff --git a/optimas/diagnostics/ax_model_manager.py b/optimas/diagnostics/ax_model_manager.py
index 32a77fec..f2064409 100644
--- a/optimas/diagnostics/ax_model_manager.py
+++ b/optimas/diagnostics/ax_model_manager.py
@@ -125,19 +125,9 @@ def _build_ax_client_from_dataframe(
         # a random sampling step will be set up by Ax, and this step does not
         # allow calling `model.predict`. Any strategy that uses a GP surrogate
         # should work.
-        ax_client = AxClient(
-            generation_strategy=GenerationStrategy(
-                [
-                    GenerationStep(
-                        model=(
-                            Models.GPEI if len(objectives) == 1 else Models.MOO
-                        ),
-                        num_trials=-1,
-                    )
-                ]
-            ),
-            verbose_logging=False,
-        )
+        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
         )

From 18f0c7f809f651dff8d00415fb76ec6681bcc686 Mon Sep 17 00:00:00 2001
From: Angel Ferran Pousa <angel.ferran.pousa@desy.de>
Date: Wed, 6 Mar 2024 16:28:24 +0100
Subject: [PATCH 31/54] Add comment

---
 optimas/diagnostics/ax_model_manager.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/optimas/diagnostics/ax_model_manager.py b/optimas/diagnostics/ax_model_manager.py
index f2064409..10f4fa66 100644
--- a/optimas/diagnostics/ax_model_manager.py
+++ b/optimas/diagnostics/ax_model_manager.py
@@ -123,8 +123,9 @@ def _build_ax_client_from_dataframe(
         # 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`. Any strategy that uses a GP surrogate
-        # should work.
+        # 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)

From 9a7a8300f4a168f92dff0bc7bd92d95cc1d6ef3a Mon Sep 17 00:00:00 2001
From: Angel Ferran Pousa <angel.ferran.pousa@desy.de>
Date: Wed, 6 Mar 2024 17:34:22 +0100
Subject: [PATCH 32/54] Return index in `get_best_point`

---
 optimas/diagnostics/ax_model_manager.py | 21 ++++++++++-----------
 1 file changed, 10 insertions(+), 11 deletions(-)

diff --git a/optimas/diagnostics/ax_model_manager.py b/optimas/diagnostics/ax_model_manager.py
index 10f4fa66..c174448e 100644
--- a/optimas/diagnostics/ax_model_manager.py
+++ b/optimas/diagnostics/ax_model_manager.py
@@ -200,7 +200,7 @@ def get_best_point(
         self,
         metric_name: Optional[str] = None,
         use_model_predictions: Optional[bool] = True,
-    ) -> Dict:
+    ) -> Tuple[int, Dict]:
         """Get the best scoring point in the sample.
 
          Parameter:
@@ -231,15 +231,14 @@ def get_best_point(
             pp = self.ax_client.get_pareto_optimal_parameters(
                 use_model_predictions=use_model_predictions
             )
-            obj_vals = [
-                objs[metric_name] for index, (vals, (objs, covs)) in pp.items()
-            ]
-            param_vals = [vals for index, (vals, (objs, covs)) in pp.items()]
-            if minimize:
-                best_obj_i = np.argmin(obj_vals)
-            else:
-                best_obj_i = np.argmax(obj_vals)
-            best_point = param_vals[best_obj_i]
+            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()
@@ -252,7 +251,7 @@ def get_best_point(
                     use_model_predictions=use_model_predictions
                 )
 
-        return best_point
+        return index, best_point
 
     def get_mid_point(
         self,

From 9b1358d157b0adc30d2127533c2320495231eaae Mon Sep 17 00:00:00 2001
From: delaossa <alberto.martinez.de.la.ossa@desy.de>
Date: Wed, 6 Mar 2024 17:36:43 +0100
Subject: [PATCH 33/54] Implement fixed point.

---
 optimas/diagnostics/ax_model_manager.py | 50 +++++++++++++++++++------
 1 file changed, 39 insertions(+), 11 deletions(-)

diff --git a/optimas/diagnostics/ax_model_manager.py b/optimas/diagnostics/ax_model_manager.py
index a0608216..2f3ed9a4 100644
--- a/optimas/diagnostics/ax_model_manager.py
+++ b/optimas/diagnostics/ax_model_manager.py
@@ -147,6 +147,7 @@ def evaluate_model(
         self,
         sample: Union[pd.DataFrame, Dict, NDArray] = None,
         metric_name: Optional[str] = None,
+        fixed_point: Optional[Dict] = None,
     ) -> Tuple[NDArray]:
         """Evaluate the model over the specified sample.
 
@@ -161,6 +162,9 @@ def evaluate_model(
         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_point: dict, optional.
+            A dictionary ``{name: val}`` with the values of the parameters
+            to be fixed in the evaluation.
 
         Returns
         -------
@@ -180,6 +184,10 @@ def evaluate_model(
         parnames = list(self.ax_client.experiment.parameters.keys())
 
         sample = convert_to_dataframe(sample)
+    
+        if fixed_point is not None:
+            for key, val in fixed_point.items():
+                sample[key] = val 
 
         # check if labels of the dataframe match the parnames
         for name in parnames:
@@ -198,7 +206,7 @@ def evaluate_model(
         sem_array = np.sqrt(cov[metric_name][metric_name])
         return m_array, sem_array
 
-    def get_best_point(
+    def get_best_evaluation(
         self,
         metric_name: Optional[str] = None,
         use_model_predictions: Optional[bool] = True,
@@ -255,10 +263,26 @@ def get_best_point(
                 )
 
         return best_point
+    
+    def get_best_point(self, metric_name: Optional[str] = None) -> Dict:
+        """Get the best scoring point in the sample.
 
-    def get_mid_point(
-        self,
-    ) -> Dict:
+         Parameter:
+        ----------
+        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
@@ -277,7 +301,7 @@ def plot_model(
         xname: Optional[str] = None,
         yname: Optional[str] = None,
         mname: Optional[str] = None,
-        p0: Optional[Dict] = None,
+        p0: Optional[Union[Dict, Literal['best', 'mid']]] = None,
         npoints: Optional[int] = 200,
         xrange: Optional[List[float]] = None,
         yrange: Optional[List[float]] = None,
@@ -364,18 +388,22 @@ def plot_model(
         X, Y = np.meshgrid(xaxis, yaxis)
         sample = {xname: X.flatten(), yname: Y.flatten()}
 
-        if p0 is None:
+        if (p0 is None) or (p0 == 'mid'):
+            # Get mid point
+            p0 = self.get_mid_point()
+        elif p0 == 'best':
             # get best point
-            p0 = self.get_best_point(
-                metric_name=mname, use_model_predictions=True
-            )
+            p0 = self.get_best_point(metric_name=mname)
 
+        fixed_point = {}
         for name, val in p0.items():
             if name not in [xname, yname]:
-                sample[name] = np.ones(npoints**2) * val
+                fixed_point[name] = p0[name]
 
         # evaluate the model
-        f_plt, sd_plt = self.evaluate_model(sample=sample, metric_name=mname)
+        f_plt, sd_plt = self.evaluate_model(sample=sample, 
+                                            metric_name=mname, 
+                                            fixed_point=p0)
 
         # select quantities to plot and set the labels
         f_plots = []

From c65177858957bf1a7f951ffe827fa2c51ec452f5 Mon Sep 17 00:00:00 2001
From: delaossa <alberto.martinez.de.la.ossa@desy.de>
Date: Wed, 6 Mar 2024 17:50:20 +0100
Subject: [PATCH 34/54] Fix typo.

---
 optimas/diagnostics/ax_model_manager.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/optimas/diagnostics/ax_model_manager.py b/optimas/diagnostics/ax_model_manager.py
index 14f43146..7f3e3821 100644
--- a/optimas/diagnostics/ax_model_manager.py
+++ b/optimas/diagnostics/ax_model_manager.py
@@ -400,7 +400,7 @@ def plot_model(
         # evaluate the model
         f_plt, sd_plt = self.evaluate_model(sample=sample, 
                                             metric_name=mname, 
-                                            fixed_point=p0)
+                                            fixed_point=fixed_point)
 
         # select quantities to plot and set the labels
         f_plots = []

From ce4219621cf6c91f2537ddbf18139bb799ee75b8 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Wed, 6 Mar 2024 16:50:39 +0000
Subject: [PATCH 35/54] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 optimas/diagnostics/ax_model_manager.py | 23 ++++++++++++-----------
 1 file changed, 12 insertions(+), 11 deletions(-)

diff --git a/optimas/diagnostics/ax_model_manager.py b/optimas/diagnostics/ax_model_manager.py
index 7f3e3821..e683dc63 100644
--- a/optimas/diagnostics/ax_model_manager.py
+++ b/optimas/diagnostics/ax_model_manager.py
@@ -182,10 +182,10 @@ def evaluate_model(
         parnames = list(self.ax_client.experiment.parameters.keys())
 
         sample = convert_to_dataframe(sample)
-    
+
         if fixed_point is not None:
             for key, val in fixed_point.items():
-                sample[key] = val 
+                sample[key] = val
 
         # check if labels of the dataframe match the parnames
         for name in parnames:
@@ -260,7 +260,7 @@ def get_best_evaluation(
                 )
 
         return index, best_point
-    
+
     def get_best_point(self, metric_name: Optional[str] = None) -> Dict:
         """Get the best scoring point in the sample.
 
@@ -275,8 +275,9 @@ def get_best_point(self, metric_name: Optional[str] = None) -> Dict:
         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)
+        _, best_point = self.get_best_evaluation(
+            metric_name=metric_name, use_model_predictions=True
+        )
         return best_point
 
     def get_mid_point(self) -> Dict:
@@ -298,7 +299,7 @@ def plot_model(
         xname: Optional[str] = None,
         yname: Optional[str] = None,
         mname: Optional[str] = None,
-        p0: Optional[Union[Dict, Literal['best', 'mid']]] = None,
+        p0: Optional[Union[Dict, Literal["best", "mid"]]] = None,
         npoints: Optional[int] = 200,
         xrange: Optional[List[float]] = None,
         yrange: Optional[List[float]] = None,
@@ -385,10 +386,10 @@ def plot_model(
         X, Y = np.meshgrid(xaxis, yaxis)
         sample = {xname: X.flatten(), yname: Y.flatten()}
 
-        if (p0 is None) or (p0 == 'mid'):
+        if (p0 is None) or (p0 == "mid"):
             # Get mid point
             p0 = self.get_mid_point()
-        elif p0 == 'best':
+        elif p0 == "best":
             # get best point
             p0 = self.get_best_point(metric_name=mname)
 
@@ -398,9 +399,9 @@ def plot_model(
                 fixed_point[name] = p0[name]
 
         # evaluate the model
-        f_plt, sd_plt = self.evaluate_model(sample=sample, 
-                                            metric_name=mname, 
-                                            fixed_point=fixed_point)
+        f_plt, sd_plt = self.evaluate_model(
+            sample=sample, metric_name=mname, fixed_point=fixed_point
+        )
 
         # select quantities to plot and set the labels
         f_plots = []

From 0a6a15db829c944d4eead55275323d9e8e96458a Mon Sep 17 00:00:00 2001
From: delaossa <alberto.martinez.de.la.ossa@desy.de>
Date: Wed, 6 Mar 2024 17:52:36 +0100
Subject: [PATCH 36/54] Small correction.

---
 optimas/diagnostics/ax_model_manager.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/optimas/diagnostics/ax_model_manager.py b/optimas/diagnostics/ax_model_manager.py
index 7f3e3821..abb2161c 100644
--- a/optimas/diagnostics/ax_model_manager.py
+++ b/optimas/diagnostics/ax_model_manager.py
@@ -156,7 +156,7 @@ def evaluate_model(
             If numpy array, it must contain the values of all the model parameres.
             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 ``p0``.
+            unless they are further specified using ``fixed_point``.
         metric_name: str, optional.
             Name of the metric to evaluate.
             If not specified, it will take the first first objective in ``self.ax_client``.

From 8520fb85a14814f1258522a6b1289e3a459482f4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=81ngel=20Ferran=20Pousa?= <angel.ferran.pousa@desy.de>
Date: Thu, 7 Mar 2024 12:19:04 +0100
Subject: [PATCH 37/54] Various fixes and changes: - Formatting. - Rename
 `plot_model` to `plot_contour`. - Make `_get_best_point` and `_get_mid_point`
 internal methods. - Rename arguments of `plot_contour` for consistency. -
 Check if Ax is installed (it should remain an optional dependency of
 optimas). - Improve docstrings. - Rename `fixed_point` to `fixed_parameters`,
 since it is not really a point (it would be slice actually). - Replace
 `RuntimeError` with more appropriate `ValueError`. - Return figure and axis
 in `plot_contour`. - Add `plot_slice` method.

---
 optimas/diagnostics/ax_model_manager.py | 461 ++++++++++++++++--------
 1 file changed, 306 insertions(+), 155 deletions(-)

diff --git a/optimas/diagnostics/ax_model_manager.py b/optimas/diagnostics/ax_model_manager.py
index 8b8995d4..8a60e420 100644
--- a/optimas/diagnostics/ax_model_manager.py
+++ b/optimas/diagnostics/ax_model_manager.py
@@ -1,42 +1,50 @@
 """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
-from optimas.core import VaryingParameter, Objective
-from optimas.utils.other import convert_to_dataframe
 
 # Ax utilities for model building
-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
+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:
-    """Manager for building and exploring GP surrogate models using ``Ax``.
+    """Class for building and exploring GP models using an ``AxClient``.
 
     Parameters
     ----------
-    source: AxClient, str or DataFrame
+    source : AxClient, str or DataFrame
         Source data for the model.
         If ``DataFrame``, the model has to be build using ``build_model``.
         If ``AxClient``, it uses the data in there to build a model.
         If ``str``, it should be the path to an ``AxClient`` json file.
-    objectives: list of `Objective`, optional
+    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
+    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
@@ -49,6 +57,11 @@ def __init__(
         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):
@@ -82,9 +95,9 @@ def _build_ax_client_from_dataframe(
         ----------
         df : DataFrame
             The source pandas ``DataFrame``.
-        objectives: list of `Objective`.
+        objectives : list of `Objective`.
             List of objectives for which a GP model should be built.
-        varying_parameters: list of `VaryingParameter`.
+        varying_parameters : list of `VaryingParameter`.
             List of parameters that were varied to scan the value of the
             objectives.
         """
@@ -141,28 +154,86 @@ def _build_ax_client_from_dataframe(
             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_point: Optional[Dict] = None,
+        fixed_parameters: Optional[Dict] = None,
     ) -> Tuple[NDArray]:
         """Evaluate the model over the specified sample.
 
-        Parameter:
+        Parameters
         ----------
-        sample: DataFrame, dict of arrays or numpy array,
+        sample : DataFrame, dict of arrays or numpy array,
             containing the data sample where to evaluate the model.
-            If numpy array, it must contain the values of all the model parameres.
+            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_point``.
-        metric_name: str, optional.
+            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_point: dict, optional.
-            A dictionary ``{name: val}`` with the values of the parameters
-            to be fixed in the evaluation.
+            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
         -------
@@ -174,17 +245,17 @@ def evaluate_model(
         else:
             metric_names = list(self.ax_client.experiment.metrics.keys())
             if metric_name not in metric_names:
-                raise RuntimeError(
-                    f"Metric name {metric_name} does not match any of the metrics. "
-                    f"Available metrics are: {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_point is not None:
-            for key, val in fixed_point.items():
+        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
@@ -211,19 +282,21 @@ def get_best_evaluation(
     ) -> Tuple[int, Dict]:
         """Get the best scoring point in the sample.
 
-         Parameter:
+        Parameters
         ----------
-        metric_name: str, optional.
+        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.
+            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
         -------
-        best_point : dict
-            A dictionary with the parameters of the best point.
+        int, dict
+            The index of the best evaluation and a dictionary with its
+            parameters.
         """
         # metric name
         if metric_name is None:
@@ -251,87 +324,64 @@ def get_best_evaluation(
             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)
+                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`.
+                # 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 get_best_point(self, metric_name: Optional[str] = None) -> Dict:
-        """Get the best scoring point in the sample.
-
-         Parameter:
-        ----------
-        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 plot_model(
+    def plot_contour(
         self,
-        xname: Optional[str] = None,
-        yname: Optional[str] = None,
-        mname: Optional[str] = None,
-        p0: Optional[Union[Dict, Literal["best", "mid"]]] = None,
-        npoints: Optional[int] = 200,
-        xrange: Optional[List[float]] = None,
-        yrange: Optional[List[float]] = None,
+        x_param: Optional[str] = None,
+        y_param: Optional[str] = None,
+        metric_name: Optional[str] = None,
+        slice_values: Optional[Union[Dict, Literal["best", "mid"]]] = "mid",
+        n_points: Optional[int] = 200,
+        x_range: Optional[List[float]] = None,
+        y_range: Optional[List[float]] = None,
         mode: Optional[Literal["mean", "sem", "both"]] = "mean",
-        clabel: Optional[bool] = False,
+        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,
-    ) -> Union[Axes, List[Axes]]:
-        """Plot model in the two selected variables, while others are fixed to the optimum.
+    ) -> Tuple[Figure, Union[Axes, List[Axes]]]:
+        """Plot a 2D slice of the surrogate model.
 
-        Parameter:
+        Parameters
         ----------
-        xname: string
-            Name of the variable to plot in x axis.
-        yname: string
-            Name of the variable to plot in y axis.
-        mname: string, optional.
+        x_param : str
+            Name of the parameter to plot in x axis.
+        y_param : 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``.
-        p0: dictionary, optional.
-            A dictionary ``{name: val}`` for the fixed values of the other
-            parameters. If not provided, then the values of the best predicted
-            parametrization will be used.
-        npoints: int, optional
+            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.
-        mode: string, optional.
-            ``mean`` plots the model mean, ``sem`` the standard error of the mean,
-            ``both`` plots both.
-        clabel: bool
+        x_range, y_range : 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"``.
+        show_contour_labels : bool
             when true labels are shown along the contour lines.
+        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
@@ -342,9 +392,9 @@ def plot_model(
 
         Returns
         -------
-        `~.axes.Axes` or array of Axes
-            Either a single `~matplotlib.axes.Axes` object or a list of Axes
-            objects if more than one subplot was created.
+        Figure, Axes or list of Axes
+            A matplotb figure and either a single `Axes` or a list of `Axes`
+            if `mode="both"`.
         """
         # get experiment info
         experiment = self.ax_client.experiment
@@ -357,50 +407,52 @@ def plot_model(
             )
 
         # select the input variables
-        if xname is None:
-            xname = parnames[0]
-        if yname is None:
-            yname = parnames[1]
+        if x_param is None:
+            x_param = parnames[0]
+        if y_param is None:
+            y_param = parnames[1]
 
         # metric name
-        if mname is None:
-            mname = self.ax_client.objective_names[0]
+        if metric_name is None:
+            metric_name = self.ax_client.objective_names[0]
 
         # set the plotting range
-        if xrange is None:
-            xrange = [None, None]
-        if yrange is None:
-            yrange = [None, None]
-        if xrange[0] is None:
-            xrange[0] = experiment.parameters[xname].lower
-        if xrange[1] is None:
-            xrange[1] = experiment.parameters[xname].upper
-        if yrange[0] is None:
-            yrange[0] = experiment.parameters[yname].lower
-        if yrange[1] is None:
-            yrange[1] = experiment.parameters[yname].upper
+        if x_range is None:
+            x_range = [None, None]
+        if y_range is None:
+            y_range = [None, None]
+        if x_range[0] is None:
+            x_range[0] = experiment.parameters[x_param].lower
+        if x_range[1] is None:
+            x_range[1] = experiment.parameters[x_param].upper
+        if y_range[0] is None:
+            y_range[0] = experiment.parameters[y_param].lower
+        if y_range[1] is None:
+            y_range[1] = experiment.parameters[y_param].upper
 
         # get grid sample of points where to evalutate the model
-        xaxis = np.linspace(xrange[0], xrange[1], npoints)
-        yaxis = np.linspace(yrange[0], yrange[1], npoints)
+        xaxis = np.linspace(x_range[0], x_range[1], n_points)
+        yaxis = np.linspace(y_range[0], y_range[1], n_points)
         X, Y = np.meshgrid(xaxis, yaxis)
-        sample = {xname: X.flatten(), yname: Y.flatten()}
+        sample = {x_param: X.flatten(), y_param: Y.flatten()}
 
-        if (p0 is None) or (p0 == "mid"):
+        if slice_values == "mid":
             # Get mid point
-            p0 = self.get_mid_point()
-        elif p0 == "best":
+            slice_values = self._get_mid_point()
+        elif slice_values == "best":
             # get best point
-            p0 = self.get_best_point(metric_name=mname)
+            slice_values = self._get_best_point(metric_name=metric_name)
 
-        fixed_point = {}
-        for name, val in p0.items():
-            if name not in [xname, yname]:
-                fixed_point[name] = p0[name]
+        fixed_parameters = {}
+        for name, val in slice_values.items():
+            if name not in [x_param, y_param]:
+                fixed_parameters[name] = slice_values[name]
 
         # evaluate the model
         f_plt, sd_plt = self.evaluate_model(
-            sample=sample, metric_name=mname, fixed_point=fixed_point
+            sample=sample,
+            metric_name=metric_name,
+            fixed_parameters=fixed_parameters,
         )
 
         # select quantities to plot and set the labels
@@ -408,10 +460,10 @@ def plot_model(
         labels = []
         if mode in ["mean", "both"]:
             f_plots.append(f_plt.reshape(X.shape))
-            labels.append(mname + ", mean")
+            labels.append(metric_name + ", mean")
         if mode in ["sem", "both"]:
             f_plots.append(sd_plt.reshape(X.shape))
-            labels.append(mname + ", sem")
+            labels.append(metric_name + ", sem")
 
         # create figure
         nplots = len(f_plots)
@@ -433,7 +485,7 @@ def plot_model(
             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=xname, ylabel=yname)
+            ax.set(xlabel=x_param, ylabel=y_param)
             # contour lines
             cset = ax.contour(
                 X,
@@ -444,35 +496,134 @@ def plot_model(
                 colors="black",
                 linestyles="solid",
             )
-            if clabel:
-                plt.clabel(cset, inline=True, fmt="%1.1f", fontsize="xx-small")
+            if show_contour_labels:
+                ax.clabel(cset, inline=True, fmt="%1.1f", fontsize="xx-small")
             # draw trials
-            ax.scatter(trials[xname], trials[yname], s=8, c="black", marker="o")
-            ax.set_xlim(xrange)
-            ax.set_ylim(yrange)
+            ax.scatter(
+                trials[x_param], trials[y_param], s=8, c="black", marker="o"
+            )
+            ax.set_xlim(x_range)
+            ax.set_ylim(y_range)
             axs.append(ax)
 
         if nplots == 1:
-            return axs[0]
+            return fig, axs[0]
         else:
-            return axs
+            return fig, axs
 
-    def get_arm_index(
+    def plot_slice(
         self,
-        arm_name: str,
-    ) -> int:
-        """Get the index of the arm by its name.
+        param: 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,
+        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
         ----------
-        arm_name: string.
-            Name of the arm. If not given, the best arm is selected.
+        param : str
+            Name of the parameter to plot in x 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 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.
+        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.
+        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
         -------
-        index: int
-            Trial index of the arm.
+        Figure, Axes
         """
-        df = self.ax_client.get_trials_data_frame()
-        index = df.loc[df["arm_name"] == arm_name, "trial_index"].iloc[0]
-        return index
+        # get experiment info
+        experiment = self.ax_client.experiment
+        parnames = list(experiment.parameters.keys())
+
+        # select the input variables
+        if param is None:
+            param = 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].lower
+        if range[1] is None:
+            range[1] = experiment.parameters[param].upper
+
+        # get sample of points where to evalutate the model
+        sample = {param: 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]:
+                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], mean, label=label, **plot_kw)
+        ax.fill_between(
+            sample[param], mean - sem, mean + sem, color="lightgray", alpha=0.5
+        )
+        ax.set_xlabel(param)
+        ax.set_ylabel(metric_name)
+        ax.legend(frameon=False)
+
+        return fig, ax

From 85a6716f4d2dd74ec782d27dce1e3416295fa93a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=81ngel=20Ferran=20Pousa?= <angel.ferran.pousa@desy.de>
Date: Thu, 7 Mar 2024 12:19:25 +0100
Subject: [PATCH 38/54] Update test

---
 tests/test_ax_model_manager.py | 76 ++++++++++++++++++++--------------
 1 file changed, 46 insertions(+), 30 deletions(-)

diff --git a/tests/test_ax_model_manager.py b/tests/test_ax_model_manager.py
index 96fd025b..396a8481 100644
--- a/tests/test_ax_model_manager.py
+++ b/tests/test_ax_model_manager.py
@@ -86,17 +86,17 @@ def test_ax_model_manager():
     assert np.allclose(mean_axcl, mean_diag, rtol=1e-2)
 
     # Make example figure with two models in 2D.
-    fig = plt.figure(figsize=(10, 4.8))
-    gs = GridSpec(1, 2, wspace=0.2, hspace=0.3)
+    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`
-    ax1 = mm_axcl.plot_model(
-        mname="f",
-        p0={"x2": x2_c},
+    # 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],
     )
@@ -105,38 +105,54 @@ def test_ax_model_manager():
     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 `f2`
-    ax2 = mm_axcl.plot_model(
-        mname="f2",
-        p0={"x2": x2_c},
-        pcolormesh_kw={"cmap": "OrRd"},
+    # 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"))
 
-    # Make example figure of the models in 1D with errors.
-    x0 = np.linspace(var1.lower_bound, var1.upper_bound, 100)
-    x1 = np.ones_like(x0) * x1_c
-    x2 = np.ones_like(x0) * x2_c
-    metric_names = mm_axcl.ax_client.objective_names
-    fig, axs = plt.subplots(len(metric_names), 1, sharex=True)
-    for i, (ax, metric_name) in enumerate(zip(axs, metric_names)):
-        mean, sed = mm_axcl.evaluate_model(
-            sample={"x0": x0, "x1": x1, "x2": x2}, metric_name=metric_name
-        )
-        ax.plot(x0, mean, color=f"C{i}", label=f"x1 = {x1_c}, x2 = {x2_c}")
-        ax.fill_between(
-            x0, mean - sed, mean + sed, color="lightgray", alpha=0.5
-        )
-        ax.set_ylabel(metric_name)
-        ax.legend(frameon=False)
-    plt.xlabel("x0")
-    plt.savefig(os.path.join(exploration_dir_path, "models_1d.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()

From 448a8f774c1d01dbb22531e45f290cdf02bc91de Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 7 Mar 2024 11:19:44 +0000
Subject: [PATCH 39/54] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 tests/test_ax_model_manager.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/test_ax_model_manager.py b/tests/test_ax_model_manager.py
index 396a8481..29d0d38f 100644
--- a/tests/test_ax_model_manager.py
+++ b/tests/test_ax_model_manager.py
@@ -154,5 +154,6 @@ def test_ax_model_manager():
     )
     fig.savefig(os.path.join(exploration_dir_path, "models_1d.png"))
 
+
 if __name__ == "__main__":
     test_ax_model_manager()

From 03096c410b567a7ff6e7f61224c77b49fc396978 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=81ngel=20Ferran=20Pousa?= <angel.ferran.pousa@desy.de>
Date: Thu, 7 Mar 2024 12:22:37 +0100
Subject: [PATCH 40/54] Add `AxModelManager` to documentation

---
 doc/source/api/diagnostics.rst | 1 +
 1 file changed, 1 insertion(+)

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

From eb1d8a2dd70e475d770421c6c2a8313f98147498 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=81ngel=20Ferran=20Pousa?= <angel.ferran.pousa@desy.de>
Date: Thu, 7 Mar 2024 12:32:24 +0100
Subject: [PATCH 41/54] Make argument names more similar to Ax

---
 optimas/diagnostics/ax_model_manager.py | 88 +++++++++++++------------
 1 file changed, 46 insertions(+), 42 deletions(-)

diff --git a/optimas/diagnostics/ax_model_manager.py b/optimas/diagnostics/ax_model_manager.py
index 8a60e420..2c64e82d 100644
--- a/optimas/diagnostics/ax_model_manager.py
+++ b/optimas/diagnostics/ax_model_manager.py
@@ -337,13 +337,13 @@ def get_best_evaluation(
 
     def plot_contour(
         self,
-        x_param: Optional[str] = None,
-        y_param: Optional[str] = None,
+        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,
-        x_range: Optional[List[float]] = None,
-        y_range: Optional[List[float]] = None,
+        range_x: Optional[List[float]] = None,
+        range_y: Optional[List[float]] = None,
         mode: Optional[Literal["mean", "sem", "both"]] = "mean",
         show_contour_labels: Optional[bool] = False,
         subplot_spec: Optional[SubplotSpec] = None,
@@ -355,9 +355,9 @@ def plot_contour(
 
         Parameters
         ----------
-        x_param : str
+        param_x : str
             Name of the parameter to plot in x axis.
-        y_param : str
+        param_y : str
             Name of the parameter to plot in y axis.
         metric_name : str, optional.
             Name of the metric to plot.
@@ -372,7 +372,7 @@ def plot_contour(
             parameter. By default, ``"mid"``.
         n_points : int, optional
             Number of points in each axis.
-        x_range, y_range : list of float, optional
+        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.
@@ -407,34 +407,34 @@ def plot_contour(
             )
 
         # select the input variables
-        if x_param is None:
-            x_param = parnames[0]
-        if y_param is None:
-            y_param = parnames[1]
+        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 x_range is None:
-            x_range = [None, None]
-        if y_range is None:
-            y_range = [None, None]
-        if x_range[0] is None:
-            x_range[0] = experiment.parameters[x_param].lower
-        if x_range[1] is None:
-            x_range[1] = experiment.parameters[x_param].upper
-        if y_range[0] is None:
-            y_range[0] = experiment.parameters[y_param].lower
-        if y_range[1] is None:
-            y_range[1] = experiment.parameters[y_param].upper
+        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(x_range[0], x_range[1], n_points)
-        yaxis = np.linspace(y_range[0], y_range[1], n_points)
+        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 = {x_param: X.flatten(), y_param: Y.flatten()}
+        sample = {param_x: X.flatten(), param_y: Y.flatten()}
 
         if slice_values == "mid":
             # Get mid point
@@ -445,7 +445,7 @@ def plot_contour(
 
         fixed_parameters = {}
         for name, val in slice_values.items():
-            if name not in [x_param, y_param]:
+            if name not in [param_x, param_y]:
                 fixed_parameters[name] = slice_values[name]
 
         # evaluate the model
@@ -485,7 +485,7 @@ def plot_contour(
             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=x_param, ylabel=y_param)
+            ax.set(xlabel=param_x, ylabel=param_y)
             # contour lines
             cset = ax.contour(
                 X,
@@ -500,10 +500,10 @@ def plot_contour(
                 ax.clabel(cset, inline=True, fmt="%1.1f", fontsize="xx-small")
             # draw trials
             ax.scatter(
-                trials[x_param], trials[y_param], s=8, c="black", marker="o"
+                trials[param_x], trials[param_y], s=8, c="black", marker="o"
             )
-            ax.set_xlim(x_range)
-            ax.set_ylim(y_range)
+            ax.set_xlim(range_x)
+            ax.set_ylim(range_y)
             axs.append(ax)
 
         if nplots == 1:
@@ -513,7 +513,7 @@ def plot_contour(
 
     def plot_slice(
         self,
-        param: Optional[str] = None,
+        range_name: Optional[str] = None,
         metric_name: Optional[str] = None,
         slice_values: Optional[Union[Dict, Literal["best", "mid"]]] = "mid",
         n_points: Optional[int] = 200,
@@ -527,7 +527,7 @@ def plot_slice(
 
         Parameters
         ----------
-        param : str
+        range_name : str
             Name of the parameter to plot in x axis.
         metric_name : str, optional.
             Name of the metric to plot.
@@ -564,8 +564,8 @@ def plot_slice(
         parnames = list(experiment.parameters.keys())
 
         # select the input variables
-        if param is None:
-            param = parnames[0]
+        if range_name is None:
+            range_name = parnames[0]
 
         # metric name
         if metric_name is None:
@@ -575,12 +575,12 @@ def plot_slice(
         if range is None:
             range = [None, None]
         if range[0] is None:
-            range[0] = experiment.parameters[param].lower
+            range[0] = experiment.parameters[range_name].lower
         if range[1] is None:
-            range[1] = experiment.parameters[param].upper
+            range[1] = experiment.parameters[range_name].upper
 
         # get sample of points where to evalutate the model
-        sample = {param: np.linspace(range[0], range[1], n_points)}
+        sample = {range_name: np.linspace(range[0], range[1], n_points)}
 
         if slice_values == "mid":
             # Get mid point
@@ -591,7 +591,7 @@ def plot_slice(
 
         fixed_parameters = {}
         for name, val in slice_values.items():
-            if name not in [param]:
+            if name not in [range_name]:
                 fixed_parameters[name] = slice_values[name]
 
         # evaluate the model
@@ -618,11 +618,15 @@ def plot_slice(
                 label += ", "
             label += f"{par} = {val}"
         ax = fig.add_subplot(gs[0])
-        ax.plot(sample[param], mean, label=label, **plot_kw)
+        ax.plot(sample[range_name], mean, label=label, **plot_kw)
         ax.fill_between(
-            sample[param], mean - sem, mean + sem, color="lightgray", alpha=0.5
+            x=sample[range_name],
+            y1=mean - sem,
+            y2=mean + sem,
+            color="lightgray",
+            alpha=0.5,
         )
-        ax.set_xlabel(param)
+        ax.set_xlabel(range_name)
         ax.set_ylabel(metric_name)
         ax.legend(frameon=False)
 

From dec20a7691331460b01baa7e6ab1bd614b0dff12 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=81ngel=20Ferran=20Pousa?= <angel.ferran.pousa@desy.de>
Date: Thu, 7 Mar 2024 12:33:13 +0100
Subject: [PATCH 42/54] Reformat imports

---
 tests/test_ax_model_manager.py | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/tests/test_ax_model_manager.py b/tests/test_ax_model_manager.py
index 29d0d38f..1b5fc6eb 100644
--- a/tests/test_ax_model_manager.py
+++ b/tests/test_ax_model_manager.py
@@ -2,13 +2,10 @@
 import numpy as np
 import matplotlib.pyplot as plt
 from matplotlib.gridspec import GridSpec
-import pytest
 
 from optimas.explorations import Exploration
 from optimas.core import VaryingParameter, Objective
-from optimas.generators import (
-    AxSingleFidelityGenerator,
-)
+from optimas.generators import AxSingleFidelityGenerator
 from optimas.evaluators import FunctionEvaluator
 from optimas.diagnostics import ExplorationDiagnostics, AxModelManager
 

From 85e32618a6294aaa90f7b2356c188fd9c3710706 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=81ngel=20Ferran=20Pousa?= <angel.ferran.pousa@desy.de>
Date: Thu, 7 Mar 2024 12:38:46 +0100
Subject: [PATCH 43/54] Improve docstrings

---
 optimas/diagnostics/ax_model_manager.py | 29 +++++++++++++------------
 1 file changed, 15 insertions(+), 14 deletions(-)

diff --git a/optimas/diagnostics/ax_model_manager.py b/optimas/diagnostics/ax_model_manager.py
index 2c64e82d..f415aaf8 100644
--- a/optimas/diagnostics/ax_model_manager.py
+++ b/optimas/diagnostics/ax_model_manager.py
@@ -218,7 +218,7 @@ def evaluate_model(
 
         Parameters
         ----------
-        sample : DataFrame, dict of arrays or numpy array,
+        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.
@@ -237,7 +237,8 @@ def evaluate_model(
 
         Returns
         -------
-        m_array, sem_array : Two numpy arrays containing the mean of the model
+        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:
@@ -381,20 +382,20 @@ def plot_contour(
         show_contour_labels : bool
             when true labels are shown along the contour lines.
         subplot_spec : SubplotSpec, optional
-            A matplotlib SubplotSpec in which to draw the axis.
+            A matplotlib ``SubplotSpec`` in which to draw the axis.
         gridspec_kw : dict, optional
-            Dict with keywords passed to the `GridSpec`.
+            Dict with keywords passed to the ``GridSpec``.
         pcolormesh_kw : dict, optional
-            Dict with keywords passed to `pcolormesh`.
+            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.
+            Additional keyword arguments to pass to ``pyplot.figure``.
+            Only used if no ``subplot_spec`` is given.
 
         Returns
         -------
         Figure, Axes or list of Axes
-            A matplotb figure and either a single `Axes` or a list of `Axes`
-            if `mode="both"`.
+            A matplotb figure and either a single ``Axes`` or a list of ``Axes``
+            if ``mode="both"``.
         """
         # get experiment info
         experiment = self.ax_client.experiment
@@ -546,14 +547,14 @@ def plot_slice(
             Range of the x axis. It not given, the lower and upper boundary
             of the x parameter will be used.
         subplot_spec : SubplotSpec, optional
-            A matplotlib SubplotSpec in which to draw the axis.
+            A matplotlib ``SubplotSpec`` in which to draw the axis.
         gridspec_kw : dict, optional
-            Dict with keywords passed to the `GridSpec`.
+            Dict with keywords passed to the ``GridSpec``.
         plot_kw : dict, optional
-            Dict with keywords passed to `ax.plot`.
+            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.
+            Additional keyword arguments to pass to ``pyplot.figure``. Only
+            used if no ``subplot_spec`` is given.
 
         Returns
         -------

From c7ee2356ca850d261ffac10b428102f56a54da27 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=81ngel=20Ferran=20Pousa?= <angel.ferran.pousa@desy.de>
Date: Thu, 7 Mar 2024 12:43:07 +0100
Subject: [PATCH 44/54] Update docstring and typo

---
 optimas/diagnostics/ax_model_manager.py | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/optimas/diagnostics/ax_model_manager.py b/optimas/diagnostics/ax_model_manager.py
index f415aaf8..5f3572df 100644
--- a/optimas/diagnostics/ax_model_manager.py
+++ b/optimas/diagnostics/ax_model_manager.py
@@ -36,10 +36,11 @@ class AxModelManager:
     Parameters
     ----------
     source : AxClient, str or DataFrame
-        Source data for the model.
-        If ``DataFrame``, the model has to be build using ``build_model``.
-        If ``AxClient``, it uses the data in there to build a model.
-        If ``str``, it should be the path to an ``AxClient`` json file.
+        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
@@ -394,8 +395,8 @@ def plot_contour(
         Returns
         -------
         Figure, Axes or list of Axes
-            A matplotb figure and either a single ``Axes`` or a list of ``Axes``
-            if ``mode="both"``.
+            A matplotlib figure and either a single ``Axes`` or a list of
+            ``Axes`` if ``mode="both"``.
         """
         # get experiment info
         experiment = self.ax_client.experiment

From 52e81cf4db81d2c12d792dceb9fe4077ddfa0a32 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=81ngel=20Ferran=20Pousa?= <angel.ferran.pousa@desy.de>
Date: Thu, 7 Mar 2024 16:20:39 +0100
Subject: [PATCH 45/54] Add documentation on building GPs

---
 .../advanced_usage/build_gp_surrogates.ipynb  | 266 ++++++++++++++++++
 doc/source/user_guide/index.rst               |   6 +
 2 files changed, 272 insertions(+)
 create mode 100644 doc/source/user_guide/advanced_usage/build_gp_surrogates.ipynb

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..cb252f74
--- /dev/null
+++ b/doc/source/user_guide/advanced_usage/build_gp_surrogates.ipynb
@@ -0,0 +1,266 @@
+{
+ "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 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 function,\n",
+    "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.AxSingleFidelityGenerator` and a simple\n",
+    ":class:`~optimas.evaluators.FunctionEvaluator` that evaluates a simple\n",
+    "analytical function."
+   ]
+  },
+  {
+   "cell_type": "raw",
+   "metadata": {},
+   "source": [
+    "Set up example optimization\n",
+    "~~~~~~~~~~~~~~~~~~~~~~~~~~~"
+   ]
+  },
+  {
+   "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 AxSingleFidelityGenerator\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[\"f\"] = 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",
+    "obj = Objective(\"f\", minimize=True)\n",
+    "obj2 = Objective(\"f2\", minimize=False)\n",
+    "\n",
+    "gen = AxSingleFidelityGenerator(\n",
+    "    varying_parameters=[var1, var2],\n",
+    "    objectives=[obj, 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=20,\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` object."
+   ]
+  },
+  {
+   "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 `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 better\n",
+    "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",
+    "f_model = diags.build_gp_model(\"f\")\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 `f`.\n",
+    "fig, ax1 = f_model.plot_contour(mode=\"both\", figsize=(6, 3), dpi=300)"
+   ]
+  },
+  {
+   "cell_type": "raw",
+   "metadata": {},
+   "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 `f`.\n",
+    "fig, ax1 = f_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=\"f\")\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, ax2 = 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 = f_model.evaluate_model(diags.history)\n",
+    "min_f, max_f = np.min(diags.history[\"f\"]), np.max(diags.history[\"f\"])\n",
+    "\n",
+    "# Make plot\n",
+    "fig, ax = plt.subplots(figsize=(5, 4), dpi=300)\n",
+    "ax.errorbar(diags.history[\"f\"], mean, yerr=sem, fmt=\"o\", 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

From a8ae726d01cdb176ed3f22df14720831d547d881 Mon Sep 17 00:00:00 2001
From: delaossa <alberto.martinez.de.la.ossa@desy.de>
Date: Thu, 7 Mar 2024 16:26:03 +0100
Subject: [PATCH 46/54] Add some options to the `AxModelManager` plots.

---
 optimas/diagnostics/ax_model_manager.py | 45 +++++++++++++++----------
 1 file changed, 28 insertions(+), 17 deletions(-)

diff --git a/optimas/diagnostics/ax_model_manager.py b/optimas/diagnostics/ax_model_manager.py
index 5f3572df..67fab4a6 100644
--- a/optimas/diagnostics/ax_model_manager.py
+++ b/optimas/diagnostics/ax_model_manager.py
@@ -347,6 +347,8 @@ def plot_contour(
         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,
@@ -380,6 +382,10 @@ def plot_contour(
         mode : str, optional.
             Whether to plot the ``"mean"`` of the model, the standard error of
             the mean ``"sem"``, or ``"both"``.
+        show_trials : bool
+            whether to show the trials used to build the model or not.
+        show_contour : bool
+            whether to show the contour or not.
         show_contour_labels : bool
             when true labels are shown along the contour lines.
         subplot_spec : SubplotSpec, optional
@@ -488,22 +494,23 @@ def plot_contour(
             cbar = plt.colorbar(im, ax=ax, location="top")
             cbar.set_label(labels[i])
             ax.set(xlabel=param_x, ylabel=param_y)
-            # contour lines
-            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")
-            # draw trials
-            ax.scatter(
-                trials[param_x], trials[param_y], s=8, c="black", marker="o"
-            )
+            # 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)
@@ -520,6 +527,7 @@ def plot_slice(
         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,
@@ -547,6 +555,8 @@ def plot_slice(
         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
+            when true a legend is shown with the slice values is shown. 
         subplot_spec : SubplotSpec, optional
             A matplotlib ``SubplotSpec`` in which to draw the axis.
         gridspec_kw : dict, optional
@@ -630,6 +640,7 @@ def plot_slice(
         )
         ax.set_xlabel(range_name)
         ax.set_ylabel(metric_name)
-        ax.legend(frameon=False)
+        if show_legend:
+            ax.legend(frameon=False)
 
         return fig, ax

From da112d854e8ae219349448c72af43280e180326b Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 7 Mar 2024 15:26:50 +0000
Subject: [PATCH 47/54] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 optimas/diagnostics/ax_model_manager.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/optimas/diagnostics/ax_model_manager.py b/optimas/diagnostics/ax_model_manager.py
index 67fab4a6..ec46d7a7 100644
--- a/optimas/diagnostics/ax_model_manager.py
+++ b/optimas/diagnostics/ax_model_manager.py
@@ -506,7 +506,9 @@ def plot_contour(
                     linestyles="solid",
                 )
                 if show_contour_labels:
-                    ax.clabel(cset, inline=True, fmt="%1.1f", fontsize="xx-small")
+                    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"
@@ -556,7 +558,7 @@ def plot_slice(
             Range of the x axis. It not given, the lower and upper boundary
             of the x parameter will be used.
         show_legend : bool
-            when true a legend is shown with the slice values is shown. 
+            when true a legend is shown with the slice values is shown.
         subplot_spec : SubplotSpec, optional
             A matplotlib ``SubplotSpec`` in which to draw the axis.
         gridspec_kw : dict, optional

From 555f0a21bb1aa7ae58a789f083954f28eb8ef2cb Mon Sep 17 00:00:00 2001
From: delaossa <alberto.martinez.de.la.ossa@desy.de>
Date: Thu, 7 Mar 2024 16:44:25 +0100
Subject: [PATCH 48/54] Typo in docs.

---
 optimas/diagnostics/ax_model_manager.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/optimas/diagnostics/ax_model_manager.py b/optimas/diagnostics/ax_model_manager.py
index ec46d7a7..75902e71 100644
--- a/optimas/diagnostics/ax_model_manager.py
+++ b/optimas/diagnostics/ax_model_manager.py
@@ -558,7 +558,7 @@ def plot_slice(
             Range of the x axis. It not given, the lower and upper boundary
             of the x parameter will be used.
         show_legend : bool
-            when true a legend is shown with the slice values is shown.
+            when true a legend is shown with the fixed slice values.
         subplot_spec : SubplotSpec, optional
             A matplotlib ``SubplotSpec`` in which to draw the axis.
         gridspec_kw : dict, optional

From f98b283c64462830fd611c5d94050c6bb7709959 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=81ngel=20Ferran=20Pousa?= <angel.ferran.pousa@desy.de>
Date: Thu, 7 Mar 2024 17:34:19 +0100
Subject: [PATCH 49/54] Update docs

---
 .../advanced_usage/build_gp_surrogates.ipynb  | 90 +++++++++++--------
 1 file changed, 53 insertions(+), 37 deletions(-)

diff --git a/doc/source/user_guide/advanced_usage/build_gp_surrogates.ipynb b/doc/source/user_guide/advanced_usage/build_gp_surrogates.ipynb
index cb252f74..094e90d5 100644
--- a/doc/source/user_guide/advanced_usage/build_gp_surrogates.ipynb
+++ b/doc/source/user_guide/advanced_usage/build_gp_surrogates.ipynb
@@ -9,27 +9,42 @@
     "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 model to any of the\n",
-    "objectives or analyzed parameters of an `optimas`\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 function,\n",
-    "make predictions, etc.\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.AxSingleFidelityGenerator` and a simple\n",
-    ":class:`~optimas.evaluators.FunctionEvaluator` that evaluates a simple\n",
-    "analytical function."
-   ]
-  },
-  {
-   "cell_type": "raw",
-   "metadata": {},
-   "source": [
+    ":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 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)"
    ]
   },
   {
@@ -41,7 +56,7 @@
     "import numpy as np\n",
     "from optimas.explorations import Exploration\n",
     "from optimas.core import VaryingParameter, Objective, Parameter\n",
-    "from optimas.generators import AxSingleFidelityGenerator\n",
+    "from optimas.generators import RandomSamplingGenerator\n",
     "from optimas.evaluators import FunctionEvaluator\n",
     "\n",
     "\n",
@@ -50,7 +65,7 @@
     "    x1 = input_params[\"x1\"]\n",
     "    x2 = input_params[\"x2\"]\n",
     "    result = -(x1 + 10 * np.cos(x1)) * (x2 + 5 * np.cos(x2))\n",
-    "    output_params[\"f\"] = result\n",
+    "    output_params[\"f1\"] = result\n",
     "    output_params[\"f2\"] = result * 2\n",
     "    output_params[\"p1\"] = np.sin(x1) + np.cos(x2)\n",
     "\n",
@@ -58,19 +73,19 @@
     "var1 = VaryingParameter(\"x1\", 0.0, 5.0)\n",
     "var2 = VaryingParameter(\"x2\", -5.0, 5.0)\n",
     "par1 = Parameter(\"p1\")\n",
-    "obj = Objective(\"f\", minimize=True)\n",
+    "obj1 = Objective(\"f1\", minimize=True)\n",
     "obj2 = Objective(\"f2\", minimize=False)\n",
     "\n",
-    "gen = AxSingleFidelityGenerator(\n",
+    "gen = RandomSamplingGenerator(\n",
     "    varying_parameters=[var1, var2],\n",
-    "    objectives=[obj, obj2],\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=20,\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",
@@ -90,7 +105,7 @@
     "~~~~~~~~~~~~~~~~~~~~~~\n",
     "\n",
     "The diagnostics class only requires the path to the exploration directory\n",
-    "as input parameter, or directly the `exploration` object."
+    "as input parameter, or directly the ``exploration`` instance."
    ]
   },
   {
@@ -113,15 +128,16 @@
     "Building a GP model of each objective and analyzed parameter\n",
     "~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n",
     "\n",
-    "To build a GP model, simply call `build_gp_model` on the diagnostics,\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",
+    "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 better\n",
-    "than higher values. This information is necessary, e.g., for determining\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."
    ]
   },
@@ -132,7 +148,7 @@
    "outputs": [],
    "source": [
     "# Build one model for each objective and analyzed parameter.\n",
-    "f_model = diags.build_gp_model(\"f\")\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)"
    ]
@@ -157,8 +173,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# plot model for `f`.\n",
-    "fig, ax1 = f_model.plot_contour(mode=\"both\", figsize=(6, 3), dpi=300)"
+    "# plot model for `f1`.\n",
+    "fig, ax1 = f1_model.plot_contour(mode=\"both\", figsize=(6, 3), dpi=300)"
    ]
   },
   {
@@ -166,7 +182,7 @@
    "metadata": {},
    "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",
+    "such as in the example below, by providing a ``subplot_spec`` where the plot\n",
     "should be drawn."
    ]
   },
@@ -182,14 +198,14 @@
     "fig = plt.figure(figsize=(10, 3), dpi=300)\n",
     "gs = GridSpec(1, 3, wspace=0.4)\n",
     "\n",
-    "# plot model for `f`.\n",
-    "fig, ax1 = f_model.plot_contour(\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=\"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",
@@ -199,7 +215,7 @@
     ")\n",
     "\n",
     "# plot model for `p1`\n",
-    "fig, ax2 = p1_model.plot_contour(\n",
+    "fig, ax3 = p1_model.plot_contour(\n",
     "    pcolormesh_kw={\"cmap\": \"PuBu\"},\n",
     "    subplot_spec=gs[0, 2],\n",
     ")"
@@ -229,12 +245,12 @@
    "outputs": [],
    "source": [
     "# Evaluate model for each point in the history\n",
-    "mean, sem = f_model.evaluate_model(diags.history)\n",
-    "min_f, max_f = np.min(diags.history[\"f\"]), np.max(diags.history[\"f\"])\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[\"f\"], mean, yerr=sem, fmt=\"o\", label=\"Data\")\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",

From 57793cd395ba5d89edd040866094ae2ef9dc088d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=81ngel=20Ferran=20Pousa?= <angel.ferran.pousa@desy.de>
Date: Thu, 7 Mar 2024 17:43:19 +0100
Subject: [PATCH 50/54] Edit docstrings

---
 optimas/diagnostics/ax_model_manager.py | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/optimas/diagnostics/ax_model_manager.py b/optimas/diagnostics/ax_model_manager.py
index 75902e71..27715153 100644
--- a/optimas/diagnostics/ax_model_manager.py
+++ b/optimas/diagnostics/ax_model_manager.py
@@ -375,19 +375,20 @@ def plot_contour(
             ``{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.
+            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"``.
+            the mean ``"sem"``, or ``"both"``. By default, ``"mean"``.
         show_trials : bool
-            whether to show the trials used to build the model or not.
+            Whether to show the trials used to build the model. By default,
+            ``True``.
         show_contour : bool
-            whether to show the contour or not.
+            Whether to show the contour lines. By default, ``True``.
         show_contour_labels : bool
-            when true labels are shown along the contour lines.
+            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
@@ -553,12 +554,13 @@ def plot_slice(
             ``{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.
+            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
-            when true a legend is shown with the fixed slice values.
+            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

From 28b794256b2d105ddc0528eced1325d152d95ec0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=81ngel=20Ferran=20Pousa?= <angel.ferran.pousa@desy.de>
Date: Thu, 7 Mar 2024 17:49:06 +0100
Subject: [PATCH 51/54] FIx typo

---
 optimas/diagnostics/ax_model_manager.py | 22 +++++++++++-----------
 1 file changed, 11 insertions(+), 11 deletions(-)

diff --git a/optimas/diagnostics/ax_model_manager.py b/optimas/diagnostics/ax_model_manager.py
index 27715153..b6cdbfef 100644
--- a/optimas/diagnostics/ax_model_manager.py
+++ b/optimas/diagnostics/ax_model_manager.py
@@ -525,7 +525,7 @@ def plot_contour(
 
     def plot_slice(
         self,
-        range_name: Optional[str] = None,
+        param_name: Optional[str] = None,
         metric_name: Optional[str] = None,
         slice_values: Optional[Union[Dict, Literal["best", "mid"]]] = "mid",
         n_points: Optional[int] = 200,
@@ -540,7 +540,7 @@ def plot_slice(
 
         Parameters
         ----------
-        range_name : str
+        param_name : str
             Name of the parameter to plot in x axis.
         metric_name : str, optional.
             Name of the metric to plot.
@@ -580,8 +580,8 @@ def plot_slice(
         parnames = list(experiment.parameters.keys())
 
         # select the input variables
-        if range_name is None:
-            range_name = parnames[0]
+        if param_name is None:
+            param_name = parnames[0]
 
         # metric name
         if metric_name is None:
@@ -591,12 +591,12 @@ def plot_slice(
         if range is None:
             range = [None, None]
         if range[0] is None:
-            range[0] = experiment.parameters[range_name].lower
+            range[0] = experiment.parameters[param_name].lower
         if range[1] is None:
-            range[1] = experiment.parameters[range_name].upper
+            range[1] = experiment.parameters[param_name].upper
 
         # get sample of points where to evalutate the model
-        sample = {range_name: np.linspace(range[0], range[1], n_points)}
+        sample = {param_name: np.linspace(range[0], range[1], n_points)}
 
         if slice_values == "mid":
             # Get mid point
@@ -607,7 +607,7 @@ def plot_slice(
 
         fixed_parameters = {}
         for name, val in slice_values.items():
-            if name not in [range_name]:
+            if name not in [param_name]:
                 fixed_parameters[name] = slice_values[name]
 
         # evaluate the model
@@ -634,15 +634,15 @@ def plot_slice(
                 label += ", "
             label += f"{par} = {val}"
         ax = fig.add_subplot(gs[0])
-        ax.plot(sample[range_name], mean, label=label, **plot_kw)
+        ax.plot(sample[param_name], mean, label=label, **plot_kw)
         ax.fill_between(
-            x=sample[range_name],
+            x=sample[param_name],
             y1=mean - sem,
             y2=mean + sem,
             color="lightgray",
             alpha=0.5,
         )
-        ax.set_xlabel(range_name)
+        ax.set_xlabel(param_name)
         ax.set_ylabel(metric_name)
         if show_legend:
             ax.legend(frameon=False)

From f944d4dd8277e587e060f98a496b7698a41857ab Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=81ngel=20Ferran=20Pousa?= <angel.ferran.pousa@desy.de>
Date: Thu, 7 Mar 2024 17:58:40 +0100
Subject: [PATCH 52/54] Add docstring details

---
 optimas/diagnostics/ax_model_manager.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/optimas/diagnostics/ax_model_manager.py b/optimas/diagnostics/ax_model_manager.py
index b6cdbfef..2ea38398 100644
--- a/optimas/diagnostics/ax_model_manager.py
+++ b/optimas/diagnostics/ax_model_manager.py
@@ -541,7 +541,8 @@ def plot_slice(
         Parameters
         ----------
         param_name : str
-            Name of the parameter to plot in x axis.
+            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

From 9ea67749fb8f35530eaadca54510de7a8baa3952 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=81ngel=20Ferran=20Pousa?= <angel.ferran.pousa@desy.de>
Date: Thu, 7 Mar 2024 17:58:53 +0100
Subject: [PATCH 53/54] Fix docs

---
 .../advanced_usage/build_gp_surrogates.ipynb  | 20 +++++++++++++++----
 1 file changed, 16 insertions(+), 4 deletions(-)

diff --git a/doc/source/user_guide/advanced_usage/build_gp_surrogates.ipynb b/doc/source/user_guide/advanced_usage/build_gp_surrogates.ipynb
index 094e90d5..8e3229e5 100644
--- a/doc/source/user_guide/advanced_usage/build_gp_surrogates.ipynb
+++ b/doc/source/user_guide/advanced_usage/build_gp_surrogates.ipynb
@@ -27,7 +27,7 @@
     "Set up example optimization\n",
     "~~~~~~~~~~~~~~~~~~~~~~~~~~~\n",
     "\n",
-    "The following cell sets up an optimization with two input parameters\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",
@@ -129,13 +129,13 @@
     "~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n",
     "\n",
     "To build a GP model, simply call\n",
-    "`:meth:`~optimas.diagnostics.Exploration.build_gp_model` on the diagnostics,\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",
+    "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."
@@ -178,8 +178,20 @@
    ]
   },
   {
-   "cell_type": "raw",
+   "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",

From 0b4d6a71367b69c3b25cebe4b80cf86720bf3be5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=81ngel=20Ferran=20Pousa?= <angel.ferran.pousa@desy.de>
Date: Thu, 7 Mar 2024 17:59:06 +0100
Subject: [PATCH 54/54] Increase version number

---
 optimas/__init__.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

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"