{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "active-ipynb",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    from openmdao.utils.notebook_utils import notebook_mode  # noqa: F401\n",
    "except ImportError:\n",
    "    !python -m pip install openmdao[notebooks]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MetaModelUnStructuredComp\n",
    "\n",
    "`MetaModelUnStructuredComp` lets you quickly create a component with surrogate models\n",
    "used to compute the outputs based on unstructured training data. Generally, this is\n",
    "used to construct a low-computational-cost replacement for computationally\n",
    "expensive components.\n",
    "\n",
    "You can define `MetaModelUnStructuredComp` with as many inputs and outputs as you like,\n",
    "and you can also use a different surrogate model for each output.\n",
    "\n",
    "```{Note}\n",
    "What's the difference between `MetaModelUnStructuredComp` and a surrogate model? In\n",
    "OpenMDAO, \"surrogate model\" refers to the model for a single response, and\n",
    "`MetaModelUnStructuredComp` represents a collection of surrogate models trained at the\n",
    "same locations in the design space.\n",
    "```\n",
    "\n",
    "## MetaModelUnStructuredComp Options\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "import openmdao.api as om\n",
    "om.show_options_table(\"openmdao.components.meta_model_unstructured_comp.MetaModelUnStructuredComp\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## MetaModelUnStructuredComp Constructor\n",
    "\n",
    "The call signature for the `MetaModelUnStructuredComp` constructor is:\n",
    "\n",
    "```{eval-rst}\n",
    "    .. automethod:: openmdao.components.meta_model_unstructured_comp.MetaModelUnStructuredComp.__init__\n",
    "        :noindex:\n",
    "```\n",
    "\n",
    "## Simple Example\n",
    "\n",
    "The following example demonstrates a simple `Problem` in which a\n",
    "`MetaModelUnStructuredComp` uses surrogates to mimic the sine and cosine functions.\n",
    "\n",
    "In this example, the `MetaModelUnStructuredComp` `trig` has a single input,\n",
    "`x`, and two outputs, `sin_x` and `cos_x`.\n",
    "\n",
    "`KrigingSurrogate` is given as the surrogate for the `sin_x` output.\n",
    "Although no surrogate has been given for the `cos_x` output, a\n",
    "`default_surrogate` is specified for the component. Any output which has\n",
    "not had a surrogate assigned will use one of the default type.\n",
    "If `default_surrogate` is not specified, then a surrogate must be\n",
    "given for all outputs.\n",
    "\n",
    "\n",
    "The first time a `MetaModelUnStructuredComp` runs, it will train the surrogates using the\n",
    "training data that has been provided, and then it will predict the output\n",
    "values. This training step only occurs on the first run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a MetaModelUnStructuredComp, specifying surrogates for the outputs\n",
    "import numpy as np\n",
    "import openmdao.api as om\n",
    "\n",
    "trig = om.MetaModelUnStructuredComp()\n",
    "\n",
    "x_train = np.linspace(0,10,20)\n",
    "\n",
    "trig.add_input('x', 0., training_data=x_train)\n",
    "\n",
    "trig.add_output('sin_x', 0.,\n",
    "                training_data=.5*np.sin(x_train),\n",
    "                surrogate=om.KrigingSurrogate())\n",
    "trig.add_output('cos_x', 0.,\n",
    "                training_data=.5*np.cos(x_train))\n",
    "\n",
    "trig.options['default_surrogate'] = om.KrigingSurrogate()\n",
    "\n",
    "# add it to a Problem, run and check the predicted values\n",
    "prob = om.Problem()\n",
    "prob.model.add_subsystem('trig', trig)\n",
    "prob.setup()\n",
    "\n",
    "prob.set_val('trig.x', 2.1)\n",
    "prob.run_model()\n",
    "\n",
    "print(prob.get_val('trig.sin_x'))\n",
    "print(prob.get_val('trig.cos_x'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "from openmdao.utils.assert_utils import assert_near_equal\n",
    "\n",
    "assert_near_equal(prob.get_val('trig.sin_x'), .5*np.sin(prob.get_val('trig.x')), 1e-4)\n",
    "assert_near_equal(prob.get_val('trig.cos_x'), .5*np.cos(prob.get_val('trig.x')), 1e-4)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Available Surrogates\n",
    "\n",
    "The following surrogates are available to use with `MetaModelUnStructuredComp`.\n",
    "\n",
    "[KrigingSurrogate](../surrogates/kriging)\n",
    "  \n",
    "  **Derivatives**: Yes\n",
    "\n",
    "  **Description**: Based on Kriging interpolation; prediction returns mean predicted value, optionally returns RMSE.\n",
    "\n",
    "\n",
    "[NearestNeighbor](../surrogates/nearestneighbor)\n",
    "  \n",
    "  **Derivatives**: Yes\n",
    "\n",
    "  **Description**: Based on the N-Dimensional Interpolation library by Stephen Marone.\n",
    "\n",
    "\n",
    "[ResponseSurface](../surrogates/responsesurface)\n",
    "  \n",
    "  **Derivatives**: Yes\n",
    "\n",
    "  **Description**: Based on second order response surface equations.\n",
    "\n",
    "\n",
    "## Advanced usage\n",
    "\n",
    "You can specify the training data after instantiation if you like, by setting the component's\n",
    "[options](../../core_features/options/options). Training data is provided in the options to the `trig`\n",
    "component using the variable names prefixed with `train_`.  This can be done anytime before\n",
    "the `MetaModelUnStructuredComp` runs for the first time.\n",
    "\n",
    "The inputs and outputs of a `MetaModelUnStructuredComp` are not limited to scalar values. The\n",
    "following modified version of the example uses an array to predict sine and\n",
    "cosine as a single output array of two values.  You will also note that the default\n",
    "surrogate can be passed as an argument to the `MetaModelUnStructuredComp` constructor, as an\n",
    "alternative to specifying it later.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a MetaModelUnStructuredComp that predicts sine and cosine as an array\n",
    "trig = om.MetaModelUnStructuredComp(default_surrogate=om.KrigingSurrogate())\n",
    "trig.add_input('x', 0)\n",
    "trig.add_output('y', np.zeros(2))\n",
    "\n",
    "# add it to a Problem\n",
    "prob = om.Problem()\n",
    "prob.model.add_subsystem('trig', trig)\n",
    "prob.setup()\n",
    "\n",
    "# provide training data\n",
    "trig.options['train_x'] = np.linspace(0, 10, 20)\n",
    "trig.options['train_y'] = np.column_stack((\n",
    "    .5*np.sin(trig.options['train_x']),\n",
    "    .5*np.cos(trig.options['train_x'])\n",
    "))\n",
    "\n",
    "# train the surrogate and check predicted value\n",
    "prob.set_val('trig.x', 2.1)\n",
    "prob.run_model()\n",
    "print(prob.get_val('trig.y'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-output",
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "assert_near_equal(prob.get_val('trig.y'), np.append(.5*np.sin(prob.get_val('trig.x')), \n",
    "                                                    .5*np.cos(prob.get_val('trig.x'))),1e-4)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In addition, it's possible to vectorize the input and output variables so that you can\n",
    "make multiple predictions for the inputs and outputs in a single execution of the\n",
    "`MetaModelUnStructuredComp` component. This is done by setting the `vec_size` argument when\n",
    "constructing the `MetaModelUnStructuredComp` component and giving it the number of predictions to make.  The following example vectorizes the `trig`\n",
    "component so that it makes three predictions at a time.  In this case, the input is\n",
    "three independent values of `x` and the output is the corresponding predicted values\n",
    "for the sine and cosine functions at those three points.  Note that a vectorized\n",
    "`MetaModelUnStructuredComp` component requires the first dimension of all input and output variables\n",
    "to be the same size as specified in the `vec_size` argument."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "size = 3\n",
    "\n",
    "# create a vectorized MetaModelUnStructuredComp for sine and cosine\n",
    "trig = om.MetaModelUnStructuredComp(vec_size=size, default_surrogate=om.KrigingSurrogate())\n",
    "trig.add_input('x', np.zeros(size))\n",
    "trig.add_output('y', np.zeros((size, 2)))\n",
    "\n",
    "# add it to a Problem\n",
    "prob = om.Problem()\n",
    "prob.model.add_subsystem('trig', trig)\n",
    "prob.setup()\n",
    "\n",
    "# provide training data\n",
    "trig.options['train_x'] = np.linspace(0, 10, 20)\n",
    "trig.options['train_y'] = np.column_stack((\n",
    "    .5*np.sin(trig.options['train_x']),\n",
    "    .5*np.cos(trig.options['train_x'])\n",
    "))\n",
    "\n",
    "# train the surrogate and check predicted value\n",
    "prob.set_val('trig.x', np.array([2.1, 3.2, 4.3]))\n",
    "prob.run_model()\n",
    "\n",
    "print(prob.get_val('trig.y'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-output",
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "assert_near_equal(prob.get_val('trig.y'), np.column_stack((\n",
    "                     .5*np.sin(prob.get_val('trig.x')),\n",
    "                     .5*np.cos(prob.get_val('trig.x')))),1e-4)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using Surrogates That Do Not Define Linearize Method\n",
    "\n",
    "In some cases, you might define surrogates but not define a `linearize` method. In this case, the\n",
    "`MetaModelUnStructuredComp` derivatives will be computed using finite differences for the output variables that use that\n",
    "surrogate. By default, the default options for the finite differencing method will be used.\n",
    "\n",
    "If you would like to specify finite differencing options, you can do so by calling the `declare_partials`\n",
    "method in the component's `setup_partials` or in a parent group's configure method.\n",
    "This example, which uses a surrogate with no\n",
    "`linearize` method and no training for simplicity, shows `declare_partials` called in `setup_partials`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from numpy import sin\n",
    "\n",
    "class SinSurrogate(om.SurrogateModel):\n",
    "    def train(self, x, y):\n",
    "        pass\n",
    "\n",
    "    def predict(self, x):\n",
    "        return sin(x)\n",
    "\n",
    "class TrigWithFdInSetup(om.MetaModelUnStructuredComp):\n",
    "    def setup(self):\n",
    "        surrogate = SinSurrogate()\n",
    "        self.add_input('x', 0.)\n",
    "        self.add_output('sin_x', 0., surrogate=surrogate)\n",
    "\n",
    "    def setup_partials(self):\n",
    "        self.declare_partials('sin_x', 'x', method='fd',\n",
    "                              form='backward', step=1e-7, step_calc='rel')\n",
    "\n",
    "# Testing explicitly setting fd inside of setup\n",
    "prob = om.Problem()\n",
    "\n",
    "trig = TrigWithFdInSetup()\n",
    "prob.model.add_subsystem('trig', trig, promotes_inputs=['x'])\n",
    "\n",
    "prob.setup(check=True)\n",
    "prob.set_val('x', 5.)\n",
    "\n",
    "trig.train = False\n",
    "prob.run_model()\n",
    "\n",
    "J = prob.compute_totals(of=['trig.sin_x'], wrt=['x'])\n",
    "deriv_using_fd = J[('trig.sin_x', 'x')]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(deriv_using_fd[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "assert_near_equal(deriv_using_fd[0], np.cos(prob['x']), 1e-4)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Complex step has not been tested with `MetaModelUnStructuredComp` and will result in an exception if used."
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "kernelspec": {
   "display_name": "dev",
   "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.13.12"
  },
  "orphan": true
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
