Merge pull request #264 from xopt-org/visualize_update

Visualize update

.github/workflows/gh-pages.yml

@@ -38,7 +38,7 @@ jobs:
           zip -r xopt-examples.zip docs/examples/
           mv xopt-examples.zip ./site/assets/
       - name: Upload docs artifact
-        uses: actions/upload-pages-artifact@v2
+        uses: actions/upload-pages-artifact@v3
         with:
           path: ./site
   deploy:

docs/examples/multi_objective_bayes_opt/mobo.ipynb

@@ -38,14 +38,11 @@
     "\n",
     "import pandas as pd\n",
     "import numpy as np\n",
-    "import torch\n",
     "\n",
     "from xopt import Xopt, Evaluator\n",
     "from xopt.generators.bayesian import MOBOGenerator\n",
     "from xopt.resources.test_functions.tnk import evaluate_TNK, tnk_vocs\n",
     "\n",
-    "from xopt.generators.bayesian.objectives import feasibility\n",
-    "\n",
     "import matplotlib.pyplot as plt\n",
     "\n",
     "# Ignore all warnings\n",
@@ -203,44 +200,8 @@
    },
    "outputs": [],
    "source": [
-    "# plot the acquisition function\n",
-    "bounds = X.generator.vocs.bounds\n",
-    "model = X.generator.model\n",
-    "\n",
-    "# create mesh\n",
-    "n = 200\n",
-    "x = torch.linspace(*bounds.T[0], n)\n",
-    "y = torch.linspace(*bounds.T[1], n)\n",
-    "xx, yy = torch.meshgrid(x, y)\n",
-    "pts = torch.hstack([ele.reshape(-1, 1) for ele in (xx, yy)]).double()\n",
-    "\n",
-    "xx, yy = xx.numpy(), yy.numpy()\n",
-    "\n",
-    "acq_func = X.generator.get_acquisition(model)\n",
-    "with torch.no_grad():\n",
-    "    acq_pts = pts.unsqueeze(1)\n",
-    "    acq = acq_func(acq_pts)\n",
-    "\n",
-    "    fig, ax = plt.subplots()\n",
-    "    c = ax.pcolor(xx, yy, acq.reshape(n, n), cmap=\"Blues\")\n",
-    "    fig.colorbar(c)\n",
-    "    ax.set_title(\"Acquisition function\")\n",
-    "\n",
-    "    ax.plot(*history[[\"x1\", \"x2\"]][history[\"feasible\"]].to_numpy().T, \".C1\")\n",
-    "    ax.plot(*history[[\"x1\", \"x2\"]][~history[\"feasible\"]].to_numpy().T, \".C2\")\n",
-    "\n",
-    "    ax.plot(*history[[\"x1\", \"x2\"]].to_numpy()[-1].T, \"+\")\n",
-    "\n",
-    "    feas = feasibility(pts.unsqueeze(1), model, tnk_vocs).flatten()\n",
-    "\n",
-    "    fig2, ax2 = plt.subplots()\n",
-    "    c = ax2.pcolor(xx, yy, feas.reshape(n, n))\n",
-    "    fig2.colorbar(c)\n",
-    "    ax2.set_title(\"Feasible Region\")\n",
-    "\n",
-    "candidate = pd.DataFrame(X.generator.generate(1), index=[0])\n",
-    "print(candidate[[\"x1\", \"x2\"]].to_numpy())\n",
-    "ax.plot(*candidate[[\"x1\", \"x2\"]].to_numpy()[0], \"o\")"
+    "## visualize model\n",
+    "X.generator.visualize_model()"
    ]
   }
  ],

docs/examples/multi_objective_bayes_opt/mobo_from_yaml.ipynb

@@ -34,9 +34,7 @@
    "outputs": [],
    "source": [
     "import os\n",
-    "import torch\n",
     "from xopt import Xopt\n",
-    "from xopt.generators.bayesian.objectives import feasibility\n",
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
     "import pandas as pd\n",
@@ -213,44 +211,8 @@
    },
    "outputs": [],
    "source": [
-    "# plot the acquisition function\n",
-    "bounds = X.generator.vocs.bounds\n",
-    "model = X.generator.model\n",
-    "\n",
-    "# create mesh\n",
-    "n = 100\n",
-    "x = torch.linspace(*bounds.T[0], n)\n",
-    "y = torch.linspace(*bounds.T[1], n)\n",
-    "xx, yy = torch.meshgrid(x, y)\n",
-    "pts = torch.hstack([ele.reshape(-1, 1) for ele in (xx, yy)]).double()\n",
-    "\n",
-    "xx, yy = xx.numpy(), yy.numpy()\n",
-    "\n",
-    "acq_func = X.generator.get_acquisition(model)\n",
-    "with torch.no_grad():\n",
-    "    acq_pts = pts.unsqueeze(1)\n",
-    "    acq = acq_func(acq_pts)\n",
-    "\n",
-    "    fig, ax = plt.subplots(figsize=(8, 8))\n",
-    "    c = ax.pcolor(xx, yy, acq.reshape(n, n), cmap=\"Blues\")\n",
-    "    fig.colorbar(c)\n",
-    "    ax.set_title(\"Acquisition function\")\n",
-    "\n",
-    "    ax.plot(*history[[\"x1\", \"x2\"]][history[\"feasible\"]].to_numpy().T, \".C1\")\n",
-    "    ax.plot(*history[[\"x1\", \"x2\"]][~history[\"feasible\"]].to_numpy().T, \".C2\")\n",
-    "\n",
-    "    ax.plot(*history[[\"x1\", \"x2\"]].to_numpy()[-1].T, \"+\")\n",
-    "\n",
-    "    feas = feasibility(pts.unsqueeze(1), model, X.vocs).flatten()\n",
-    "\n",
-    "    fig2, ax2 = plt.subplots(figsize=(8, 8))\n",
-    "    c = ax2.pcolor(xx, yy, feas.reshape(n, n))\n",
-    "    fig2.colorbar(c)\n",
-    "    ax2.set_title(\"Feasible Region\")\n",
-    "\n",
-    "candidate = pd.DataFrame(X.generator.generate(1), index=[0])\n",
-    "print(candidate[[\"x1\", \"x2\"]].to_numpy())\n",
-    "ax.plot(*candidate[[\"x1\", \"x2\"]].to_numpy()[0], \"o\")"
+    "## visualize model\n",
+    "X.generator.visualize_model()"
    ]
   }
  ],

xopt/generators/bayesian/bayesian_generator.py

@@ -535,7 +535,48 @@ def get_optimum(self):
         return self._process_candidates(result)
 
     def visualize_model(self, **kwargs):
-        """displays the GP models"""
+        """Display GP model predictions for the selected output(s).
+
+        The GP models are displayed with respect to the named variables. If None are given, the list of variables in
+        vocs is used. Feasible samples are indicated with a filled orange "o", infeasible samples with a hollow
+        red "o". Feasibility is calculated with respect to all constraints unless the selected output is a
+        constraint itself, in which case only that one is considered.
+
+        Parameters
+        ----------
+        **kwargs: dict, optional
+            Supported keyword arguments:
+            - output_names : List[str]
+                Outputs for which the GP models are displayed. Defaults to all outputs in vocs.
+            - variable_names : List[str]
+                The variables with respect to which the GP models are displayed (maximum of 2).
+                Defaults to vocs.variable_names.
+            - idx : int
+                Index of the last sample to use. This also selects the point of reference in
+                higher dimensions unless an explicit reference_point is given.
+            - reference_point : dict
+                Reference point determining the value of variables in vocs.variable_names, but not in variable_names
+                (slice plots in higher dimensions). Defaults to last used sample.
+            - show_samples : bool, optional
+                Whether samples are shown.
+            - show_prior_mean : bool, optional
+                Whether the prior mean is shown.
+            - show_feasibility : bool, optional
+                Whether the feasibility region is shown.
+            - show_acquisition : bool, optional
+                Whether the acquisition function is computed and shown (only if acquisition function is not None).
+            - n_grid : int, optional
+                Number of grid points per dimension used to display the model predictions.
+            - axes : Axes, optional
+                Axes object used for plotting.
+            - exponentiate : bool, optional
+                Flag to exponentiate acquisition function before plotting.
+
+        Returns
+        -------
+        result : tuple
+            The matplotlib figure and axes objects.
+        """
         return visualize_generator_model(self, **kwargs)
 
     def _get_initial_conditions(self, n_candidates=1) -> Union[Tensor, None]:

xopt/generators/bayesian/visualize.py

@@ -15,25 +15,53 @@
 
 
 def visualize_generator_model(generator, **kwargs) -> tuple:
-    """Displays GP model predictions for the selected output(s).
+    """Visualizes GP model predictions for the specified output(s).
 
-    The GP models are displayed with respect to the named variables. If None are given, the list of variables in
-    generator.vocs is used. Feasible samples are indicated with a filled orange "o", infeasible samples with a
-    hollow red "o". Feasibility is calculated with respect to all constraints unless the selected output is a
-    constraint itself, in which case only that one is considered.
+    This function generates a visualization of the Gaussian Process (GP) models associated with the provided generator.
+    It plots model predictions for selected outputs, showing feasible samples with filled orange circles ("o") and
+    infeasible samples with hollow red circles ("o"). Feasibility is calculated with respect to all constraints,
+    except when the selected output is a constraint itself, in which case only that constraint is considered.
 
     Parameters
     ----------
     generator : BayesianGenerator
-            Bayesian generator for which the GP model shall be visualized.
-    **kwargs : visualization parameters
-        See parameters of :func:`visualize_model`.
+        The Bayesian generator whose GP model is to be visualized. The generator must have a trained model.
+    **kwargs : dict, optional
+        Additional visualization parameters to customize the plots. Refer to the parameters of :func:`visualize_model`
+        for more details.
+
+    Raises
+    ------
+    ValueError
+        If the generator's model has not been trained (i.e., `generator.model` is None).
 
     Returns
     -------
     tuple
-        The matplotlib figure and axes objects.
+        A tuple containing the matplotlib figure and axes objects used for the visualization.
+
+    Examples
+    --------
+    >>> from xopt.generators.bayesian import ExpectedImprovementGenerator
+    >>> from xopt.generators.bayesian.visualize import visualize_generator_model
+    >>> generator = BayesianGenerator(...)
+    >>> generator.train_model()
+    >>> fig, ax = visualize_generator_model(generator, output_names=['output1'], show_acquisition=False)
+    >>> fig.show()
+
+    Notes
+    -----
+    - The visualization is limited to at most two variables at a time.
+    - The generator should be trained (via `generator.train_model()`) before calling this function.
+    - The function internally calls `visualize_model` to handle the actual plotting.
+
+    The following documentation is inherited from `visualize_model`:
+
     """
+
+    # append docstring dynamically
+    visualize_generator_model.__doc__ += visualize_model.__doc__
+
     if generator.model is None:
         raise ValueError(
             "The generator.model doesn't exist, try calling generator.train_model()."
@@ -62,6 +90,7 @@ def visualize_model(
     show_acquisition: bool = True,
     n_grid: int = 50,
     axes: Optional[Axes] = None,
+    exponentiate: bool = True,
 ) -> tuple:
     """Displays GP model predictions for the selected output(s).
 
@@ -103,6 +132,8 @@ def visualize_model(
         Number of grid points per dimension used to display the model predictions.
     axes : Axes, optional
         Axes object used for plotting.
+    exponentiate : bool, optional
+        Flag to exponentiate acquisition function before plotting.
 
     Returns
     -------
@@ -413,6 +444,7 @@ def plot_acquisition_function(
     show_legend: bool = True,
     n_grid: int = 100,
     axis=None,
+    exponentiate: bool = True,
     **_,
 ):
     """Displays the given acquisition function.
@@ -441,6 +473,8 @@ def plot_acquisition_function(
         See eponymous parameter of :func:`visualize_model`.
     axis : Axes, optional
         The axis to use for plotting. If None is given, a new one is generated.
+    exponentiate : bool, optional
+        Flag to exponentiate acquisition function before plotting.
     _
 
     Returns
@@ -453,6 +487,12 @@ def plot_acquisition_function(
     reference_point = _get_reference_point(reference_point, vocs, data, idx)
     kwargs = locals()
     input_mesh = _generate_input_mesh(**kwargs)
+
+    if exponentiate:
+        y_label = r"$\exp[ \alpha]$"
+    else:
+        y_label = r"$\alpha$"
+
     if len(variable_names) == 1:
         x_axis = (
             input_mesh[:, vocs.variable_names.index(variable_names[0])]
@@ -468,6 +508,10 @@ def plot_acquisition_function(
                 .numpy()
             )
         acq = acquisition_function(input_mesh.unsqueeze(1)).detach().squeeze().numpy()
+
+        if exponentiate:
+            acq = np.exp(acq)
+
         if base_acq is None:
             axis.plot(x_axis, acq, "C0-")
         else:
@@ -479,7 +523,8 @@ def plot_acquisition_function(
             if show_legend:
                 axis.legend()
         axis.set_xlabel(variable_names[0])
-        axis.set_ylabel(r"$\alpha\,$[{}]".format(vocs.output_names[0]))
+
+        axis.set_ylabel(y_label)
     else:
         if only_base_acq:
             if not hasattr(acquisition_function, "base_acquisition"):
@@ -496,6 +541,10 @@ def plot_acquisition_function(
             acq = (
                 acquisition_function(input_mesh.unsqueeze(1)).detach().squeeze().numpy()
             )
+
+        if exponentiate:
+            acq = np.exp(acq)
+
         if only_base_acq:
             title = "Base Acq. Function"
         elif hasattr(acquisition_function, "base_acquisition"):
@@ -506,7 +555,7 @@ def plot_acquisition_function(
             prediction=acq,
             input_mesh=input_mesh,
             title=title,
-            cbar_label=r"$\alpha\,$[{}]".format(vocs.output_names[0]),
+            cbar_label=y_label,
             output_name=vocs.output_names[0],
             **kwargs,
         )