{
 "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]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# KSComp\n",
    "\n",
    "KSComp provides a way to aggregate many constraints into a single constraint. This is usually done for performance\n",
    "reasons, in particular, to reduce the calculation time needed for the total derivatives of your model. The KSComp\n",
    "implements the Kreisselmeier-Steinhauser Function to aggregate constraint vector input \"g\" into a single scalar output 'KS'.\n",
    "\n",
    "By default, the constraint vector \"g\" is assumed be of the form where g<=0 satisfies the constraints, but other forms can be specified using the \"upper\" and \"lower_flag\" options.\n",
    "\n",
    "The output \"KS\" should be constrained with an upper-bound of zero to satisfy the aggregated constraint.\n",
    "By default, it is left to the user to provide this constraint.  However, setting option \"add_constraint\"\n",
    "to True will cause the KSComp to automatically add a constraint to the optimization.\n",
    "\n",
    "## KSComp Options"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "import openmdao.api as om\n",
    "om.show_options_table(\"openmdao.components.ks_comp.KSComp\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## KSComp Constructor\n",
    "\n",
    "The call signature for the `KSComp` constructor is:\n",
    "\n",
    "```{eval-rst}\n",
    "    .. automethod:: openmdao.components.ks_comp.KSComp.__init__\n",
    "        :noindex:\n",
    "```\n",
    "\n",
    "## KSComp Example\n",
    "\n",
    "The following example is perhaps the simplest possible. It shows a component that represents a constraint\n",
    "of width two. We would like to aggregate the values of this constraint vector into a single scalar\n",
    "value using the KSComp."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import openmdao.api as om\n",
    "\n",
    "prob = om.Problem()\n",
    "model = prob.model\n",
    "\n",
    "model.add_subsystem('comp', om.ExecComp('y = 3.0*x',\n",
    "                                        x=np.zeros((2, )),\n",
    "                                        y=np.zeros((2, ))), promotes_inputs=['x'])\n",
    "\n",
    "model.add_subsystem('ks', om.KSComp(width=2))\n",
    "\n",
    "model.connect('comp.y', 'ks.g')\n",
    "\n",
    "prob.setup()\n",
    "prob.set_val('x', np.array([5.0, 4.0]))\n",
    "prob.run_model()\n",
    "\n",
    "print(prob.get_val('ks.KS'))"
   ]
  },
  {
   "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('ks.KS'), [[15.0]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A more practical example that uses the KSComp can be found in the [beam optimization](../../../examples/beam_optimization_example_part_2) example.\n",
    "\n",
    "You can also independently aggregate multiple rows of an output as separate constraints by declaring the vec_size argument:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "import openmdao.api as om\n",
    "\n",
    "prob = om.Problem()\n",
    "model = prob.model\n",
    "\n",
    "model.add_subsystem('comp', om.ExecComp('y = 3.0*x',\n",
    "                                        x=np.zeros((2, 2)),\n",
    "                                        y=np.zeros((2, 2))), promotes_inputs=['x'])\n",
    "model.add_subsystem('ks', om.KSComp(width=2, vec_size=2))\n",
    "\n",
    "model.connect('comp.y', 'ks.g')\n",
    "\n",
    "prob.setup()\n",
    "prob.set_val('x', np.array([[5.0, 4.0], [10.0, 8.0]]))\n",
    "prob.run_model()\n",
    "\n",
    "print(prob.get_val('ks.KS'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "assert_near_equal(prob.get_val('ks.KS'), np.array([[15], [30]]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## KSComp Option Examples\n",
    "\n",
    "Normally, the input constraint vector is assumed to be of the form g<=0 is satisfied. If you would like to set a\n",
    "different upper bound for the constraint, you can declare it in the \"upper\" option in the options dictionary.\n",
    "\n",
    "In the following example, we specify a new upper bound of 16 for the constraint vector. Note that the KS output\n",
    "is still satisfied if it is less than zero.\n",
    "\n",
    "**upper**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "prob = om.Problem()\n",
    "model = prob.model\n",
    "\n",
    "model.add_subsystem('comp', om.ExecComp('y = 3.0*x',\n",
    "                                        x=np.zeros((2, )),\n",
    "                                        y=np.zeros((2, ))), promotes_inputs=['x'])\n",
    "model.add_subsystem('ks', om.KSComp(width=2))\n",
    "\n",
    "model.connect('comp.y', 'ks.g')\n",
    "\n",
    "model.ks.options['upper'] = 16.0\n",
    "prob.setup()\n",
    "prob.set_val('x', np.array([5.0, 4.0]))\n",
    "prob.run_model()\n",
    "\n",
    "print(prob['ks.KS'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "assert_near_equal(prob['ks.KS'], np.array([[-1.0]]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Normally, the input constraint vector is satisfied if it is negative and violated if it is positive. You can\n",
    "reverse this behavior by setting the \"lower_flag\" option to True. In the following example, we turn on the\n",
    "\"lower_flag\" so that positive values of the input constraint are considered satisfied. Note that the KS output\n",
    "is still satisfied if it is less than zero.\n",
    "\n",
    "**lower_flag**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "prob = om.Problem()\n",
    "model = prob.model\n",
    "\n",
    "model.add_subsystem('comp', om.ExecComp('y = 3.0*x',\n",
    "                                        x=np.zeros((2, )),\n",
    "                                        y=np.zeros((2, ))), promotes_inputs=['x'])\n",
    "\n",
    "model.add_subsystem('ks', om.KSComp(width=2))\n",
    "\n",
    "model.connect('comp.y', 'ks.g')\n",
    "\n",
    "model.ks.options['lower_flag'] = True\n",
    "prob.setup()\n",
    "prob.set_val('x', np.array([5.0, 4.0]))\n",
    "prob.run_model()\n",
    "\n",
    "print(prob.get_val('ks.KS'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "assert_near_equal(prob.get_val('ks.KS'), [[-12.0]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Typically, the KSComp is used to provide a constraint which aggregates many values into a single scalar constraint.\n",
    "Consider the following simple example, where we seek to maximize the peak of a parabola but also\n",
    "keep the peak of the parabola below a certain threshold value.  Clearly, the solution here is to have the peak of\n",
    "the parabola lie on the peak constraint.\n",
    "\n",
    "Note the resulting value of the offset \"k\" is not exactly 4.0 as we might expect.  The KS function\n",
    "provides a differentiable constraint aggregation, but the resulting scalar constraint is slightly\n",
    "conservative.\n",
    "\n",
    "**add_constraint**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "n = 50\n",
    "prob = om.Problem()\n",
    "model = prob.model\n",
    "\n",
    "prob.driver = om.ScipyOptimizeDriver()\n",
    "\n",
    "model.add_subsystem('comp', om.ExecComp('y = -3.0*x**2 + k',\n",
    "                                        x=np.zeros((n, )),\n",
    "                                        y=np.zeros((n, )),\n",
    "                                        k=0.0), promotes_inputs=['x', 'k'])\n",
    "\n",
    "model.add_subsystem('ks', om.KSComp(width=n, upper=4.0, add_constraint=True))\n",
    "\n",
    "model.add_design_var('k', lower=-10, upper=10)\n",
    "model.add_objective('k', scaler=-1)\n",
    "\n",
    "model.connect('comp.y', 'ks.g')\n",
    "\n",
    "prob.setup()\n",
    "prob.set_val('x', np.linspace(-np.pi/2, np.pi/2, n))\n",
    "prob.set_val('k', 5.)\n",
    "\n",
    "prob.run_driver()\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "x = prob.get_val('x')\n",
    "y = prob.get_val('comp.y')\n",
    "\n",
    "ax.plot(x, y, 'r.')\n",
    "ax.plot(x, 4.0*np.ones_like(x), 'k--')\n",
    "ax.set_xlabel('x')\n",
    "ax.set_ylabel('y')\n",
    "ax.grid(True)\n",
    "ax.text(-0.25, 0, f\"k = {prob.get_val('k')[0]:6.3f}\")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "assert(max(prob.get_val('comp.y')) <= 4.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**units**\n",
    "\n",
    "Finally, note that you can pass a units option to the KSComp that will define units on its input and output variables. There is only one unit, shared between both inputs and outputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 10\n",
    "\n",
    "model = om.Group()\n",
    "\n",
    "model.add_subsystem('ks', om.KSComp(width=n, units='m'), promotes_inputs=[('g', 'x')])\n",
    "model.set_input_defaults('x', range(n), units='ft')\n",
    "\n",
    "prob = om.Problem(model=model)\n",
    "prob.setup()\n",
    "prob.run_model()\n",
    "\n",
    "print(prob.get_val('ks.KS', indices=0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-output",
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "assert_near_equal(prob.get_val('ks.KS', indices=0), np.amax(prob.get_val('x')), tolerance=1e-8)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Example: KSComp for aggregating a constraint with a lower bound.**\n",
    "\n",
    "When you are using the KSComp to aggregate a constraint with a lower bound, set the \"lower_flag\"\n",
    "to True, and it will create a constraint that is satisfied when it is greater than the value\n",
    "specified in \"upper\". Because KSComp is a smooth max function, you can pass in any values to\n",
    "\"g\" and it will approximate both the min and max."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "n = 50\n",
    "prob = om.Problem()\n",
    "model = prob.model\n",
    "\n",
    "prob.driver = om.ScipyOptimizeDriver()\n",
    "\n",
    "model.add_subsystem('comp', om.ExecComp('y = -3.0*x**2 + k',\n",
    "                                        x=np.zeros((n, )),\n",
    "                                        y=np.zeros((n, )),\n",
    "                                        k=0.0), promotes_inputs=['x', 'k'])\n",
    "\n",
    "model.add_subsystem('ks', om.KSComp(width=n, upper=1.0, lower_flag=True))\n",
    "\n",
    "model.add_design_var('k', lower=-10, upper=10)\n",
    "model.add_objective('k', scaler=-1)\n",
    "\n",
    "model.connect('comp.y', 'ks.g')\n",
    "\n",
    "prob.setup()\n",
    "prob.set_val('x', np.linspace(-np.pi/2, np.pi/2, n))\n",
    "prob.set_val('k', 5.)\n",
    "\n",
    "prob.run_driver()\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "x = prob.get_val('x')\n",
    "y = prob.get_val('comp.y')\n",
    "\n",
    "ax.plot(x, y, 'r.')\n",
    "ax.set_xlabel('x')\n",
    "ax.set_ylabel('y')\n",
    "ax.grid(True)\n",
    "ax.text(-0.25, 0, f\"k = {prob.get_val('k')[0]:6.3f}\")\n",
    "\n",
    "plt.show()\n",
    "\n",
    "print(prob['ks.KS'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-output",
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "assert(prob['ks.KS'] <= 0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Example: KSComp for aggregating a constraint with a non-zero upper bound.**\n",
    "\n",
    "Here we use the KSComp to aggregate a constraint with an upper bound."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "n = 50\n",
    "prob = om.Problem()\n",
    "model = prob.model\n",
    "\n",
    "prob.driver = om.ScipyOptimizeDriver()\n",
    "\n",
    "model.add_subsystem('comp', om.ExecComp('y = -3.0*x**2 + k',\n",
    "                                        x=np.zeros((n, )),\n",
    "                                        y=np.zeros((n, )),\n",
    "                                        k=0.0), promotes_inputs=['x', 'k'])\n",
    "\n",
    "model.add_subsystem('ks', om.KSComp(width=n, upper=15.0))\n",
    "\n",
    "model.add_design_var('k', lower=-10, upper=10)\n",
    "model.add_objective('k', scaler=-1)\n",
    "\n",
    "model.connect('comp.y', 'ks.g')\n",
    "\n",
    "prob.setup()\n",
    "prob.set_val('x', np.linspace(-np.pi/2, np.pi/2, n))\n",
    "prob.set_val('k', 5.)\n",
    "\n",
    "prob.run_driver()\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "x = prob.get_val('x')\n",
    "y = prob.get_val('comp.y')\n",
    "\n",
    "ax.plot(x, y, 'r.')\n",
    "ax.set_xlabel('x')\n",
    "ax.set_ylabel('y')\n",
    "ax.grid(True)\n",
    "ax.text(-0.25, 0, f\"k = {prob.get_val('k')[0]:6.3f}\")\n",
    "\n",
    "plt.show()\n",
    "\n",
    "print(prob['ks.KS'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-output",
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "assert(prob['ks.KS'] <= 0)"
   ]
  }
 ],
 "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.1"
  },
  "orphan": true
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
