{
 "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`` if 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.utils.AxModelManager.plot_contour`\n",
    "and :meth:`~optimas.utils.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.utils.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
}