{
 "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\n",
    "except ImportError:\n",
    "    !python -m pip install openmdao[notebooks]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MetaModelStructuredComp\n",
    "\n",
    "`MetaModelStructuredComp` is a smooth interpolation Component for data that exists on a regular,\n",
    "structured, grid. This differs from [MetaModelUnStructured](metamodelunstructured_comp)\n",
    "which accepts unstructured data as collections of points.\n",
    "\n",
    "```{Note}\n",
    "OpenMDAO contains two components that perform interpolation: `SplineComp` and `MetaModelStructuredComp`.\n",
    "While they provide access to mostly the same algorithms, their usage is subtly different.\n",
    "The fundamental differences between them are as follows:\n",
    "\n",
    "`MetaModelStructuredComp` is used when you have a set of known data values y on a structured grid x and\n",
    "want to interpolate a new y value at a new x location that lies inside the grid. In this case, you\n",
    "generally start with a known set of fixed \"training\" values and their locations.\n",
    "\n",
    "[SplineComp](spline_comp.ipynb) is used when you want to create a smooth curve with a large number of points, but you\n",
    "want to control the shape of the curve with a small number of control points. The x locations of\n",
    "the interpolated points (and where applicable, the control points) are fixed and known, but the\n",
    "y values at the control points vary as the curve shape is modified by an upstream connection.\n",
    "\n",
    "MetaModelStructuredComp can be used for multi-dimensional design spaces, whereas SplineComp is\n",
    "restricted to one dimension.\n",
    "```\n",
    "\n",
    "\n",
    "`MetaModelStructuredComp` produces smooth fits through provided training data using polynomial\n",
    "splines of various orders. The interpolation methods include three that wrap methods in\n",
    "scipy.interpolate, as well as five methods that are written in pure python. For all methods,\n",
    "derivatives are automatically computed.  The following table summarizes the methods and gives\n",
    "the number of points required for each.\n",
    "\n",
    "```{eval-rst}\n",
    "+---------------+--------+------------------------------------------------------------------+\n",
    "| Method        | Order  | Description                                                      |\n",
    "+===============+========+==================================================================+\n",
    "| slinear       | 1      | Basic linear interpolation                                       |\n",
    "+---------------+--------+------------------------------------------------------------------+\n",
    "| lagrange2     | 2      | Second order Lagrange polynomial                                 |\n",
    "+---------------+--------+------------------------------------------------------------------+\n",
    "| lagrange3     | 3      | Third order Lagrange polynomial                                  |\n",
    "+---------------+--------+------------------------------------------------------------------+\n",
    "| akima         | 3      | Interpolation using Akima splines                                |\n",
    "+---------------+--------+------------------------------------------------------------------+\n",
    "| cubic         | 3      | Cubic spline, with continuity of derivatives between segments    |\n",
    "+---------------+--------+------------------------------------------------------------------+\n",
    "| scipy_slinear | 1      | Scipy linear interpolation. Same as slinear, though slower       |\n",
    "+---------------+--------+------------------------------------------------------------------+\n",
    "| scipy_cubic   | 3      | Scipy cubic interpolation. More accurate than cubic, but slower  |\n",
    "+---------------+--------+------------------------------------------------------------------+\n",
    "| scipy_quintic | 5      | Scipy quintic interpolation. Most accurate, but slowest          |\n",
    "+---------------+--------+------------------------------------------------------------------+\n",
    "| 3D-slinear    | 1      | Linear on a fixed 3D table                                       |\n",
    "+---------------+--------+------------------------------------------------------------------+\n",
    "| 1D-akima      | 3      | Akima on a fixed 1D table                                        |\n",
    "+---------------+--------+------------------------------------------------------------------+\n",
    "| 3D-lagrange2  | 2      | Third order Lagrange on a fixed 3D table                         |\n",
    "+---------------+--------+------------------------------------------------------------------+\n",
    "| 3D-lagrange3  | 3      | Third order Lagrange on a fixed 3D table                         |\n",
    "+---------------+--------+------------------------------------------------------------------+\n",
    "```\n",
    "\n",
    "Note that `MetaModelStructuredComp` only accepts scalar inputs and outputs. If you have a\n",
    "multivariable function, each input variable needs its own named OpenMDAO input.\n",
    "\n",
    "For multi-dimensional data, fits are computed on a separable per-axis basis. A single interpolation\n",
    "method is used for all dimensions, so the minimum table dimension must be high enough to use\n",
    "the chosen interpolate. However, if you choose one of the scipy methods, then automatic order\n",
    "reduction is supported. In this case, if a particular dimension does not have enough training data\n",
    "points to support a selected spline order (e.g. 3 sample points, but an order 5 'scipy_quintic'\n",
    "spline is specified), then the order of the fitted spline will be automatically reduced to one of the\n",
    "lower order scipy methods ('scipy_cubic' or 'scipy_slinear') for that dimension alone.\n",
    "\n",
    "The available methods include two (\"trilinear\" and \"akima1D\") that operate on a fixed number of dimensions in the table. These methods sacrifice flexibility in favor of computational efficiency. Some of the efficiency is gained by caching coefficients in a bin once they are computed, so these methods also assume that the table values are fixed.\n",
    "\n",
    "Extrapolation is supported, but disabled by default. It can be enabled via the `extrapolate`\n",
    "option (see below).\n",
    "\n",
    "## MetaModelStructuredComp Options\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "import openmdao.api as om\n",
    "om.show_options_table(\"openmdao.components.meta_model_structured_comp.MetaModelStructuredComp\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## MetaModelStructuredComp Constructor\n",
    "\n",
    "The call signature for the `MetaModelStructuredComp` constructor is:\n",
    "\n",
    "```{eval-rst}\n",
    "    .. automethod:: openmdao.components.meta_model_structured_comp.MetaModelStructuredComp.__init__\n",
    "        :noindex:\n",
    "```\n",
    "\n",
    "## MetaModelStructuredComp Examples\n",
    "\n",
    "A simple quick-start example is fitting the exclusive-or (\"XOR\") operator between\n",
    "two inputs, `x` and `y`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import openmdao.api as om\n",
    "\n",
    "# Create regular grid interpolator instance\n",
    "xor_interp = om.MetaModelStructuredComp(method='scipy_slinear')\n",
    "\n",
    "# set up inputs and outputs\n",
    "xor_interp.add_input('x', 0.0, training_data=np.array([0.0, 1.0]), units=None)\n",
    "xor_interp.add_input('y', 1.0, training_data=np.array([0.0, 1.0]), units=None)\n",
    "\n",
    "\n",
    "xor_interp.add_output('xor', 1.0, training_data=np.array([[0.0, 1.0], [1.0, 0.0]]), units=None)\n",
    "\n",
    "# Set up the OpenMDAO model\n",
    "model = om.Group()\n",
    "model.add_subsystem('comp', xor_interp, promotes=[\"*\"])\n",
    "prob = om.Problem(model)\n",
    "prob.setup()\n",
    "\n",
    "prob.set_val('x', 0)\n",
    "\n",
    "# Now test out a 'fuzzy' XOR\n",
    "prob.set_val('x', 0.9)\n",
    "prob.set_val('y', 0.001242)\n",
    "\n",
    "prob.run_model()\n",
    "\n",
    "print(prob.get_val('xor'))\n",
    "\n",
    "# we can verify all gradients by checking against finite-difference\n",
    "prob.check_partials(compact_print=True);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-output",
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "from numpy.testing import assert_almost_equal\n",
    "\n",
    "assert_almost_equal(prob.get_val('xor'), 0.8990064)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "An important consideration for multi-dimensional input is that the order in which\n",
    "the input variables are added sets the expected dimension of the output\n",
    "training data. For example, if inputs `x`, `y` and `z` are added to the component\n",
    "with training data array lengths of 5, 12, and 20 respectively, and are added\n",
    "in `x`, `y`, and `z` order, than the output training data must be an ndarray\n",
    "with shape (5, 12, 20).\n",
    "\n",
    "This is illustrated by the example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create input param training data, of sizes 25, 5, and 10 points resp.\n",
    "p1 = np.linspace(0, 100, 25)\n",
    "p2 = np.linspace(-10, 10, 5)\n",
    "p3 = np.linspace(0, 1, 10)\n",
    "\n",
    "# can use meshgrid to create a 3D array of test data\n",
    "P1, P2, P3 = np.meshgrid(p1, p2, p3, indexing='ij')\n",
    "f = np.sqrt(P1) + P2 * P3\n",
    "\n",
    "# verify the shape matches the order and size of the input params\n",
    "print(f.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create regular grid interpolator instance\n",
    "interp = om.MetaModelStructuredComp(method='scipy_cubic')\n",
    "interp.add_input('p1', 0.5, training_data=p1)\n",
    "interp.add_input('p2', 0.0, training_data=p2)\n",
    "interp.add_input('p3', 3.14, training_data=p3)\n",
    "\n",
    "interp.add_output('f', 0.0, training_data=f)\n",
    "\n",
    "# Set up the OpenMDAO model\n",
    "model = om.Group()\n",
    "model.add_subsystem('comp', interp, promotes=[\"*\"])\n",
    "prob = om.Problem(model)\n",
    "prob.setup()\n",
    "\n",
    "# set inputs\n",
    "prob.set_val('p1', 55.12)\n",
    "prob.set_val('p2', -2.14)\n",
    "prob.set_val('p3', 0.323)\n",
    "\n",
    "prob.run_model()\n",
    "\n",
    "print(prob.get_val('f'))\n",
    "\n",
    "# we can verify all gradients by checking against finite-difference\n",
    "prob.check_partials(compact_print=True);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "assert_almost_equal(prob.get_val('f'), 6.73306472)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can also predict multiple independent output points by setting the `vec_size` argument to be equal to the number of points you want to predict. Here, we set it to 2 and predict 2 points with `MetaModelStructuredComp`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create input param training data, of sizes 25, 5, and 10 points resp.\n",
    "p1 = np.linspace(0, 100, 25)\n",
    "p2 = np.linspace(-10, 10, 5)\n",
    "p3 = np.linspace(0, 1, 10)\n",
    "\n",
    "# can use meshgrid to create a 3D array of test data\n",
    "P1, P2, P3 = np.meshgrid(p1, p2, p3, indexing='ij')\n",
    "f = np.sqrt(P1) + P2 * P3\n",
    "\n",
    "# Create regular grid interpolator instance\n",
    "interp = om.MetaModelStructuredComp(method='scipy_cubic', vec_size=2)\n",
    "interp.add_input('p1', 0.5, training_data=p1)\n",
    "interp.add_input('p2', 0.0, training_data=p2)\n",
    "interp.add_input('p3', 3.14, training_data=p3)\n",
    "\n",
    "interp.add_output('f', 0.0, training_data=f)\n",
    "\n",
    "# Set up the OpenMDAO model\n",
    "model = om.Group()\n",
    "model.add_subsystem('comp', interp, promotes=[\"*\"])\n",
    "prob = om.Problem(model)\n",
    "prob.setup()\n",
    "\n",
    "# set inputs\n",
    "prob.set_val('p1', np.array([55.12, 12.0]))\n",
    "prob.set_val('p2', np.array([-2.14, 3.5]))\n",
    "prob.set_val('p3', np.array([0.323, 0.5]))\n",
    "\n",
    "prob.run_model()\n",
    "\n",
    "print(prob.get_val('f'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "assert_almost_equal(prob.get_val('f'), np.array([6.73306472, 5.2118645]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dynamic Training Points\n",
    "\n",
    "Finally, it is possible for MetaModelStructuredComp to be passed the values for each point on the grid from other components in a model. To enable this feature, set the option `training_data_gradients` to `True`. Every time `add_input` is called to add an input grid dimension \"f\", an additional coresponding input named \"f_train\" is added to the component. This input can be connected to an output somewhere in your model. The output should be a multi-dimensional array whose shape tuple contains the lengths of each dimension of the grid as declared when calling `add_input` on the MetaModelStructuredComp.\n",
    "\n",
    "Note that the grid point locations are always fixed and cannot be provided dynamically. Only the values at those grid points can be provided this way. \n",
    "\n",
    "When `training_data_gradients` is set to True, the partial derivative of the interpolated output with respect to the training inputs is also computed by the metamodel component. This may have a small performance impact for some of the interpolation methods.  Again, the grid point locations are fixed, so the derivative of the interpolated output with respect to those is not computed.\n",
    "\n",
    "The following example shows the use of dynamic training. Here, the component `TableGen` computes the values of the table each iteration. The formula from the previous example is slightly modified to add an input variable \"k\" that adjusts the table values.  This input could be declared as a design variable and adjusted by an optimizer.  At the end of the example, we run the model and do a `check_totals` to verify that we are computing correct derivatives across the training point inputs (\"comp.f\" with respect to \"tab.k\")."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "import openmdao.api as om\n",
    "\n",
    "\n",
    "# Define grids of sizes 25, 5, and 10 points for the dimensions of our training data.\n",
    "p1 = np.linspace(0, 100, 25)\n",
    "p2 = np.linspace(-10, 10, 5)\n",
    "p3 = np.linspace(0, 1, 10)\n",
    "\n",
    "\n",
    "class TableGen(om.ExplicitComponent):\n",
    "\n",
    "    def setup(self):\n",
    "        self.add_input('k', 1.0)\n",
    "        self.add_output('values', np.zeros((25, 5, 10)))\n",
    "\n",
    "        # These are fixed locations.\n",
    "        self.P1, self.P2, self.P3 = np.meshgrid(p1, p2, p3, indexing='ij')\n",
    "\n",
    "        self.declare_partials('values', 'k')\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        P1, P2, P3 = self.P1, self.P2, self.P3\n",
    "        k = inputs['k']\n",
    "\n",
    "        outputs['values'] = np.sqrt(P1) + P2 * P3 * k\n",
    "\n",
    "    def compute_partials(self, inputs, partials):\n",
    "        P2, P3 = self.P2, self.P3\n",
    "        partials['values', 'k'] = P2 * P3\n",
    "\n",
    "\n",
    "prob = om.Problem()\n",
    "model = prob.model\n",
    "\n",
    "model.add_subsystem('tab', TableGen())\n",
    "\n",
    "interp = om.MetaModelStructuredComp(method='lagrange3', training_data_gradients=True)\n",
    "interp.add_input('p1', 0.5, p1)\n",
    "interp.add_input('p2', 0.0, p2)\n",
    "interp.add_input('p3', 3.14, p3)\n",
    "\n",
    "# No need to pass training data into this call, because it is provided by 'tab'.\n",
    "interp.add_output('f', 0.0)\n",
    "\n",
    "model.add_subsystem('comp', interp)\n",
    "\n",
    "# The new input for the training data is called \"f_train\".\n",
    "model.connect('tab.values', 'comp.f_train')\n",
    "\n",
    "prob.setup(force_alloc_complex=True)\n",
    "\n",
    "# set inputs\n",
    "prob.set_val('comp.p1', 55.12)\n",
    "prob.set_val('comp.p2', -2.14)\n",
    "prob.set_val('comp.p3', 0.323)\n",
    "\n",
    "prob.run_model()\n",
    "\n",
    "print(prob.get_val('comp.f'))\n",
    "\n",
    "# we can verify all gradients by checking against finite-difference\n",
    "prob.check_totals(of='comp.f', wrt=['tab.k', 'comp.p1', 'comp.p2', 'comp.p3'],\n",
    "                  method='cs', compact_print=True);\n"
   ]
  },
  {
   "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('comp.f'), 6.73306, 1e-5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Standalone Interface for Table Interpolation\n",
    "\n",
    "The underlying interpolation algorithms can be used standalone (i.e., outside of the\n",
    "MetaModelStructuredComp) through the `InterpND` class. This can be useful for inclusion in another\n",
    "component.  The following component shows how to perform interpolation on the same table\n",
    "as in the previous example using standalone code. This time, we choose 'lagrange3' as the\n",
    "interpolation algorithm."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openmdao.components.interp_util.interp import InterpND\n",
    "\n",
    "# create input param training data, of sizes 25, 5, and 10 points resp.\n",
    "p1 = np.linspace(0, 100, 25)\n",
    "p2 = np.linspace(-10, 10, 5)\n",
    "p3 = np.linspace(0, 1, 10)\n",
    "\n",
    "# can use meshgrid to create a 3D array of test data\n",
    "P1, P2, P3 = np.meshgrid(p1, p2, p3, indexing='ij')\n",
    "f = np.sqrt(P1) + P2 * P3\n",
    "\n",
    "interp = InterpND(method='lagrange3', points=(p1, p2, p3), values=f)\n",
    "\n",
    "x = np.array([55.12, -2.14, 0.323])\n",
    "f, df_dx = interp.interpolate(x, compute_derivative=True)\n",
    "\n",
    "print(f)\n",
    "print(df_dx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-output",
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "actual = np.array([6.73306794])\n",
    "deriv_actual = np.array([[ 0.06734927, 0.323 , -2.14]])\n",
    "\n",
    "assert_near_equal(f, actual, tolerance=1e-7)\n",
    "assert_near_equal(df_dx, deriv_actual, tolerance=1e-7)"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.10"
  },
  "orphan": true
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
