Merge pull request #178 from delaossa/feature/ax_model_manager

Implement `AxModelManager`  and allow building GP models from diagnostics

doc/source/api/diagnostics.rst

@@ -7,3 +7,4 @@ Diagnostics
     :toctree: _autosummary
 
     ExplorationDiagnostics
+    AxModelManager

doc/source/user_guide/advanced_usage/build_gp_surrogates.ipynb

@@ -0,0 +1,294 @@
+{
+ "cells": [
+  {
+   "cell_type": "raw",
+   "metadata": {
+    "raw_mimetype": "text/restructuredtext"
+   },
+   "source": [
+    "Building GP surrogate models from optimization data\n",
+    "===================================================\n",
+    "\n",
+    "The :class:`~optimas.diagnostics.ExplorationDiagnostics` class\n",
+    "provides a simple way of fitting a Gaussian process (GP) model to any of the\n",
+    "objectives or analyzed parameters of an ``optimas``\n",
+    ":class:`~optimas.explorations.Exploration`, independently of which generator\n",
+    "was used. This is useful to get a better understanding of the underlying\n",
+    "function, make predictions, etc.\n",
+    "\n",
+    "In this example, we will illustrate how to build GP models by using\n",
+    "a basic optimization that runs directly on\n",
+    "a Jupyter notebook. This optimization uses an\n",
+    ":class:`~optimas.generators.RandomSamplingGenerator` and a simple\n",
+    ":class:`~optimas.evaluators.FunctionEvaluator` that evaluates an\n",
+    "analytical function.\n",
+    "\n",
+    "\n",
+    "Set up example optimization\n",
+    "~~~~~~~~~~~~~~~~~~~~~~~~~~~\n",
+    "\n",
+    "The following cell sets up and runs an optimization with two input parameters\n",
+    "``x1`` and ``x2``, two objectives ``f1`` and ``f2``, and one additional\n",
+    "analyzed parameter ``p1``.\n",
+    "At each evaluation, the ``eval_func_sf_moo`` function is run,\n",
+    "which assigns a value to each outcome parameter according to the analytical\n",
+    "formulas\n",
+    "\n",
+    ".. math::\n",
+    "\n",
+    "   f_1(x_1, x_2) = -(x_1 + 10 \\cos(x_1)) (x_2 + 5\\cos(x_2))\n",
+    "\n",
+    ".. math::\n",
+    "\n",
+    "   f_2(x_1, x_2) = 2 f_1(x_1, x_2)\n",
+    "\n",
+    ".. math::\n",
+    "\n",
+    "   p_1(x_1, x_2) = \\sin(x_1) + \\cos(x_2)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "from optimas.explorations import Exploration\n",
+    "from optimas.core import VaryingParameter, Objective, Parameter\n",
+    "from optimas.generators import RandomSamplingGenerator\n",
+    "from optimas.evaluators import FunctionEvaluator\n",
+    "\n",
+    "\n",
+    "def eval_func_sf_moo(input_params, output_params):\n",
+    "    \"\"\"Example multi-objective function.\"\"\"\n",
+    "    x1 = input_params[\"x1\"]\n",
+    "    x2 = input_params[\"x2\"]\n",
+    "    result = -(x1 + 10 * np.cos(x1)) * (x2 + 5 * np.cos(x2))\n",
+    "    output_params[\"f1\"] = result\n",
+    "    output_params[\"f2\"] = result * 2\n",
+    "    output_params[\"p1\"] = np.sin(x1) + np.cos(x2)\n",
+    "\n",
+    "\n",
+    "var1 = VaryingParameter(\"x1\", 0.0, 5.0)\n",
+    "var2 = VaryingParameter(\"x2\", -5.0, 5.0)\n",
+    "par1 = Parameter(\"p1\")\n",
+    "obj1 = Objective(\"f1\", minimize=True)\n",
+    "obj2 = Objective(\"f2\", minimize=False)\n",
+    "\n",
+    "gen = RandomSamplingGenerator(\n",
+    "    varying_parameters=[var1, var2],\n",
+    "    objectives=[obj1, obj2],\n",
+    "    analyzed_parameters=[par1],\n",
+    ")\n",
+    "ev = FunctionEvaluator(function=eval_func_sf_moo)\n",
+    "exploration = Exploration(\n",
+    "    generator=gen,\n",
+    "    evaluator=ev,\n",
+    "    max_evals=50,\n",
+    "    sim_workers=1,\n",
+    "    exploration_dir_path=\"./exploration\",\n",
+    "    libe_comms=\"threads\",  #this is only needed to run on a Jupyter notebook.\n",
+    ")\n",
+    "\n",
+    "# Run exploration.\n",
+    "exploration.run()"
+   ]
+  },
+  {
+   "cell_type": "raw",
+   "metadata": {
+    "raw_mimetype": "text/restructuredtext"
+   },
+   "source": [
+    "Initialize diagnostics\n",
+    "~~~~~~~~~~~~~~~~~~~~~~\n",
+    "\n",
+    "The diagnostics class only requires the path to the exploration directory\n",
+    "as input parameter, or directly the ``exploration`` instance."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from optimas.diagnostics import ExplorationDiagnostics\n",
+    "\n",
+    "diags = ExplorationDiagnostics(exploration)"
+   ]
+  },
+  {
+   "cell_type": "raw",
+   "metadata": {
+    "raw_mimetype": "text/restructuredtext"
+   },
+   "source": [
+    "Building a GP model of each objective and analyzed parameter\n",
+    "~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n",
+    "\n",
+    "To build a GP model, simply call\n",
+    ":meth:`~optimas.diagnostics.Exploration.build_gp_model` on the diagnostics,\n",
+    "indicating the name of the variable to which the model should be fitted.\n",
+    "This variable can be any ``objective`` or ``analyzed_parameter`` of the\n",
+    "optimization.\n",
+    "\n",
+    "Note that when building a surrogate model of an analyzed parameter, it is\n",
+    "required to provide a value to the ``minimize`` argument. This parameter\n",
+    "should therefore be ``True`` is lower values of the analyzed parameter are\n",
+    "better than higher values. This information is necessary, e.g., for determining\n",
+    "the best point in the model."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Build one model for each objective and analyzed parameter.\n",
+    "f1_model = diags.build_gp_model(\"f1\")\n",
+    "f2_model = diags.build_gp_model(\"f2\")\n",
+    "p1_model = diags.build_gp_model(\"p1\", minimize=False)"
+   ]
+  },
+  {
+   "cell_type": "raw",
+   "metadata": {
+    "raw_mimetype": "text/restructuredtext"
+   },
+   "source": [
+    "Visualizing the surrogate models\n",
+    "~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n",
+    "\n",
+    "The models provide some basic plotting methods for easy visualization, like\n",
+    ":meth:`~optimas.diagnostics.AxModelManager.plot_contour`\n",
+    "and :meth:`~optimas.diagnostics.AxModelManager.plot_slice`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# plot model for `f1`.\n",
+    "fig, ax1 = f1_model.plot_contour(mode=\"both\", figsize=(6, 3), dpi=300)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# plot 1D slice of `f1`.\n",
+    "fig, ax1 = f1_model.plot_slice(\"x1\", figsize=(6, 3), dpi=300)"
+   ]
+  },
+  {
+   "cell_type": "raw",
+   "metadata": {
+    "raw_mimetype": "text/restructuredtext"
+   },
+   "source": [
+    "These methods also allow more complex plot compositions to be created,\n",
+    "such as in the example below, by providing a ``subplot_spec`` where the plot\n",
+    "should be drawn."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "from matplotlib.gridspec import GridSpec\n",
+    "\n",
+    "fig = plt.figure(figsize=(10, 3), dpi=300)\n",
+    "gs = GridSpec(1, 3, wspace=0.4)\n",
+    "\n",
+    "# plot model for `f1`.\n",
+    "fig, ax1 = f1_model.plot_contour(\n",
+    "    pcolormesh_kw={\"cmap\": \"GnBu\"},\n",
+    "    subplot_spec=gs[0, 0],\n",
+    ")\n",
+    "\n",
+    "# Get and draw top 3 evaluations for `f`\n",
+    "df_top = diags.get_best_evaluations(top=3, objective=\"f1\")\n",
+    "ax1.scatter(df_top[\"x1\"], df_top[\"x2\"], c=\"red\", marker=\"x\")\n",
+    "\n",
+    "# plot model for `f2`\n",
+    "fig, ax2 = f2_model.plot_contour(\n",
+    "    pcolormesh_kw={\"cmap\": \"OrRd\"},\n",
+    "    subplot_spec=gs[0, 1],\n",
+    ")\n",
+    "\n",
+    "# plot model for `p1`\n",
+    "fig, ax3 = p1_model.plot_contour(\n",
+    "    pcolormesh_kw={\"cmap\": \"PuBu\"},\n",
+    "    subplot_spec=gs[0, 2],\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "raw",
+   "metadata": {
+    "raw_mimetype": "text/restructuredtext"
+   },
+   "source": [
+    "Evaluating the surrogate model\n",
+    "~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n",
+    "\n",
+    "In addition to plotting, it is also possible to evaluate the model at any\n",
+    "point by using the :meth:`~optimas.diagnostics.AxModelManager.evaluate_model`\n",
+    "method.\n",
+    "\n",
+    "In the example below, this method is used to evaluate the model in all the\n",
+    "history points to create a cross-validation plot."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Evaluate model for each point in the history\n",
+    "mean, sem = f1_model.evaluate_model(diags.history)\n",
+    "min_f, max_f = np.min(diags.history[\"f1\"]), np.max(diags.history[\"f1\"])\n",
+    "\n",
+    "# Make plot\n",
+    "fig, ax = plt.subplots(figsize=(5, 4), dpi=300)\n",
+    "ax.errorbar(diags.history[\"f1\"], mean, yerr=sem, fmt=\"o\", ms=4, label=\"Data\")\n",
+    "ax.plot([min_f, max_f], [min_f, max_f], color=\"k\", ls=\"--\", label=\"Ideal correlation\")\n",
+    "ax.set_xlabel(\"Observations\")\n",
+    "ax.set_ylabel(\"Model predictions\")\n",
+    "ax.legend(frameon=False)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "optimas_env_py11",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

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

optimas/__init__.py

@@ -1 +1 @@
-__version__ = "0.4.1"
+__version__ = "0.5.0"

optimas/diagnostics/__init__.py

@@ -1,3 +1,4 @@
 from .exploration_diagnostics import ExplorationDiagnostics
+from .ax_model_manager import AxModelManager
 
-__all__ = ["ExplorationDiagnostics"]
+__all__ = ["ExplorationDiagnostics", "AxModelManager"]